Provide non_maximum_suppression in osam.apis

osam/_models/efficientsam/__init__.py

@@ -80,7 +80,7 @@ def generate(self, request: types.GenerateRequest) -> types.GenerateResponse:
         mask = masks[0, 0, 0, :, :]  # (1, 1, 3, H, W) -> (H, W)
         mask = mask > 0.0
 
-        bbox = imgviz.instances.mask_to_bbox([mask])[0].astype(int)
+        bbox = imgviz.instances.masks_to_bboxes(masks=[mask])[0].astype(int)
 
         return types.GenerateResponse(
             model=self.name,

osam/_models/yoloworld/__init__.py

@@ -2,9 +2,9 @@
 
 import imgviz
 import numpy as np
-import onnxruntime
 from loguru import logger
 
+from ... import apis
 from ... import types
 from . import clip
 
@@ -54,8 +54,7 @@ def generate(self, request: types.GenerateRequest) -> types.GenerateResponse:
         max_annotations = (
             len(bboxes) if prompt.max_annotations is None else prompt.max_annotations
         )
-        bboxes, scores, labels = _non_maximum_suppression(
-            inference_session=self._inference_sessions["nms"],
+        bboxes, scores, labels = apis.non_maximum_suppression(
             boxes=bboxes,
             scores=scores,
             iou_threshold=iou_threshold,
@@ -98,10 +97,6 @@ class YoloWorldXL(_YoloWorld):
             url="https://github.com/wkentaro/yolo-world-onnx/releases/download/v0.1.0/yolo_world_v2_xl_vlpan_bn_2e-3_100e_4x8gpus_obj365v1_goldg_train_lvis_minival.onnx",
             hash="sha256:92660c6456766439a2670cf19a8a258ccd3588118622a15959f39e253731c05d",
         ),
-        "nms": types.Blob(
-            url="https://github.com/wkentaro/yolo-world-onnx/releases/download/v0.1.0/non_maximum_suppression.onnx",
-            hash="sha256:328310ba8fdd386c7ca63fc9df3963cc47b1268909647abd469e8ebdf7f3d20a",
-        ),
     }
 
 
@@ -145,37 +140,3 @@ def _untransform_bboxes(
     bboxes[:, 1::2] = np.clip(bboxes[:, 1::2], 0, original_image_hw[0])
     bboxes = bboxes.round().astype(int)
     return bboxes
-
-
-def _non_maximum_suppression(
-    inference_session: onnxruntime.InferenceSession,
-    boxes: np.ndarray,
-    scores: np.ndarray,
-    iou_threshold: float,
-    score_threshold: float,
-    max_num_detections: int,
-) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
-    selected_indices = inference_session.run(
-        output_names=["selected_indices"],
-        input_feed={
-            "boxes": boxes[None, :, :],
-            "scores": scores[None, :, :].transpose(0, 2, 1),
-            "max_output_boxes_per_class": np.array(
-                [max_num_detections], dtype=np.int64
-            ),
-            "iou_threshold": np.array([iou_threshold], dtype=np.float32),
-            "score_threshold": np.array([score_threshold], dtype=np.float32),
-        },
-    )[0]
-    labels = selected_indices[:, 1]
-    box_indices = selected_indices[:, 2]
-    boxes = boxes[box_indices]
-    scores = scores[box_indices, labels]
-
-    if len(boxes) > max_num_detections:
-        keep_indices = np.argsort(scores)[-max_num_detections:]
-        boxes = boxes[keep_indices]
-        scores = scores[keep_indices]
-        labels = labels[keep_indices]
-
-    return boxes, scores, labels

osam/apis.py

@@ -1,7 +1,11 @@
 from typing import List
 from typing import Optional
+from typing import Tuple
 from typing import Type
 
+import numpy as np
+import onnxruntime
+
 from . import _models
 from . import types
 
@@ -42,3 +46,53 @@ def generate(request: types.GenerateRequest) -> types.GenerateResponse:
 
     response: types.GenerateResponse = running_model.generate(request=request)
     return response
+
+
+_non_maximum_suppression_inference_session: Optional[onnxruntime.InferenceSession] = (
+    None
+)
+
+
+def non_maximum_suppression(
+    boxes: np.ndarray,
+    scores: np.ndarray,
+    iou_threshold: float,
+    score_threshold: float,
+    max_num_detections: int,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    global _non_maximum_suppression_inference_session
+    if _non_maximum_suppression_inference_session is None:
+        blob = types.Blob(
+            url="https://github.com/wkentaro/yolo-world-onnx/releases/download/v0.1.0/non_maximum_suppression.onnx",  # noqa
+            hash="sha256:328310ba8fdd386c7ca63fc9df3963cc47b1268909647abd469e8ebdf7f3d20a",
+        )
+        blob.pull()
+        _non_maximum_suppression_inference_session = onnxruntime.InferenceSession(
+            blob.path, providers=["CPUExecutionProvider"]
+        )
+    inference_session = _non_maximum_suppression_inference_session
+
+    selected_indices = inference_session.run(
+        output_names=["selected_indices"],
+        input_feed={
+            "boxes": boxes[None, :, :],
+            "scores": scores[None, :, :].transpose(0, 2, 1),
+            "max_output_boxes_per_class": np.array(
+                [max_num_detections], dtype=np.int64
+            ),
+            "iou_threshold": np.array([iou_threshold], dtype=np.float32),
+            "score_threshold": np.array([score_threshold], dtype=np.float32),
+        },
+    )[0]
+    labels = selected_indices[:, 1]
+    box_indices = selected_indices[:, 2]
+    boxes = boxes[box_indices]
+    scores = scores[box_indices, labels]
+
+    if len(boxes) > max_num_detections:
+        keep_indices = np.argsort(scores)[-max_num_detections:]
+        boxes = boxes[keep_indices]
+        scores = scores[keep_indices]
+        labels = labels[keep_indices]
+
+    return boxes, scores, labels

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