xiaoxue_admin
/
supervision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #1877 from roboflow/feat/gemini-2_5_segmentation 

feat: 🚀 add support for Google Gemini 2.5 bounding box and mask parsing for "from_vlm"
This commit is contained in:
Piotr Skalski 2025-07-14 23:04:33 +02:00 committed by GitHub
parent 7b8560016e efe59cb91c
commit daaadaa1f9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 687 additions and 45 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

257
supervision/detection/core.py
View File

@ -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 = "<loc0256><loc0256><loc0768><loc0768> cat"
@ -849,7 +860,7 @@ class Detections:
# {'class_name': array(['cat'], dtype='<U10')}
```
Examples:
!!! example "Qwen2.5-VL"
```python
import supervision as sv
@ -879,7 +890,7 @@ class Detections:
# array([0, 1])
```
Examples:
!!! example "Gemini 2.0"
```python
import supervision as sv
@ -900,8 +911,103 @@ class Detections:
detections.xyxy
# array([[543., 40., 728., 200.], [653., 352., 820., 522.]])
detections.data
# {'class_name': array(['cat', 'dog'], dtype='<U26')}
detections.class_id
# array([0, 1])
```
!!! example "Gemini 2.5"
??? tip "Prompt engineering"
To get the best results from Google Gemini 2.5, use the following prompt.
This prompt is designed to detect all visible objects in the image,
including small, distant, or partially visible ones, and to return
tight bounding boxes.
```
Carefully examine this image and detect ALL visible objects, including
small, distant, or partially visible ones.
IMPORTANT: Focus on finding as many objects as possible, even if you are
only moderately confident.
Make sure each bounding box is as tight as possible.
Valid object classes: {class_list}
For each detected object, provide:
- "label": the exact class name from the list above
- "confidence": your certainty (between 0.0 and 1.0)
- "box_2d": the bounding box [ymin, xmin, ymax, xmax] normalized to 0-1000
- "mask": the binary mask of the object as a base64-encoded string
Detect everything that matches the valid classes. Do not be
conservative; include objects even with moderate confidence.
Return a JSON array, for example:
[
{
"label": "person",
"confidence": 0.95,
"box_2d": [100, 200, 300, 400],
"mask": "..."
},
{
"label": "kite",
"confidence": 0.80,
"box_2d": [50, 150, 250, 350],
"mask": "..."
}
]
```
When using the google-genai library, it is recommended to set
thinking_budget=0 in thinking_config for more direct and faster responses.
```python
from google.generativeai import types
model.generate_content(
...,
generation_config=generation_config,
safety_settings=safety_settings,
thinking_config=types.ThinkingConfig(
thinking_budget=0
)
)
```
For a shorter prompt focused only on segmentation masks, you can use:
```
Return a JSON list of segmentation masks. Each entry should include the
2D bounding box in the "box_2d" key, the segmentation mask in the "mask"
key, and the text label in the "label" key. Use descriptive labels.
```
```python
import supervision as sv
gemini_response_text = \"\"\"```json
[
{"box_2d": [543, 40, 728, 200], "label": "cat", "id": 1},
{"box_2d": [653, 352, 820, 522], "label": "dog", "id": 2}
]
```\"\"\"
detections = sv.Detections.from_lmm(
sv.LMM.GOOGLE_GEMINI_2_5,
gemini_response_text,
resolution_wh=(1000, 1000),
classes=['cat', 'dog'],
)
detections.xyxy
# array([[543., 40., 728., 200.], [653., 352., 820., 522.]])
detections.data
# {'class_name': array(['cat', 'dog'], dtype='<U26')}
@ -910,7 +1016,7 @@ class Detections:
# array([0, 1])
```
Examples:
!!! example "Moondream"
```python
import supervision as sv
@ -931,7 +1037,7 @@ class Detections:
]
}
detections = sv.Detections.from_vmm(
detections = sv.Detections.from_lmm(
sv.LMM.MOONDREAM,
moondream_result,
resolution_wh=(1000, 1000),
@ -941,7 +1047,7 @@ class Detections:
# array([[1752.28, 818.82, 2165.72, 1229.14],
# [1908.01, 1346.67, 2585.99, 2024.11]])
```
"""
""" # noqa: E501
# filler logic mapping old from_lmm to new from_vlm
lmm_to_vlm = {
@ -978,11 +1084,21 @@ class Detections:
cls, vlm: Union[VLM, str], result: Union[str, dict], **kwargs: Any
) -> 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 = "<loc0256><loc0256><loc0768><loc0768> cat"
@ -1015,7 +1132,7 @@ class Detections:
# {'class_name': array(['cat'], dtype='<U10')}
```
Examples:
!!! example "Qwen2.5-VL"
```python
import supervision as sv
@ -1045,7 +1162,7 @@ class Detections:
# array([0, 1])
```
Examples:
!!! example "Gemini 2.0"
```python
import supervision as sv
@ -1066,8 +1183,103 @@ class Detections:
detections.xyxy
# array([[543., 40., 728., 200.], [653., 352., 820., 522.]])
detections.data
# {'class_name': array(['cat', 'dog'], dtype='<U26')}
detections.class_id
# array([0, 1])
```
!!! example "Gemini 2.5"
??? tip "Prompt engineering"
To get the best results from Google Gemini 2.5, use the following prompt.
This prompt is designed to detect all visible objects in the image,
including small, distant, or partially visible ones, and to return
tight bounding boxes.
```
Carefully examine this image and detect ALL visible objects, including
small, distant, or partially visible ones.
IMPORTANT: Focus on finding as many objects as possible, even if you are
only moderately confident.
Make sure each bounding box is as tight as possible.
Valid object classes: {class_list}
For each detected object, provide:
- "label": the exact class name from the list above
- "confidence": your certainty (between 0.0 and 1.0)
- "box_2d": the bounding box [ymin, xmin, ymax, xmax] normalized to 0-1000
- "mask": the binary mask of the object as a base64-encoded string
Detect everything that matches the valid classes. Do not be
conservative; include objects even with moderate confidence.
Return a JSON array, for example:
[
{
"label": "person",
"confidence": 0.95,
"box_2d": [100, 200, 300, 400],
"mask": "..."
},
{
"label": "kite",
"confidence": 0.80,
"box_2d": [50, 150, 250, 350],
"mask": "..."
}
]
```
When using the google-genai library, it is recommended to set
thinking_budget=0 in thinking_config for more direct and faster responses.
```python
from google.generativeai import types
model.generate_content(
...,
generation_config=generation_config,
safety_settings=safety_settings,
thinking_config=types.ThinkingConfig(
thinking_budget=0
)
)
```
For a shorter prompt focused only on segmentation masks, you can use:
```
Return a JSON list of segmentation masks. Each entry should include the
2D bounding box in the "box_2d" key, the segmentation mask in the "mask"
key, and the text label in the "label" key. Use descriptive labels.
```
```python
import supervision as sv
gemini_response_text = \"\"\"```json
[
{"box_2d": [543, 40, 728, 200], "label": "cat", "id": 1},
{"box_2d": [653, 352, 820, 522], "label": "dog", "id": 2}
]
```\"\"\"
detections = sv.Detections.from_vlm(
sv.VLM.GOOGLE_GEMINI_2_5,
gemini_response_text,
resolution_wh=(1000, 1000),
classes=['cat', 'dog'],
)
detections.xyxy
# array([[543., 40., 728., 200.], [653., 352., 820., 522.]])
detections.data
# {'class_name': array(['cat', 'dog'], dtype='<U26')}
@ -1076,7 +1288,7 @@ class Detections:
# array([0, 1])
```
Examples:
!!! example "Moondream"
```python
import supervision as sv
@ -1106,10 +1318,10 @@ class Detections:
detections.xyxy
# array([[1752.28, 818.82, 2165.72, 1229.14],
# [1908.01, 1346.67, 2585.99, 2024.11]])
```
"""
""" # noqa: E501
vlm = validate_vlm_parameters(vlm, result, kwargs)
if vlm == VLM.PALIGEMMA:
@ -1135,8 +1347,8 @@ class Detections:
return cls(xyxy=xyxy, mask=mask, data=data)
if vlm == VLM.GOOGLE_GEMINI_2_0 or vlm == VLM.GOOGLE_GEMINI_2_5:
xyxy, class_id, class_name = from_google_gemini(result, **kwargs)
if vlm == VLM.GOOGLE_GEMINI_2_0:
xyxy, class_id, class_name = from_google_gemini_2_0(result, **kwargs)
data = {CLASS_NAME_DATA_FIELD: class_name}
return cls(xyxy=xyxy, class_id=class_id, data=data)
@ -1144,6 +1356,19 @@ class Detections:
xyxy = from_moondream(result, **kwargs)
return cls(xyxy=xyxy)
if vlm == VLM.GOOGLE_GEMINI_2_5:
xyxy, class_id, class_name, confidence, mask = from_google_gemini_2_5(
result, **kwargs
)
data = {CLASS_NAME_DATA_FIELD: class_name}
return cls(
xyxy=xyxy,
class_id=class_id,
mask=mask,
confidence=confidence,
data=data,
)
return cls.empty()
@classmethod

