Added test cases for the tools added

tests/test_tools.py

@@ -2,6 +2,13 @@
 from PIL import Image
 
 from vision_agent.tools.tools import CLIP, GroundingDINO, GroundingSAM, ImageCaption
+from vision_agent.tools.tools_v2 import (
+ clip,
+ zero_shot_counting,
+ visual_prompt_counting,
+ image_question_answering,
+ ocr,
+)
 
 
 def test_grounding_dino():
@@ -28,15 +35,49 @@ def test_grounding_sam():
 
 
 def test_clip():
- img = Image.fromarray(ski.data.coins())
- result = CLIP()(
- prompt="coins",
+ img = ski.data.coins()
+ result = clip(
+ classes=["coins", "notes"],
  image=img,
  )
- assert result["scores"] == [1.0]
+ assert result["scores"] == [0.99, 0.01]
 
 
 def test_image_caption() -> None:
  img = Image.fromarray(ski.data.coins())
  result = ImageCaption()(image=img)
  assert result["text"]
+
+
+def test_zero_shot_counting() -> None:
+ img = Image.fromarray(ski.data.coins())
+ result = zero_shot_counting(
+ image=img,
+ )
+ assert result["count"] == 24
+
+
+def test_visual_prompt_counting() -> None:
+ img = Image.fromarray(ski.data.checkerboard())
+ result = visual_prompt_counting(
+ visual_prompt={"bbox": [0.125, 0, 0.25, 0.125]},
+ image=img,
+ )
+ assert result["count"] == 32
+
+
+def test_image_question_answering() -> None:
+ img = Image.fromarray(ski.data.rocket())
+ result = image_question_answering(
+ prompt="Is the scene captured during day or night ?",
+ image=img,
+ )
+ assert result == "night"
+
+
+def test_ocr() -> None:
+ img = Image.fromarray(ski.data.page())
+ result = ocr(
+ image=img,
+ )
+ assert result[0]["label"] == "Region-based segmentation"