fixed integ tests

landing-ai · Aug 10, 2024 · 6f0dd09 · 6f0dd09
1 parent 6ac6e58
commit 6f0dd09
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/tests/integ/test_tools.py b/tests/integ/test_tools.py
@@ -118,7 +118,7 @@ def test_nsfw_classification():
  result = vit_nsfw_classification(
  image=img,
  )
- assert result["labels"] == "normal"
+ assert result["label"] == "normal"
 
 
 def test_image_caption() -> None:

diff --git a/vision_agent/tools/tools.py b/vision_agent/tools/tools.py
@@ -537,7 +537,7 @@ def vit_nsfw_classification(image: np.ndarray) -> Dict[str, Any]:
  Example
  -------
  >>> vit_nsfw_classification(image)
- {"labels": "normal", "scores": 0.68},
+ {"label": "normal", "scores": 0.68},
  """
 
  image_b64 = convert_to_b64(image)
@@ -603,7 +603,7 @@ def florencev2_image_caption(image: np.ndarray, detail_caption: bool = True) ->
  }
 
  answer = send_inference_request(data, "florence2", v2=True)
- return answer["text"][0] # type: ignore
+ return answer[task] # type: ignore
 
 
 def florencev2_object_detection(image: np.ndarray) -> List[Dict[str, Any]]: