fix type errors

landing-ai · Mar 6, 2024 · 4767d6d · 4767d6d
1 parent 91d046e
commit 4767d6d
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/vision_agent/tools/tools.py b/vision_agent/tools/tools.py
@@ -1,22 +1,27 @@
+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
 
 
 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
 
 
 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