fix countgd

Author: Dayof

.github/workflows/ci_cd.yml

@@ -83,9 +83,6 @@ jobs:
     - name: Test with pytest
       run: |
         poetry run pytest -v tests/integ
-    - name: Test with pytest, dev env
-      run: |
-        LANDINGAI_API_KEY=$LANDINGAI_DEV_API_KEY LANDINGAI_URL=https://api.dev.landing.ai poetry run pytest -v tests/integration_dev
 
   release:
     name: Release

tests/integ/test_tools.py

@@ -32,6 +32,8 @@
     template_match,
     vit_image_classification,
     vit_nsfw_classification,
+    countgd_counting,
+    countgd_example_based_counting,
 )
 
 FINE_TUNE_ID = "65ebba4a-88b7-419f-9046-0750e30250da"
@@ -387,3 +389,18 @@ def test_generate_hed():
     )
 
     assert result.shape == img.shape
+
+
+def test_countgd_counting() -> None:
+    img = ski.data.coins()
+    result = countgd_counting(image=img, prompt="coin")
+    assert len(result) == 24
+
+
+def test_countgd_example_based_counting() -> None:
+    img = ski.data.coins()
+    result = countgd_example_based_counting(
+        visual_prompts=[[85, 106, 122, 145]],
+        image=img,
+    )
+    assert len(result) == 24

tests/integration_dev/__init__.py

@@ -700,22 +700,18 @@ def countgd_counting(
             {'score': 0.98, 'label': 'flower', 'bbox': [0.44, 0.24, 0.49, 0.58},
         ]
     """
-    buffer_bytes = numpy_to_bytes(image)
-    files = [("image", buffer_bytes)]
+    image_b64 = convert_to_b64(image)
     prompt = prompt.replace(", ", " .")
-    payload = {"prompts": [prompt], "model": "countgd"}
+    payload = {"prompt": prompt, "image": image_b64}
     metadata = {"function_name": "countgd_counting"}
-    resp_data = send_task_inference_request(
-        payload, "text-to-object-detection", files=files, metadata=metadata
-    )
-    bboxes_per_frame = resp_data[0]
+    resp_data = send_task_inference_request(payload, "countgd", metadata=metadata)
     bboxes_formatted = [
         ODResponseData(
             label=bbox["label"],
-            bbox=list(map(lambda x: round(x, 2), bbox["bounding_box"])),
+            bbox=list(map(lambda x: round(x, 2), bbox["bbox"])),
             score=round(bbox["score"], 2),
         )
-        for bbox in bboxes_per_frame
+        for bbox in resp_data
     ]
     filtered_bboxes = filter_bboxes_by_threshold(bboxes_formatted, box_threshold)
     return [bbox.model_dump() for bbox in filtered_bboxes]