fix return values from countgd endpoint

vision_agent/tools/tool_utils.py

@@ -1,7 +1,7 @@
 import inspect
 import logging
 import os
-from typing import Any, Callable, Dict, List, MutableMapping, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, MutableMapping, Optional, Tuple
 
 import pandas as pd
 from IPython.display import display
@@ -34,7 +34,7 @@ def send_inference_request(
  files: Optional[List[Tuple[Any, ...]]] = None,
  v2: bool = False,
  metadata_payload: Optional[Dict[str, Any]] = None,
-) -> Union[Dict[str, Any], List[Dict[str, Any]]]:
+) -> Dict[str, Any]:
  # TODO: runtime_tag and function_name should be metadata_payload and now included
  # in the service payload
  try:

vision_agent/tools/tools.py

@@ -540,10 +540,10 @@ def countgd_counting(
  "box_threshold": box_threshold,
  }
  metadata_payload = {"function_name": "countgd_counting"}
- data: List[Dict[str, Any]] = send_inference_request(
+ resp: List[Dict[str, Any]] = send_inference_request(
  payload, "countgd", v2=True, metadata_payload=metadata_payload
- )
- return data
+ ) # type: ignore
+ return resp["data"]
 
 
 def countgd_example_based_counting(
@@ -589,10 +589,10 @@ def countgd_example_based_counting(
  "box_threshold": box_threshold,
  }
  metadata_payload = {"function_name": "countgd_example_based_counting"}
- data: List[Dict[str, Any]] = send_inference_request(
+ resp: List[Dict[str, Any]] = send_inference_request(
  payload, "countgd", v2=True, metadata_payload=metadata_payload
- )
- return data
+ ) # type: ignore
+ return resp["data"]
 
 
 def florence2_roberta_vqa(prompt: str, image: np.ndarray) -> str: