Fix a bug in result parsing

landing-ai · Apr 4, 2024 · b1f7bdc · b1f7bdc
1 parent e48df0e
commit b1f7bdc
Showing 1 changed file with 6 additions and 3 deletions.
diff --git 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