Add osam_yoloworld

osam/__main__.py

@@ -5,6 +5,7 @@
 import sys
 
 import click
+import imgviz
 import numpy as np
 import PIL.Image
 import uvicorn
@@ -127,15 +128,23 @@ def run(model_name: str, image_path: str, prompt, json: bool) -> None:
         click.echo(response.model_dump_json())
     else:
         visualization: np.ndarray
-        if len(response.masks) == 1:
+        if response.masks is not None and len(response.masks) == 1:
             visualization = (
                 0.5 * image
                 + 0.5
                 * np.array([0, 255, 0])[None, None, :]
                 * (response.masks[0] > 0)[:, :, None]
             ).astype(np.uint8)
         else:
-            raise NotImplementedError
+            visualization = imgviz.instances2rgb(
+                image=image,
+                labels=[
+                    1 + request.prompt.texts.index(text) for text in response.texts
+                ],
+                bboxes=response.bounding_boxes,
+                captions=response.texts,
+            )
+
         sys.stdout.buffer.write(_image_ndarray_to_data(visualization))
 
 

osam/_models.py

@@ -1,3 +1,4 @@
 # import modules to register the models
 import osam_efficientsam  # noqa: F401
 import osam_sam  # noqa: F401
+import osam_yoloworld  # noqa: F401