Spaces:
Sleeping
Sleeping
| 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)) | |