From b1f7bdc63fdb469485ed71493579dc24d2adb786 Mon Sep 17 00:00:00 2001 From: Yazhou Cao Date: Thu, 4 Apr 2024 15:02:18 -0700 Subject: [PATCH] Fix a bug in result parsing --- vision_agent/tools/tools.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/vision_agent/tools/tools.py b/vision_agent/tools/tools.py index 32c250eb..164ca6ea 100644 --- a/vision_agent/tools/tools.py +++ b/vision_agent/tools/tools.py @@ -292,12 +292,15 @@ def __call__(self, prompt: str, image: Union[str, ImageType]) -> Dict: data: Dict[str, Any] = resp_json["data"] ret_pred: Dict[str, List] = {"labels": [], "bboxes": [], "masks": []} if "bboxes" in data: - data["bboxes"] = [normalize_bbox(box, image_size) for box in data["bboxes"]] + ret_pred["bboxes"] = [ + normalize_bbox(box, image_size) for box in data["bboxes"] + ] if "masks" in data: - data["masks"] = [ + ret_pred["masks"] = [ rle_decode(mask_rle=mask, shape=data["mask_shape"]) - for mask in data["masks"][0] + for mask in data["masks"] ] + ret_pred["labels"] = data["labels"] return ret_pred