diff --git a/vision_agent/tools/tools.py b/vision_agent/tools/tools.py
index 8f3dfd50..bd640e76 100644
--- a/vision_agent/tools/tools.py
+++ b/vision_agent/tools/tools.py
@@ -1,8 +1,13 @@
+from typing import Dict, List, Union
+
+from PIL.Image import Image as ImageType
+
+
 class Classifier:
     def __init__(self, prompt: str):
         self.prompt = prompt
 
-    def __call__(self: image: Union[str, Image]) -> List[Dict]:
+    def __call__(self, image: Union[str, ImageType]) -> List[Dict]:
         raise NotImplementedError
 
 
@@ -10,7 +15,7 @@ class Detector:
     def __init__(self, prompt: str):
         self.prompt = prompt
 
-    def __call__(self: image: Union[str, Image]) -> List[Dict]:
+    def __call__(self, image: Union[str, ImageType]) -> List[Dict]:
         raise NotImplementedError
 
 
@@ -18,5 +23,5 @@ class Segmentor:
     def __init__(self, prompt: str):
         self.prompt = prompt
 
-    def __call__(self: image: Union[str, Image]) -> List[Dict]:
+    def __call__(self, image: Union[str, ImageType]) -> List[Dict]:
         raise NotImplementedError