From b18eabe46a43e8cffb6df285a183028a41e5dfa5 Mon Sep 17 00:00:00 2001
From: Yazhou Cao <yazhou.cao@landing.ai>
Date: Tue, 9 Apr 2024 09:51:23 -0700
Subject: [PATCH] Fix lint errors

---
 vision_agent/tools/tools.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/vision_agent/tools/tools.py b/vision_agent/tools/tools.py
index 6e1f1bcd..6d2111f0 100644
--- a/vision_agent/tools/tools.py
+++ b/vision_agent/tools/tools.py
@@ -348,9 +348,9 @@ class DINOv(Tool):
                 "parameters": {
                     "prompt": [
                         {"mask": "balloon_mask1.jpg", "image": "balloon.jpg"},
-                        {"mask": "balloon_mask2.jpg", "image": "balloon.jpg"}
+                        {"mask": "balloon_mask2.jpg", "image": "balloon.jpg"},
                     ],
-                    "image": "input.jpg"
+                    "image": "input.jpg",
                 },
             },
             {
@@ -359,13 +359,15 @@ class DINOv(Tool):
                     "prompt": [
                         {"mask": "obj_mask1.jpg", "image": "background.jpg"},
                     ],
-                    "image": "input.jpg"
+                    "image": "input.jpg",
                 },
             },
         ],
     }
 
-    def __call__(self, prompt: List[Dict[str, str]], image: Union[str, ImageType]) -> Dict:
+    def __call__(
+        self, prompt: List[Dict[str, str]], image: Union[str, ImageType]
+    ) -> Dict:
         """Invoke the DINOv model.
 
         Parameters:
@@ -394,7 +396,7 @@ def __call__(self, prompt: List[Dict[str, str]], image: Union[str, ImageType]) -
         ) or "statusCode" not in resp_json:
             _LOGGER.error(f"Request failed: {resp_json}")
             raise ValueError(f"Request failed: {resp_json}")
-        rets = resp_json["data"]
+        rets: Dict[str, Any] = resp_json["data"]
         shape = rets.pop("mask_shape")
         mask_files = []
         for encoded_mask in rets["masks"]: