diff --git a/osam/_data/dogs.jpg b/osam/_data/dogs.jpg new file mode 100644 index 0000000..d834f61 Binary files /dev/null and b/osam/_data/dogs.jpg differ diff --git a/osam/apis_test.py b/osam/apis_test.py index e69de29..5b8f66f 100644 --- a/osam/apis_test.py +++ b/osam/apis_test.py @@ -0,0 +1,48 @@ +import pathlib + +import imgviz +import pytest + +from . import apis +from . import types + +here = pathlib.Path(__file__).resolve().parent + + +@pytest.mark.parametrize( + "model", + ["efficientsam:10m", "efficientsam:latest", "sam:100m", "sam:300m", "sam:latest"], +) +def test_generate_point_to_mask(model: str) -> None: + image = imgviz.io.imread(here / "_data" / "dogs.jpg") + request: types.GenerateRequest = types.GenerateRequest(model=model, image=image) + response: types.GenerateResponse = apis.generate(request=request) + + assert response.model == model + + assert len(response.annotations) == 1 + annotation: types.Annotation = response.annotations[0] + assert annotation.text is None + assert annotation.score is None + assert annotation.mask is not None + assert annotation.mask.dtype == bool + assert annotation.mask.shape == image.shape[:2] + + +def test_generate_text_to_bounding_box() -> None: + model: str = "yoloworld:latest" + + image = imgviz.io.imread(here / "_data" / "dogs.jpg") + request: types.GenerateRequest = types.GenerateRequest( + model=model, image=image, prompt=types.Prompt(texts=["dog"]) + ) + response: types.GenerateResponse = apis.generate(request=request) + + assert response.model == model + + assert len(response.annotations) == 3 + for annotation in response.annotations: + assert annotation.bounding_box is not None + assert annotation.text == "dog" + assert isinstance(annotation.score, float) + assert annotation.mask is None