fix mypy errors

landing-ai · Apr 9, 2024 · 3f02965 · 3f02965
1 parent 81cdb3f
commit 3f02965
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/vision_agent/tools/tools.py b/vision_agent/tools/tools.py
@@ -124,7 +124,7 @@ def __call__(self, prompt: str, image: Union[str, ImageType]) -> Dict:
  round(prob, 4) for prob in resp_json["data"]["scores"]
  ]
 
- return resp_json["data"]
+ return resp_json["data"] # type: ignore
 
 
 class GroundingDINO(Tool):
@@ -340,8 +340,14 @@ class AgentGroundingSAM(GroundingSAM):
  returns the file name. This makes it easier for agents to use.
  """
 
- def __call__(self, prompt: str, image: Union[str, ImageType]) -> Dict:
- rets = super().__call__(prompt, image)
+ def __call__(
+ self,
+ prompt: str,
+ image: Union[str, ImageType],
+ box_threshold: float = 0.2,
+ iou_threshold: float = 0.75,
+ ) -> Dict:
+ rets = super().__call__(prompt, image, box_threshold, iou_threshold)
  mask_files = []
  for mask in rets["masks"]:
  with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp: