Spaces:
Sleeping
Sleeping
File size: 4,486 Bytes
d725335 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | from __future__ import annotations
import re
from typing import Any, Protocol
from .models import Detection
def parse_class_prompt(class_prompt: str | list[str] | tuple[str, ...]) -> list[str]:
if isinstance(class_prompt, (list, tuple)):
parts = [str(item) for item in class_prompt]
else:
parts = re.split(r"[,;\n]+", class_prompt)
seen: set[str] = set()
classes: list[str] = []
for part in parts:
label = " ".join(part.strip().split())
key = label.lower()
if label and key not in seen:
seen.add(key)
classes.append(label)
return classes
class Detector(Protocol):
class_names: list[str]
def detect(
self,
frame: Any,
*,
frame_index: int,
timestamp_sec: float,
confidence: float,
image_size: int | None = None,
max_detections: int | None = None,
) -> list[Detection]:
...
class UltralyticsYOLOEDetector:
def __init__(
self,
*,
class_names: list[str],
model_name: str = "yoloe-26s-seg.pt",
device: str | None = None,
) -> None:
if not class_names:
raise ValueError("YOLOE needs at least one open-vocabulary class.")
try:
from ultralytics import YOLOE
except ImportError as exc: # pragma: no cover - optional heavy dependency
raise RuntimeError("Install ultralytics to use the YOLOE detector.") from exc
self.class_names = class_names
self.model_name = model_name
self.device = device
self.model = YOLOE(model_name)
self.model.set_classes(class_names)
def detect(
self,
frame: Any,
*,
frame_index: int,
timestamp_sec: float,
confidence: float,
image_size: int | None = None,
max_detections: int | None = None,
) -> list[Detection]:
kwargs: dict[str, Any] = {"conf": confidence, "verbose": False}
if self.device:
kwargs["device"] = self.device
if image_size:
kwargs["imgsz"] = image_size
if max_detections:
kwargs["max_det"] = max_detections
results = self.model.predict(frame, **kwargs)
if not results:
return []
return detections_from_ultralytics_result(
results[0],
frame_shape=frame.shape,
frame_index=frame_index,
timestamp_sec=timestamp_sec,
fallback_names=self.class_names,
)
def detections_from_ultralytics_result(
result: Any,
*,
frame_shape: tuple[int, int, int],
frame_index: int,
timestamp_sec: float,
fallback_names: list[str],
) -> list[Detection]:
boxes = getattr(result, "boxes", None)
if boxes is None or boxes.xyxy is None:
return []
height, width = frame_shape[:2]
xyxy_values = boxes.xyxy.cpu().tolist()
confidences = boxes.conf.cpu().tolist()
class_ids = boxes.cls.cpu().tolist()
names = getattr(result, "names", {}) or {}
detections: list[Detection] = []
for bbox, score, class_id in zip(xyxy_values, confidences, class_ids):
label = _label_from_names(names, int(class_id), fallback_names)
detections.append(
Detection(
frame_index=frame_index,
timestamp_sec=timestamp_sec,
label=label,
confidence=float(score),
bbox_xyxy=tuple(float(value) for value in bbox),
bbox_xyxy_norm=_normalize_box(bbox, width, height),
)
)
return detections
def _label_from_names(names: Any, class_id: int, fallback_names: list[str]) -> str:
if isinstance(names, dict) and class_id in names:
return str(names[class_id])
if isinstance(names, list) and 0 <= class_id < len(names):
return str(names[class_id])
if 0 <= class_id < len(fallback_names):
return fallback_names[class_id]
return f"class_{class_id}"
def _normalize_box(bbox: list[float], width: int, height: int) -> tuple[float, float, float, float]:
x1, y1, x2, y2 = bbox
if width <= 0 or height <= 0:
return (0.0, 0.0, 0.0, 0.0)
return (
_clamp01(float(x1) / width),
_clamp01(float(y1) / height),
_clamp01(float(x2) / width),
_clamp01(float(y2) / height),
)
def _clamp01(value: float) -> float:
return max(0.0, min(1.0, value))
|