207
supervision/detection/vlm.py
View File

@ -1,9 +1,12 @@
import base64
import io
import json
import re
from enum import Enum
from typing import Any, Dict, List, Optional, Tuple, Union
import numpy as np
from PIL import Image
from supervision.detection.utils import (
denormalize_boxes,
@ -11,6 +14,7 @@ from supervision.detection.utils import (
polygon_to_xyxy,
)
from supervision.utils.internal import deprecated
from supervision.validators import validate_resolution
@deprecated(
@ -23,6 +27,7 @@ class LMM(Enum):
QWEN_2_5_VL = "qwen_2_5_vl"
GOOGLE_GEMINI_2_0 = "gemini_2_0"
GOOGLE_GEMINI_2_5 = "gemini_2_5"
MOONDREAM = "moondream"
class VLM(Enum):
@ -58,6 +63,7 @@ ALLOWED_ARGUMENTS: Dict[VLM, List[str]] = {
VLM.QWEN_2_5_VL: ["input_wh", "resolution_wh", "classes"],
VLM.GOOGLE_GEMINI_2_0: ["resolution_wh", "classes"],
VLM.GOOGLE_GEMINI_2_5: ["resolution_wh", "classes"],
VLM.MOONDREAM: ["resolution_wh"],
}
SUPPORTED_TASKS_FLORENCE_2 = [
@ -126,11 +132,7 @@ def from_paligemma(
the class labels for each bounding box.
"""
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)
pattern = re.compile(
r"(?<!<loc\d{4}>)<loc(\d{4})><loc(\d{4})><loc(\d{4})><loc(\d{4})> ([\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 = []

25
supervision/validators/__init__.py
View File

@ -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

243
test/detection/test_vlm.py
View File

@ -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])