diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 7d142000..ed3cdf15 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -37,7 +37,8 @@ from supervision.detection.vlm import ( LMM, VLM, from_florence_2, - from_google_gemini, + from_google_gemini_2_0, + from_google_gemini_2_5, from_moondream, from_paligemma, from_qwen_2_5_vl, @@ -815,6 +816,15 @@ class Detections: Creates a Detections object from the given result string based on the specified Large Multimodal Model (LMM). + | Name | Enum (sv.LMM) | Tasks | Required parameters | Optional parameters | + |---------------------|----------------------|-------------------------|-----------------------------|---------------------| + | PaliGemma | `PALIGEMMA` | detection | `resolution_wh` | `classes` | + | PaliGemma 2 | `PALIGEMMA` | detection | `resolution_wh` | `classes` | + | Qwen2.5-VL | `QWEN_2_5_VL` | detection | `resolution_wh`, `input_wh` | `classes` | + | Google Gemini 2.0 | `GOOGLE_GEMINI_2_0` | detection | `resolution_wh` | `classes` | + | Google Gemini 2.5 | `GOOGLE_GEMINI_2_5` | detection, segmentation | `resolution_wh` | `classes` | + | Moondream | `MOONDREAM` | detection | `resolution_wh` | | + Args: lmm (Union[LMM, str]): The type of LMM (Large Multimodal Model) to use. result (str): The result string containing the detection data. @@ -828,8 +838,9 @@ class Detections: disallowed arguments are provided. ValueError: If the specified LMM is not supported. - Examples: + !!! example "PaliGemma" ```python + import supervision as sv paligemma_result = " cat" @@ -849,7 +860,7 @@ class Detections: # {'class_name': array(['cat'], dtype=' Detections: """ + Creates a Detections object from the given result string based on the specified Vision Language Model (VLM). + | Name | Enum (sv.VLM) | Tasks | Required parameters | Optional parameters | + |---------------------|----------------------|-------------------------|-----------------------------|---------------------| + | PaliGemma | `PALIGEMMA` | detection | `resolution_wh` | `classes` | + | PaliGemma 2 | `PALIGEMMA` | detection | `resolution_wh` | `classes` | + | Qwen2.5-VL | `QWEN_2_5_VL` | detection | `resolution_wh`, `input_wh` | `classes` | + | Google Gemini 2.0 | `GOOGLE_GEMINI_2_0` | detection | `resolution_wh` | `classes` | + | Google Gemini 2.5 | `GOOGLE_GEMINI_2_5` | detection, segmentation | `resolution_wh` | `classes` | + | Moondream | `MOONDREAM` | detection | `resolution_wh` | | + Args: - vlm (Union[VLM, str]): The type of VLM (Large Multimodal Model) to use. + vlm (Union[VLM, str]): The type of VLM (Vision Language Model) to use. result (str): The result string containing the detection data. **kwargs (Any): Additional keyword arguments required by the specified VLM. @@ -994,8 +1110,9 @@ class Detections: disallowed arguments are provided. ValueError: If the specified VLM is not supported. - Examples: + !!! example "PaliGemma" ```python + import supervision as sv paligemma_result = " cat" @@ -1015,7 +1132,7 @@ class Detections: # {'class_name': array(['cat'], dtype=') ([\w\s\-]+)" @@ -189,14 +191,9 @@ def from_qwen_2_5_vl( class_name (np.ndarray): An array of shape `(n,)` containing the class labels for each bounding box """ - in_w, in_h = input_wh - out_w, out_h = resolution_wh - if in_w <= 0 or in_h <= 0 or out_w <= 0 or out_h <= 0: - raise ValueError( - f"Both input and resolution dimensions must be positive. " - f"Got input_wh=({in_w}, {in_h}), resolution_wh=({out_w}, {out_h})." - ) + in_w, in_h = validate_resolution(input_wh) + out_w, out_h = validate_resolution(resolution_wh) pattern = re.compile(r"```json\s*(.*?)\s*```", re.DOTALL) @@ -325,7 +322,7 @@ def from_florence_2( f"Expected string to end in location tags, but got {result}" ) - w, h = resolution_wh + w, h = validate_resolution(resolution_wh) xyxy = np.array([match.groups()], dtype=np.float32) xyxy *= np.array([w, h, w, h]) / 1000 result_string = result[: match.start()] @@ -335,7 +332,7 @@ def from_florence_2( assert False, f"Unimplemented task: {task}" -def from_google_gemini( +def from_google_gemini_2_0( result: str, resolution_wh: Tuple[int, int], classes: Optional[List[str]] = None, @@ -377,11 +374,7 @@ def from_google_gemini( """ - w, h = resolution_wh - if w <= 0 or h <= 0: - raise ValueError( - f"Both dimensions in resolution_wh must be positive. Got ({w}, {h})." - ) + w, h = validate_resolution(resolution_wh) lines = result.splitlines() for i, line in enumerate(lines): @@ -396,14 +389,15 @@ def from_google_gemini( return np.empty((0, 4)), None, np.empty((0,), dtype=str) labels = [] - xyxy = [] + boxes_list = [] + for item in data: if "box_2d" not in item or "label" not in item: continue labels.append(item["label"]) box = item["box_2d"] # Gemini bbox order is [y_min, x_min, y_max, x_max] - xyxy.append( + boxes_list.append( denormalize_boxes( np.array([box[1], box[0], box[3], box[2]]).astype(np.float64), resolution_wh=(w, h), @@ -411,10 +405,10 @@ def from_google_gemini( ) ) - if not xyxy: + if not boxes_list: return np.empty((0, 4)), None, np.empty((0,), dtype=str) - xyxy = np.array(xyxy) + xyxy = np.array(boxes_list) class_name = np.array(labels) class_id = None @@ -427,6 +421,168 @@ def from_google_gemini( return xyxy, class_id, class_name +def from_google_gemini_2_5( + result: str, + resolution_wh: Tuple[int, int], + classes: Optional[List[str]] = None, +) -> Tuple[ + np.ndarray, + Optional[np.ndarray], + np.ndarray, + Optional[np.ndarray], + Optional[np.ndarray], +]: + """ + Parse and scale bounding boxes and masks from Google Gemini 2.5 style + [JSON output](https://ai.google.dev/gemini-api/docs/vision?lang=python). + + The JSON is expected to be enclosed in triple backticks with the format: + ```json + [ + { + "box_2d": [x1, y1, x2, y2], + "mask": "data:image/png;base64,...", + "label": "some class name", + "confidence": 0.95, + }, + ... + ] + ``` + + Args: + result: String containing the JSON snippet enclosed by triple backticks. + resolution_wh: (output_width, output_height) to which we rescale the boxes. + classes: Optional list of valid class names. If provided, returned boxes/labels + are filtered to only those classes found here. + + Returns: + xyxy (np.ndarray): An array of shape `(n, 4)` containing + the bounding boxes coordinates in format `[x1, y1, x2, y2]` + class_id (np.ndarray): An array of shape `(n,)` containing + the class indices for each bounding box + class_name (np.ndarray): An array of shape `(n,)` containing + the class labels for each bounding box + confidence: Optional[np.ndarray]: An array of shape `(n,)` containing + the confidence scores for each bounding box. If not provided, + it defaults to 0.0 for each box. + masks (Optional[np.ndarray]): An array of shape `(n, h, w)` containing + the segmentation masks for each bounding box + """ + w, h = validate_resolution(resolution_wh) + + lines = result.splitlines() + for i, line in enumerate(lines): + if line == "```json": + result = "\n".join(lines[i + 1 :]) + result = result.split("```")[0] + break + + try: + data = json.loads(result) + except json.JSONDecodeError: + return ( + np.empty((0, 4)), + np.array([], dtype=int), + np.array([], dtype=str), + np.array([], dtype=float), + None, + ) + + boxes_list: list = [] + labels_list: list = [] + confidence_list: Optional[list] = [] + masks_list: Optional[list] = [] + + for item in data: + if "box_2d" not in item or "label" not in item: + continue + labels_list.append(item["label"]) + box = item["box_2d"] + # Gemini bbox order is [y_min, x_min, y_max, x_max] + absolute_bbox = denormalize_boxes( + np.array([box[1], box[0], box[3], box[2]]).astype(np.float64), + resolution_wh=(w, h), + normalization_factor=1000, + ) + boxes_list.append(absolute_bbox) + + if "mask" in item: + if masks_list is not None: + png_str = item["mask"] + if not png_str.startswith("data:image/png;base64,"): + masks_list.append(np.zeros((h, w), dtype=bool)) + continue + + png_str = png_str.removeprefix("data:image/png;base64,") + png_str = base64.b64decode(png_str) + mask_img = Image.open(io.BytesIO(png_str)) + + y_min, y_max = int(absolute_bbox[1]), int(absolute_bbox[3]) + x_min, x_max = int(absolute_bbox[0]), int(absolute_bbox[2]) + + bbox_height = y_max - y_min + bbox_width = x_max - x_min + + if bbox_height > 0 and bbox_width > 0: + mask_img = mask_img.resize( + (bbox_width, bbox_height), resample=Image.Resampling.BILINEAR + ) + np_mask = np.zeros((h, w), dtype=bool) + np_mask[y_min:y_max, x_min:x_max] = np.array(mask_img) > 0 + masks_list.append(np_mask) + else: + masks_list.append(np.zeros((h, w), dtype=bool)) + else: + masks_list = None + + if "confidence" in item: + if confidence_list is not None: + confidence_list.append(item["confidence"]) + else: + confidence_list = None + + if not boxes_list: + return ( + np.empty((0, 4)), + np.array([], dtype=int), + np.array([], dtype=str), + np.array([], dtype=float), + None, + ) + + xyxy = np.array(boxes_list, dtype=float) + class_name = np.array(labels_list) + class_id: np.ndarray + + if classes is not None: + mask = np.array([name in classes for name in class_name], dtype=bool) + xyxy = xyxy[mask] + class_name = class_name[mask] + class_id = np.array([classes.index(name) for name in class_name]) + if masks_list is not None: + masks_list = [m for m, keep in zip(masks_list, mask) if keep] + + if confidence_list is not None: + confidence_list = [c for c, keep in zip(confidence_list, mask) if keep] + else: + unique_labels = sorted(list(set(class_name))) + label_to_id = {label: i for i, label in enumerate(unique_labels)} + class_id = np.array([label_to_id[name] for name in class_name]) + + confidence = ( + np.array(confidence_list, dtype=float) if confidence_list is not None else None + ) + masks = np.array(masks_list) if masks_list is not None else None + + return ( + xyxy, + class_id, + class_name, + confidence, + masks, + ) + + def from_moondream( result: dict, resolution_wh: Tuple[int, int], @@ -450,7 +606,6 @@ def from_moondream( ] } - Args: result: Dictionary containing the JSON output from the model. resolution_wh: (output_width, output_height) to which we rescale the boxes. @@ -458,7 +613,7 @@ def from_moondream( Returns: xyxy (np.ndarray): An array of shape `(n, 4)` containing the bounding boxes coordinates in format `[x1, y1, x2, y2]` - """ # docs + """ w, h = resolution_wh if w <= 0 or h <= 0: @@ -467,7 +622,7 @@ def from_moondream( ) if "objects" not in result or not isinstance(result["objects"], list): - return np.empty((0, 4)) + return np.empty((0, 4), dtype=float) denormalize_xyxy = [] diff --git a/supervision/validators/__init__.py b/supervision/validators/__init__.py index 29ac0da9..f40b6f83 100644 --- a/supervision/validators/__init__.py +++ b/supervision/validators/__init__.py @@ -1,4 +1,4 @@ -from typing import Any, Dict +from typing import Any, Dict, Tuple import numpy as np @@ -138,3 +138,26 @@ def validate_keypoints_fields( validate_class_id(class_id, n) validate_keypoint_confidence(confidence, n, m) validate_data(data, n) + + +def validate_resolution(resolution: Any) -> Tuple[int, int]: + if not (isinstance(resolution, tuple) and len(resolution) == 2): + raise ValueError( + f""" + resolution must be a tuple of two integers, got + {type(resolution)} with value {resolution} + """ + ) + w, h = resolution + if not (isinstance(w, int) and isinstance(h, int)): + raise ValueError( + f""" + Both elements in resolution must be integers. + Got types ({type(w)}, {type(h)}) + """ + ) + if w <= 0 or h <= 0: + raise ValueError( + f"Both dimensions in resolution must be positive. Got ({w}, {h})." + ) + return w, h diff --git a/test/detection/test_vlm.py b/test/detection/test_vlm.py index 7b9acb0c..6464c234 100644 --- a/test/detection/test_vlm.py +++ b/test/detection/test_vlm.py @@ -7,7 +7,8 @@ import pytest from supervision.detection.vlm import ( from_florence_2, - from_google_gemini, + from_google_gemini_2_0, + from_google_gemini_2_5, from_moondream, from_paligemma, from_qwen_2_5_vl, @@ -492,7 +493,7 @@ def test_from_google_gemini( expected_results: Tuple[np.ndarray, Optional[np.ndarray], np.ndarray], ) -> None: with exception: - xyxy, class_id, class_name = from_google_gemini( + xyxy, class_id, class_name = from_google_gemini_2_0( result=result, resolution_wh=resolution_wh, classes=classes ) if expected_results is not None: @@ -883,3 +884,241 @@ def test_florence_2( assert result[3] is None else: np.testing.assert_array_equal(result[3], expected_results[3]) + + +@pytest.mark.parametrize( + "exception, result, resolution_wh, classes, expected_results", + [ + ( + does_not_raise(), + "random text", + (1000, 1000), + None, + ( + np.empty((0, 4)), + np.empty(0, dtype=int), + np.empty(0, dtype=str), + np.empty(0, dtype=float), + None, + ), + ), + ( + does_not_raise(), + "```json\ninvalid json\n```", + (1000, 1000), + None, + ( + np.empty((0, 4)), + np.empty(0, dtype=int), + np.empty(0, dtype=str), + np.empty(0, dtype=float), + None, + ), + ), + ( + does_not_raise(), + "```json\n[]\n```", + (1000, 1000), + None, + ( + np.empty((0, 4)), + np.empty(0, dtype=int), + np.empty(0, dtype=str), + np.empty(0, dtype=float), + None, + ), + ), + ( + does_not_raise(), + """```json + [ + {"box_2d": [100, 200, 300, 400], "label": "cat", "confidence": 0.8} + ] + ```""", + (1000, 500), + None, + ( + np.array([[200.0, 50.0, 400.0, 150.0]]), + np.array([0]), + np.array(["cat"], dtype=str), + np.array([0.8]), + None, + ), + ), + ( + does_not_raise(), + """```json + [ + {"box_2d": [10, 20, 110, 120], "label": "cat", "confidence": 0.8}, + {"box_2d": [50, 100, 150, 200], "label": "dog", "confidence": 0.9} + ] + ```""", + (640, 480), + None, + ( + np.array([[12.8, 4.8, 76.8, 52.8], [64.0, 24.0, 128.0, 72.0]]), + np.array([0, 1]), + np.array(["cat", "dog"], dtype=str), + np.array([0.8, 0.9]), + None, + ), + ), + ( + does_not_raise(), + """```json + [ + {"box_2d": [10, 20, 110, 120], "label": "cat", "confidence": 0.8} + ] + ```""", + (640, 480), + ["dog", "person"], + ( + np.empty((0, 4)), + np.empty(0, dtype=int), + np.empty(0, dtype=str), + np.empty(0, dtype=float), + None, + ), + ), + ( + does_not_raise(), + """```json + [ + {"box_2d": [10, 20, 110, 120], "label": "cat", "confidence": 0.8}, + {"box_2d": [50, 100, 150, 200], "label": "dog", "confidence": 0.9} + ] + ```""", + (640, 480), + ["person", "dog"], + ( + np.array([[64.0, 24.0, 128.0, 72.0]]), + np.array([1]), + np.array(["dog"], dtype=str), + np.array([0.9]), + None, + ), + ), + ( + does_not_raise(), + """```json + [ + {"box_2d": [10, 20, 110, 120], "label": "cat", "confidence": 0.8}, + {"box_2d": [50, 100, 150, 200], "label": "dog", "confidence": 0.9} + ] + ```""", + (640, 480), + ["cat", "dog"], + ( + np.array([[12.8, 4.8, 76.8, 52.8], [64.0, 24.0, 128.0, 72.0]]), + np.array([0, 1]), + np.array(["cat", "dog"]), + np.array([0.8, 0.9]), + None, + ), + ), + ( + pytest.raises(ValueError), + """```json + [ + {"box_2d": [10, 20, 110, 120], "label": "cat"} + ] + ```""", + (0, 480), + None, + None, + ), + ( + pytest.raises(ValueError), + """```json + [ + {"box_2d": [10, 20, 110, 120], "label": "cat"} + ] + ```""", + (640, -100), + None, + None, + ), + ( + does_not_raise(), + """```json + [ + {"box_2d": [10, 20, 110, 120], "mask": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAAAAACoWZBhAAAADElEQVR4nGNgoCcAAABuAAFIXXpjAAAAAElFTkSuQmCC", "label": "cat"} + ] + ```""", # noqa E501 // docs + (10, 10), + ["cat"], + ( + np.array([[0.2, 0.1, 1.2, 1.1]]), + np.array([0]), + np.array(["cat"]), + None, + np.array([np.zeros((10, 10), dtype=bool)]), + ), + ), + ( + does_not_raise(), + """```json + [ + {"box_2d": [100, 100, 200, 200], "mask": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAAAAACoWZBhAAAADElEQVR4nGNgoCcAAABuAAFIXXpjAAAAAElFTkSuQmCC", "label": "cat", "confidence": 0.8}, + {"box_2d": [300, 300, 400, 400], "mask": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAAAAACoWZBhAAAADElEQVR4nGNgoCcAAABuAAFIXXpjAAAAAElFTkSuQmCC", "label": "dog", "confidence": 0.9} + ] + ```""", # noqa E501 // docs + (10, 10), + ["cat", "dog"], + ( + np.array([[1.0, 1.0, 2.0, 2.0], [3.0, 3.0, 4.0, 4.0]]), + np.array([0, 1]), + np.array(["cat", "dog"]), + np.array([0.8, 0.9]), + np.array( + [np.zeros((10, 10), dtype=bool), np.zeros((10, 10), dtype=bool)], + ), + ), + ), + ], +) +def test_from_google_gemini_2_5( + exception, + result: str, + resolution_wh: Tuple[int, int], + classes: Optional[List[str]], + expected_results: Optional[ + Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray] + ], +): + with exception: + ( + xyxy, + class_id, + class_name, + confidence, + masks, + ) = from_google_gemini_2_5( + result=result, resolution_wh=resolution_wh, classes=classes + ) + + if expected_results is None: + return + + assert xyxy.shape == expected_results[0].shape + assert np.allclose(xyxy, expected_results[0]) + + assert class_id.shape == expected_results[1].shape + assert np.array_equal(class_id, expected_results[1]) + + assert class_name.shape == expected_results[2].shape + assert np.array_equal(class_name, expected_results[2]) + + if confidence is None: + assert expected_results[3] is None + else: + assert expected_results[3] is not None + assert confidence.shape == expected_results[3].shape + assert np.allclose(confidence, expected_results[3]) + + if masks is None: + assert expected_results[4] is None + else: + assert masks is not None + assert masks.shape == expected_results[4].shape + assert np.array_equal(masks, expected_results[4])