add: tests

This commit is contained in:
soumik12345 2025-07-24 15:17:05 +05:30
parent 8fa21dd165
commit 137078602b
1 changed files with 110 additions and 0 deletions

View File

@ -6,7 +6,10 @@ from contextlib import nullcontext as does_not_raise
import numpy as np
import pytest
from supervision.config import CLASS_NAME_DATA_FIELD
from supervision.detection.core import Detections
from supervision.detection.vlm import (
VLM,
from_florence_2,
from_google_gemini_2_0,
from_google_gemini_2_5,
@ -1122,3 +1125,110 @@ def test_from_google_gemini_2_5(
assert masks is not None
assert masks.shape == expected_results[4].shape
assert np.array_equal(masks, expected_results[4])
@pytest.mark.parametrize(
"exception, result, resolution_wh, classes, expected_detections",
[
(
pytest.raises(ValueError),
"",
(100, 100),
None,
None,
), # empty text
(
pytest.raises(ValueError),
"random text",
(100, 100),
None,
None,
), # random text
(
does_not_raise(),
"<|ref|>cat<|/ref|><|det|>[[100, 200, 300, 400]]<|/det|>",
(1000, 1000),
None,
Detections(
xyxy=np.array([[100.1, 200.2, 300.3, 400.4]]),
class_id=np.array([0]),
data={CLASS_NAME_DATA_FIELD: np.array(["cat"])},
),
), # single box, no classes
(
does_not_raise(),
"<|ref|>cat<|/ref|><|det|>[[100, 200, 300, 400]]<|/det|>",
(1000, 1000),
["cat", "dog"],
Detections(
xyxy=np.array([[100.1, 200.2, 300.3, 400.4]]),
class_id=np.array([0]),
data={CLASS_NAME_DATA_FIELD: np.array(["cat"])},
),
), # single box, with classes
(
does_not_raise(),
"<|ref|>person<|/ref|><|det|>[[100, 200, 300, 400]]<|/det|>",
(1000, 1000),
["cat", "dog"],
Detections.empty(),
), # single box, wrong class
(
does_not_raise(),
(
"<|ref|>cat<|/ref|><|det|>[[100, 200, 300, 400]]<|/det|>"
"<|ref|>dog<|/ref|><|det|>[[500, 600, 700, 800]]<|/det|>"
),
(1000, 1000),
["cat"],
Detections(
xyxy=np.array([[100.1, 200.2, 300.3, 400.4]]),
class_id=np.array([0]),
data={CLASS_NAME_DATA_FIELD: np.array(["cat"])},
),
), # multiple boxes, one class correct
(
pytest.raises(ValueError),
"<|ref|>cat<|/ref|>",
(100, 100),
None,
None,
), # only ref
(
pytest.raises(ValueError),
"<|det|>[[100, 200, 300, 400]]<|/det|>",
(100, 100),
None,
None,
), # only det
],
)
def test_from_deepseek_vl_2(
exception,
result: str,
resolution_wh: tuple[int, int],
classes: list[str] | None,
expected_detections: Detections,
):
with exception:
detections = Detections.from_vlm(
vlm=VLM.DEEPSEEK_VL_2,
result=result,
resolution_wh=resolution_wh,
classes=classes,
)
if expected_detections is None:
return
assert len(detections) == len(expected_detections)
if len(detections) == 0:
return
assert np.allclose(detections.xyxy, expected_detections.xyxy, atol=1e-1)
assert np.array_equal(detections.class_id, expected_detections.class_id)
assert np.array_equal(
detections.data[CLASS_NAME_DATA_FIELD],
expected_detections.data[CLASS_NAME_DATA_FIELD],
)