diff --git a/libraries/getitrack/src/getitrack/adapters/__init__.py b/libraries/getitrack/src/getitrack/adapters/__init__.py new file mode 100644 index 00000000000..2dfae290ba7 --- /dev/null +++ b/libraries/getitrack/src/getitrack/adapters/__init__.py @@ -0,0 +1,13 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +"""Adapters between third-party detector frameworks and getitrack types. + +Each adapter is a `DetectionAdapter` subclass wrapping one detector +instance. Framework imports stay lazy or duck-typed, so no adapter adds +a hard dependency to getitrack. +""" + +from getitrack.adapters.base import DetectionAdapter +from getitrack.adapters.geti import GetiAdapter + +__all__ = ["DetectionAdapter", "GetiAdapter"] diff --git a/libraries/getitrack/src/getitrack/adapters/base.py b/libraries/getitrack/src/getitrack/adapters/base.py new file mode 100644 index 00000000000..de540d6bf32 --- /dev/null +++ b/libraries/getitrack/src/getitrack/adapters/base.py @@ -0,0 +1,48 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +"""Adapter interface between detector frameworks and getitrack. + +A `DetectionAdapter` owns one detector instance and translates between +raw BGR frames and getitrack `Detections`, so trackers and pipelines +consume every framework through the same two members: `detect` and +`class_names`. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import numpy as np + + from getitrack.core.detection import Detections + + +class DetectionAdapter(ABC): + """Wraps one detector behind a framework-agnostic interface. + + Concrete adapters keep their framework imports lazy or duck-typed so + getitrack stays installable without the framework. + """ + + @abstractmethod + def detect(self, frame_bgr: np.ndarray, frame_id: int) -> Detections: + """Run the detector on one BGR frame. + + Args: + frame_bgr: ``(H, W, 3)`` uint8 frame in BGR order. + frame_id: Frame index to stamp on the returned `Detections`. + + Returns: + `Detections` in original frame coordinates. + """ + + @property + def class_names(self) -> dict[int, str] | None: + """Class-id-to-name mapping carried by the detector, if any. + + Returns None when the framework does not expose meaningful names, + in which case callers supply their own table + """ + return None diff --git a/libraries/getitrack/src/getitrack/adapters/geti.py b/libraries/getitrack/src/getitrack/adapters/geti.py new file mode 100644 index 00000000000..e20760fcf83 --- /dev/null +++ b/libraries/getitrack/src/getitrack/adapters/geti.py @@ -0,0 +1,175 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +"""Adapter for getitune detection models. + +getitune and torch are imported lazily inside the methods that need them, so +getitrack imports without either installed. +""" + +from __future__ import annotations + +import re +from typing import TYPE_CHECKING, Literal + +import cv2 +import numpy as np + +from getitrack.adapters.base import DetectionAdapter +from getitrack.core.detection import Detections + +if TYPE_CHECKING: + import torch + from getitune.backend.lightning.models.detection.base import LightningDetectionModel + from getitune.config.data import IntensityConfig + from getitune.data.entity import PredictionBatch, SampleBatch + +# Matches getitune placeholder class names like "label_0", "label_1". +_PLACEHOLDER_NAME = re.compile(r"^label_\d+$") + +# Raw-value divisor per storage dtype for ``scale_to_unit`` intensity mapping. +_DTYPE_MAX_VALUE = {"uint8": 255.0, "uint16": 65535.0, "int16": 32767.0, "float32": 1.0} + + +def _intensity_scale(intensity_config: IntensityConfig | None) -> float: + """Return the divisor mapping raw pixels to ``[0, 1]`` for the model's inputs. + + Reads the getitune ``intensity_config``; defaults to uint8 (255) when unset. + Only ``scale_to_unit`` mode is supported, with the raw max taken from + ``max_value`` or the storage dtype. + """ + if intensity_config is None: + return 255.0 + if intensity_config.mode != "scale_to_unit": + msg = f"GetiAdapter supports intensity mode 'scale_to_unit', got '{intensity_config.mode}'" + raise NotImplementedError(msg) + if intensity_config.max_value is not None: + return float(intensity_config.max_value) + return _DTYPE_MAX_VALUE.get(intensity_config.storage_dtype, 255.0) + + +class GetiAdapter(DetectionAdapter): + """Runs a getitune detection model on raw BGR frames. + + Wraps the preprocess, ``predict_step``, and postprocess round trip + of a getitune Lightning detection model (RF-DETR, YOLOX, ...) for + inference outside the getitune Trainer. + + Example: + >>> adapter = GetiAdapter(model, device="cuda") + >>> detections = adapter.detect(frame, frame_id=0) + """ + + def __init__(self, model: LightningDetectionModel, device: Literal["cpu", "cuda", "mps"] = "cpu") -> None: + """Wrap a getitune detection model. + + Args: + model: getitune detection model in eval mode, exposing + ``data_input_params``, ``predict_step``, and ``label_info``. + device: Torch device the model lives on. + """ + self.model = model + self.device = device + # (mean, std) normalization tensors, built lazily on the first frame and + # reused since the model's mean/std are fixed. None keeps torch out of __init__. + self._norm: tuple[torch.Tensor, torch.Tensor] | None = None + + @property + def class_names(self) -> dict[int, str] | None: + """Class names read from the model's ``label_info``. + + getitune models trained or loaded against a dataset carry real + names; models built from a bare class count carry generated + ``label_N`` placeholders, for which this returns None so callers + fall back to their own table. + """ + names = getattr(getattr(self.model, "label_info", None), "label_names", None) + if not names or all(_PLACEHOLDER_NAME.match(name) for name in names): + return None + return dict(enumerate(names)) + + def detect(self, frame_bgr: np.ndarray, frame_id: int) -> Detections: + """Run one BGR frame through the model. + + Composes `preprocess`, the model's ``predict_step``, and + `to_detections` into the full frame-to-detections round trip. + Requires getitune and torch. + + Args: + frame_bgr: ``(H, W, 3)`` uint8 frame in BGR order. + frame_id: Frame index to stamp on the returned `Detections`. + + Returns: + `Detections` for the frame, in original frame coordinates. + """ + import torch + + with torch.no_grad(): + preds = self.model.predict_step(self.preprocess(frame_bgr), batch_idx=0) + return self.to_detections(preds, frame_id=frame_id) + + def preprocess(self, frame_bgr: np.ndarray) -> SampleBatch: + """Preprocess a BGR frame into a getitune ``SampleBatch``. + + Resizes to the model's input size, maps raw pixels to ``[0, 1]`` using + the model's ``intensity_config`` (uint8 or uint16 inputs), applies the + model's mean/std, and sets ``scale_factor`` so predicted boxes map back + to the original frame coordinates. Requires getitune and torch. + + Args: + frame_bgr: ``(H, W, 3)`` uint8 or uint16 frame in BGR order, e.g. + from ``cv2.VideoCapture``. + + Returns: + A single-image getitune ``SampleBatch`` on ``self.device``. + """ + import torch + from getitune.data.entity import ImageInfo, SampleBatch + + params = self.model.data_input_params + inp_h, inp_w = params.input_size + ori_h, ori_w = frame_bgr.shape[:2] + + if self._norm is None: + self._norm = (torch.tensor(params.mean).view(3, 1, 1), torch.tensor(params.std).view(3, 1, 1)) + mean_t, std_t = self._norm + + rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) + resized = cv2.resize(rgb, (inp_w, inp_h)).astype(np.float32) + tensor = torch.from_numpy(resized).permute(2, 0, 1) + tensor = tensor / _intensity_scale(params.intensity_config) + tensor = (tensor - mean_t) / std_t + + img_info = ImageInfo( + img_idx=0, + img_shape=(inp_h, inp_w), + ori_shape=(ori_h, ori_w), + scale_factor=(inp_h / ori_h, inp_w / ori_w), + ) + return SampleBatch(images=tensor.unsqueeze(0).to(self.device), imgs_info=[img_info]) + + @staticmethod + def to_detections(batch: PredictionBatch, frame_id: int, image_index: int = 0) -> Detections: + """Convert one image's predictions from a getitune ``PredictionBatch``. + + Args: + batch: A getitune ``PredictionBatch`` with per-image ``bboxes``, + ``scores``, and ``labels`` torch tensors. + frame_id: Frame index to stamp on the returned `Detections`. + image_index: Which image of the batch to convert. + + Returns: + `Detections` with float32 xyxy boxes, float32 scores, and int64 + class ids. + + Raises: + ValueError: If the batch is missing bboxes, scores, or labels. + """ + if batch.bboxes is None or batch.scores is None or batch.labels is None: + msg = "PredictionBatch must carry bboxes, scores, and labels" + raise ValueError(msg) + return Detections( + bboxes=batch.bboxes[image_index].detach().cpu().numpy().reshape(-1, 4).astype(np.float32), + scores=batch.scores[image_index].detach().cpu().numpy().reshape(-1).astype(np.float32), + class_ids=batch.labels[image_index].detach().cpu().numpy().reshape(-1).astype(np.int64), + frame_id=frame_id, + ) diff --git a/libraries/getitrack/src/getitrack/io/__init__.py b/libraries/getitrack/src/getitrack/io/__init__.py new file mode 100644 index 00000000000..92b49a732c6 --- /dev/null +++ b/libraries/getitrack/src/getitrack/io/__init__.py @@ -0,0 +1,7 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +"""Video input and output built on OpenCV.""" + +from getitrack.io.video import VideoReader, VideoWriter + +__all__ = ["VideoReader", "VideoWriter"] diff --git a/libraries/getitrack/src/getitrack/io/video.py b/libraries/getitrack/src/getitrack/io/video.py new file mode 100644 index 00000000000..c5213200eee --- /dev/null +++ b/libraries/getitrack/src/getitrack/io/video.py @@ -0,0 +1,168 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +"""OpenCV-backed video reader and writer. + +`VideoReader` iterates BGR uint8 frames from a video file. `VideoWriter` +writes BGR uint8 frames to a video file, creating parent directories on +demand. Both are context managers and release their OpenCV handles on close. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING + +import cv2 + +if TYPE_CHECKING: + from collections.abc import Iterator + from types import TracebackType + + import numpy as np + +_FRAME_CHANNELS = 3 + + +class VideoReader: + """Sequential frame reader over a video file. + + Iterating yields BGR uint8 arrays of shape ``(height, width, 3)`` in + decode order. Metadata (fps, frame size, frame count) is exposed as + properties read from the container header. + + Example: + >>> with VideoReader("input.mp4") as reader: + ... for frame in reader: + ... process(frame) + """ + + def __init__(self, path: str | Path) -> None: + self.path = Path(path) + if not self.path.is_file(): + msg = f"video file not found: {self.path}" + raise FileNotFoundError(msg) + self._cap = cv2.VideoCapture(str(self.path)) + if not self._cap.isOpened(): + self._cap.release() + msg = f"OpenCV could not open video: {self.path}" + raise ValueError(msg) + + @property + def fps(self) -> float: + """Frames per second reported by the container.""" + return float(self._cap.get(cv2.CAP_PROP_FPS)) + + @property + def width(self) -> int: + """Frame width in pixels.""" + return int(self._cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + + @property + def height(self) -> int: + """Frame height in pixels.""" + return int(self._cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + @property + def frame_count(self) -> int: + """Number of frames reported by the container header. + + Some containers report an estimate; the true count is the number + of frames actually yielded by iteration. + """ + return int(self._cap.get(cv2.CAP_PROP_FRAME_COUNT)) + + def __iter__(self) -> Iterator[np.ndarray]: + while True: + ok, frame = self._cap.read() + if not ok: + return + yield frame + + def close(self) -> None: + """Release the underlying OpenCV capture handle.""" + self._cap.release() + + def __enter__(self) -> VideoReader: + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + self.close() + + +class VideoWriter: + """Video writer for annotated tracking output. + + Frames must be BGR uint8 arrays whose size matches ``frame_size``. + Parent directories of ``path`` are created automatically. + + Example: + >>> with VideoWriter("out.mp4", fps=30.0, frame_size=(640, 480)) as writer: + ... writer.write(frame) + """ + + def __init__( + self, + path: str | Path, + fps: float, + frame_size: tuple[int, int], + codec: str = "mp4v", + ) -> None: + """Open a video file for writing. + + Args: + path: Destination video path. + fps: Output frame rate. + frame_size: ``(width, height)`` of every frame. + codec: FourCC codec identifier. + """ + self.path = Path(path) + self.path.parent.mkdir(parents=True, exist_ok=True) + self._frame_size = frame_size + fourcc = cv2.VideoWriter.fourcc(*codec) + self._writer = cv2.VideoWriter(str(self.path), fourcc, fps, frame_size) + if not self._writer.isOpened(): + self._writer.release() + msg = f"OpenCV could not open video for writing: {self.path} (codec '{codec}')" + raise ValueError(msg) + self.frames_written = 0 + + def write(self, frame: np.ndarray) -> None: + """Append one BGR uint8 frame. + + Args: + frame: ``(height, width, 3)`` uint8 array matching the + ``frame_size`` given at construction. + + Raises: + ValueError: If the frame dtype is not uint8 or its shape does not + match ``frame_size``. + """ + if frame.dtype != "uint8": + msg = f"frame dtype {frame.dtype} does not match expected uint8" + raise ValueError(msg) + expected = (self._frame_size[1], self._frame_size[0], _FRAME_CHANNELS) + if frame.shape != expected: + msg = f"frame shape {frame.shape} does not match expected {expected}" + raise ValueError(msg) + self._writer.write(frame) + self.frames_written += 1 + + def close(self) -> None: + """Finalise the container and release the OpenCV writer handle.""" + self._writer.release() + + def __enter__(self) -> VideoWriter: + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + self.close() diff --git a/libraries/getitrack/src/getitrack/visualization/__init__.py b/libraries/getitrack/src/getitrack/visualization/__init__.py new file mode 100644 index 00000000000..c9898b85190 --- /dev/null +++ b/libraries/getitrack/src/getitrack/visualization/__init__.py @@ -0,0 +1,7 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +"""Drawing utilities for tracker output.""" + +from getitrack.visualization.annotator import TrackAnnotator, VideoAnnotator, color_for_track + +__all__ = ["TrackAnnotator", "VideoAnnotator", "color_for_track"] diff --git a/libraries/getitrack/src/getitrack/visualization/annotator.py b/libraries/getitrack/src/getitrack/visualization/annotator.py new file mode 100644 index 00000000000..3cf5ea4fbbd --- /dev/null +++ b/libraries/getitrack/src/getitrack/visualization/annotator.py @@ -0,0 +1,207 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +"""Bounding-box and track-id overlay drawing. + +`TrackAnnotator` draws one frame's `TrackedDetections` onto a BGR image. +Colors are deterministic per track id (golden-ratio hue stepping), so the +same track keeps the same color across frames and across runs. +`VideoAnnotator` composes a `TrackAnnotator` with a `VideoWriter` for the +common case of writing the annotated frames straight to a video file. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import TYPE_CHECKING + +import cv2 +import numpy as np + +from getitrack.io import VideoWriter + +if TYPE_CHECKING: + from collections.abc import Sequence + from pathlib import Path + from types import TracebackType + + from getitrack.core.detection import TrackedDetections + +_GOLDEN_RATIO_CONJUGATE = 0.61803398875 +_HUE_MAX = 179 +_LABEL_PADDING = 3 +_LUMINANCE_THRESHOLD = 140.0 + + +def color_for_track(track_id: int) -> tuple[int, int, int]: + """Return a deterministic BGR color for a track id. + + Hues are stepped by the golden-ratio conjugate so consecutive ids get + visually distant colors. + + Args: + track_id: Stable track identifier. + + Returns: + ``(B, G, R)`` tuple with components in ``[0, 255]``. + """ + hue = (track_id * _GOLDEN_RATIO_CONJUGATE) % 1.0 + hsv = np.array([[[int(hue * _HUE_MAX), 200, 255]]], dtype=np.uint8) + b, g, r = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)[0, 0] + return int(b), int(g), int(r) + + +class TrackAnnotator: + """Draws tracked bounding boxes, ids, class names, and scores onto frames. + + Labels render as `` # ``; the class part appears only + when ``class_names`` is provided, the score only when ``show_score`` is + true. + + Attributes: + thickness: Box outline thickness in pixels. + font_scale: OpenCV font scale for labels. + show_score: Append the detection score to each label when true. + class_names: Optional class-id-to-name lookup, either a sequence + indexed by class id or a (possibly sparse) mapping of class id + to name. Unknown class ids fall back to the numeric id. + """ + + def __init__( + self, + thickness: int = 2, + font_scale: float = 0.4, + show_score: bool = True, + class_names: Sequence[str] | Mapping[int, str] | None = None, + ) -> None: + self.thickness = thickness + self.font_scale = font_scale + self.show_score = show_score + self.class_names = class_names + + def annotate(self, frame: np.ndarray, tracked: TrackedDetections) -> np.ndarray: + """Return a copy of ``frame`` with all tracks drawn on it. + + Args: + frame: ``(H, W, 3)`` BGR uint8 image. Not modified. + tracked: One frame's tracker output. + + Returns: + Annotated copy of the input frame. + """ + out = frame.copy() + for bbox, track_id, score, class_id in zip( + tracked.bboxes, + tracked.track_ids, + tracked.scores, + tracked.class_ids, + strict=True, + ): + color = color_for_track(int(track_id)) + x1, y1, x2, y2 = (round(float(v)) for v in bbox) + cv2.rectangle(out, (x1, y1), (x2, y2), color, self.thickness) + label = self._label(int(track_id), float(score), int(class_id)) + self._draw_label(out, label, x1, y1, color) + return out + + def _label(self, track_id: int, score: float, class_id: int) -> str: + parts = [] + if isinstance(self.class_names, Mapping): + parts.append(self.class_names.get(class_id, str(class_id))) + elif self.class_names is not None: + parts.append(self.class_names[class_id] if 0 <= class_id < len(self.class_names) else str(class_id)) + parts.append(f"#{track_id}") + if self.show_score: + parts.append(f"{score:.2f}") + return " ".join(parts) + + def _draw_label(self, image: np.ndarray, label: str, x: int, y: int, color: tuple[int, int, int]) -> None: + (text_w, text_h), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, self.font_scale, 1) + top = max(y - text_h - baseline - _LABEL_PADDING, 0) + bottom_right = (x + text_w + _LABEL_PADDING, top + text_h + baseline + _LABEL_PADDING) + cv2.rectangle(image, (x, top), bottom_right, color, -1) + # Black or white text, whichever contrasts with the box color. + luminance = 0.114 * color[0] + 0.587 * color[1] + 0.299 * color[2] + text_color = (0, 0, 0) if luminance > _LUMINANCE_THRESHOLD else (255, 255, 255) + cv2.putText( + image, + label, + (x + _LABEL_PADDING, top + text_h + _LABEL_PADDING // 2), + cv2.FONT_HERSHEY_SIMPLEX, + self.font_scale, + text_color, + 1, + cv2.LINE_AA, + ) + + +class VideoAnnotator: + """Writes annotated tracking output straight to a video file. + + Composes a `TrackAnnotator` with a `VideoWriter` for the common case. + Use the two pieces directly when the annotated frames are needed for + anything other than a video file (live preview, image dumps). + + Example: + >>> with VideoAnnotator("out.mp4", fps=30.0, frame_size=(640, 480), class_names=names) as out: + ... out.write(frame, tracked) + """ + + def __init__( + self, + path: str | Path, + fps: float, + frame_size: tuple[int, int], + codec: str = "mp4v", + thickness: int = 2, + font_scale: float = 0.4, + show_score: bool = True, + class_names: Sequence[str] | Mapping[int, str] | None = None, + ) -> None: + """Open a video file for annotated writing. + + Args: + path: Destination video path. + fps: Output frame rate. + frame_size: ``(width, height)`` of every frame. + codec: FourCC codec identifier. + thickness: Box outline thickness in pixels. + font_scale: OpenCV font scale for labels. + show_score: Append the detection score to each label when true. + class_names: Optional class-id-to-name lookup for labels. + """ + self._annotator = TrackAnnotator( + thickness=thickness, + font_scale=font_scale, + show_score=show_score, + class_names=class_names, + ) + self._writer = VideoWriter(path, fps=fps, frame_size=frame_size, codec=codec) + + @property + def path(self) -> Path: + """Destination video path.""" + return self._writer.path + + @property + def frames_written(self) -> int: + """Number of frames written so far.""" + return self._writer.frames_written + + def write(self, frame: np.ndarray, tracked: TrackedDetections) -> None: + """Annotate one frame and append it to the video.""" + self._writer.write(self._annotator.annotate(frame, tracked)) + + def close(self) -> None: + """Finalise the container and release the writer handle.""" + self._writer.close() + + def __enter__(self) -> VideoAnnotator: + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + self.close() diff --git a/libraries/getitrack/tests/unit/test_adapters.py b/libraries/getitrack/tests/unit/test_adapters.py new file mode 100644 index 00000000000..6d2af79432b --- /dev/null +++ b/libraries/getitrack/tests/unit/test_adapters.py @@ -0,0 +1,156 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +"""Tests for getitrack.adapters.""" + +import importlib.util +from types import SimpleNamespace + +import numpy as np +import pytest + +from getitrack.adapters import DetectionAdapter, GetiAdapter + +_NEEDS_GETITUNE = pytest.mark.skipif( + importlib.util.find_spec("torch") is None or importlib.util.find_spec("getitune") is None, + reason="torch and getitune not installed", +) + + +@_NEEDS_GETITUNE +class TestToDetections: + """Requires torch; ``to_detections`` converts ``PredictionBatch`` tensors.""" + + def _batch(self, n=2) -> SimpleNamespace: + import torch + + return SimpleNamespace( + bboxes=[torch.tensor([[10.0, 10.0, 50.0, 50.0], [60.0, 60.0, 90.0, 90.0]][:n])], + scores=[torch.tensor([0.9, 0.4][:n])], + labels=[torch.tensor([3, 8][:n])], + ) + + def test_converts_tensors(self): + d = GetiAdapter.to_detections(self._batch(), frame_id=7) + assert d.frame_id == 7 + assert d.bboxes.dtype == np.float32 + assert d.scores.dtype == np.float32 + assert d.class_ids.dtype == np.int64 + assert d.bboxes.shape == (2, 4) + assert d.class_ids.tolist() == [3, 8] + + def test_empty_image(self): + import torch + + batch = SimpleNamespace( + bboxes=[torch.empty((0, 4))], + scores=[torch.empty((0,))], + labels=[torch.empty((0,))], + ) + d = GetiAdapter.to_detections(batch, frame_id=1) + assert len(d) == 0 + + def test_missing_fields_raise(self): + batch = SimpleNamespace(bboxes=None, scores=None, labels=None) + with pytest.raises(ValueError, match="must carry"): + GetiAdapter.to_detections(batch, frame_id=0) + + def test_image_index_selects_batch_element(self): + import torch + + batch = SimpleNamespace( + bboxes=[torch.zeros((1, 4)), torch.tensor([[5.0, 5.0, 9.0, 9.0]])], + scores=[torch.tensor([0.5]), torch.tensor([0.7])], + labels=[torch.tensor([0]), torch.tensor([2])], + ) + d = GetiAdapter.to_detections(batch, frame_id=0, image_index=1) + assert d.class_ids.tolist() == [2] + assert d.scores.tolist() == pytest.approx([0.7]) + + +class TestClassNames: + def test_is_detection_adapter(self): + assert issubclass(GetiAdapter, DetectionAdapter) + + def test_real_names_become_mapping(self): + model = SimpleNamespace(label_info=SimpleNamespace(label_names=["person", "car"])) + assert GetiAdapter(model).class_names == {0: "person", 1: "car"} + + def test_placeholder_names_return_none(self): + model = SimpleNamespace(label_info=SimpleNamespace(label_names=[f"label_{i}" for i in range(5)])) + assert GetiAdapter(model).class_names is None + + def test_missing_label_info_returns_none(self): + assert GetiAdapter(SimpleNamespace()).class_names is None + + def test_empty_names_return_none(self): + model = SimpleNamespace(label_info=SimpleNamespace(label_names=[])) + assert GetiAdapter(model).class_names is None + + +@_NEEDS_GETITUNE +class TestPreprocess: + """Requires torch + getitune; skipped in getitrack-only environments.""" + + def _adapter( + self, input_size=(64, 96), mean=(0.0, 0.0, 0.0), std=(1.0, 1.0, 1.0), intensity_config=None + ) -> GetiAdapter: + params = SimpleNamespace(input_size=input_size, mean=mean, std=std, intensity_config=intensity_config) + return GetiAdapter(SimpleNamespace(data_input_params=params)) + + def _frame(self, h=480, w=640) -> np.ndarray: + return np.random.default_rng(0).integers(0, 255, size=(h, w, 3), dtype=np.uint8) + + def test_output_shape_and_scale_factor(self): + batch = self._adapter(input_size=(64, 96)).preprocess(self._frame(480, 640)) + assert tuple(batch.images.shape) == (1, 3, 64, 96) + info = batch.imgs_info[0] + assert info.ori_shape == (480, 640) + assert info.scale_factor == pytest.approx((64 / 480, 96 / 640)) + + def test_unit_range_normalization_for_small_means(self): + batch = self._adapter(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)).preprocess(self._frame()) + # (x/255 - 0.5) / 0.5 lies in [-1, 1]. + assert float(batch.images.min()) >= -1.0 + assert float(batch.images.max()) <= 1.0 + + def test_uint16_scaled_by_intensity_config(self): + from getitune.config.data import IntensityConfig + + cfg = IntensityConfig(storage_dtype="uint16", max_value=65535.0) + frame = np.full((10, 10, 3), 30000, dtype=np.uint16) + batch = self._adapter(input_size=(10, 10), intensity_config=cfg).preprocess(frame) + # Raw uint16 value scaled by max_value, not 255. + assert float(batch.images.max()) == pytest.approx(30000 / 65535, abs=1e-4) + + def test_bgr_to_rgb_conversion(self): + frame = np.zeros((10, 10, 3), dtype=np.uint8) + frame[:, :, 0] = 255 # blue channel in BGR + batch = self._adapter(input_size=(10, 10)).preprocess(frame) + # No intensity config -> uint8 scaling, so blue (255/255) becomes 1.0 in channel 2. + assert float(batch.images[0, 2].min()) == 1.0 + assert float(batch.images[0, 0].max()) == 0.0 + + +@_NEEDS_GETITUNE +class TestDetect: + def test_full_round_trip(self): + class _FakeModel: + data_input_params = SimpleNamespace( + input_size=(32, 32), mean=(0.0, 0.0, 0.0), std=(1.0, 1.0, 1.0), intensity_config=None + ) + + def predict_step(self, batch, batch_idx) -> SimpleNamespace: + import torch + + assert tuple(batch.images.shape) == (1, 3, 32, 32) + return SimpleNamespace( + bboxes=[torch.tensor([[1.0, 2.0, 3.0, 4.0]])], + scores=[torch.tensor([0.9])], + labels=[torch.tensor([1])], + ) + + frame = np.zeros((48, 64, 3), dtype=np.uint8) + d = GetiAdapter(_FakeModel()).detect(frame, frame_id=11) + assert d.frame_id == 11 + assert len(d) == 1 + assert d.class_ids.tolist() == [1] diff --git a/libraries/getitrack/tests/unit/test_annotator.py b/libraries/getitrack/tests/unit/test_annotator.py new file mode 100644 index 00000000000..ab33037ed55 --- /dev/null +++ b/libraries/getitrack/tests/unit/test_annotator.py @@ -0,0 +1,114 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +"""Tests for getitrack.visualization.annotator.""" + +import numpy as np + +from getitrack.core.detection import TrackedDetections +from getitrack.core.track import TrackState +from getitrack.io import VideoReader +from getitrack.visualization import TrackAnnotator, VideoAnnotator, color_for_track + + +def _frame(w=128, h=96) -> np.ndarray: + return np.full((h, w, 3), 30, dtype=np.uint8) + + +def _tracked(boxes, frame_id=1) -> TrackedDetections: + n = len(boxes) + return TrackedDetections( + bboxes=np.asarray(boxes, dtype=np.float32), + scores=np.full(n, 0.9, dtype=np.float32), + class_ids=np.zeros(n, dtype=np.int64), + track_ids=np.arange(1, n + 1, dtype=np.int64), + track_states=np.full(n, TrackState.ACTIVE, dtype=np.int8), + frame_id=frame_id, + ) + + +class TestTrackAnnotator: + def test_returns_modified_copy(self): + frame = _frame() + original = frame.copy() + out = TrackAnnotator().annotate(frame, _tracked([[20, 20, 60, 60]])) + assert out is not frame + assert np.array_equal(frame, original) + assert not np.array_equal(out, frame) + assert out.shape == frame.shape + assert out.dtype == np.uint8 + + def test_empty_tracks_changes_nothing(self): + frame = _frame() + out = TrackAnnotator().annotate(frame, TrackedDetections.create_empty(frame_id=1)) + assert np.array_equal(out, frame) + + def test_show_score_false(self): + out = TrackAnnotator(show_score=False).annotate(_frame(), _tracked([[10, 10, 40, 40]])) + assert not np.array_equal(out, _frame()) + + def test_multiple_tracks(self): + out = TrackAnnotator().annotate(_frame(), _tracked([[5, 5, 30, 30], [60, 40, 100, 80]])) + assert not np.array_equal(out, _frame()) + + +class TestColorForTrack: + def test_deterministic(self): + assert color_for_track(7) == color_for_track(7) + + def test_distinct_for_consecutive_ids(self): + colors = {color_for_track(i) for i in range(1, 11)} + assert len(colors) == 10 + + def test_valid_bgr_range(self): + for i in range(1, 50): + assert all(0 <= c <= 255 for c in color_for_track(i)) + + +class TestClassNames: + def test_label_includes_class_name(self): + ann = TrackAnnotator(class_names=["background", "person", "car"]) + assert ann._label(7, 0.86, 1) == "person #7 0.86" + + def test_unknown_class_id_falls_back_to_numeric(self): + ann = TrackAnnotator(class_names=["background"]) + assert ann._label(3, 0.5, 42) == "42 #3 0.50" + + def test_no_class_names_keeps_id_only_label(self): + assert TrackAnnotator(show_score=False)._label(5, 0.9, 1) == "#5" + + def test_annotate_with_class_names_draws(self): + ann = TrackAnnotator(class_names=["background", "person"]) + out = ann.annotate(_frame(), _tracked([[10, 10, 60, 60]])) + assert not np.array_equal(out, _frame()) + + +class TestVideoAnnotator: + def test_writes_annotated_frames(self, tmp_path): + path = tmp_path / "out.mp4" + with VideoAnnotator(path, fps=30.0, frame_size=(128, 96)) as out: + for _ in range(5): + out.write(_frame(), _tracked([[20, 20, 60, 60]])) + assert out.frames_written == 5 + with VideoReader(path) as reader: + frames = list(reader) + assert len(frames) == 5 + # The stored frames must differ from the raw input: boxes were drawn. + assert not np.array_equal(frames[0], _frame()) + + def test_style_params_forwarded(self, tmp_path): + with VideoAnnotator( + tmp_path / "out.mp4", + fps=30.0, + frame_size=(128, 96), + class_names={0: "thing"}, + show_score=False, + thickness=3, + ) as out: + assert out._annotator.class_names == {0: "thing"} + assert out._annotator.show_score is False + assert out._annotator.thickness == 3 + out.write(_frame(), _tracked([[10, 10, 50, 50]])) + + def test_path_property(self, tmp_path): + with VideoAnnotator(tmp_path / "a" / "out.mp4", fps=30.0, frame_size=(128, 96)) as out: + assert out.path == tmp_path / "a" / "out.mp4" diff --git a/libraries/getitrack/tests/unit/test_video_io.py b/libraries/getitrack/tests/unit/test_video_io.py new file mode 100644 index 00000000000..cdab070b63b --- /dev/null +++ b/libraries/getitrack/tests/unit/test_video_io.py @@ -0,0 +1,77 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +"""Tests for getitrack.io.video.""" + +from pathlib import Path + +import numpy as np +import pytest + +from getitrack.io import VideoReader, VideoWriter + +_W, _H = 64, 48 + + +def _solid_frame(value: int) -> np.ndarray: + return np.full((_H, _W, 3), value, dtype=np.uint8) + + +def _write_video(path, n_frames=10, fps=30.0) -> Path: + with VideoWriter(path, fps=fps, frame_size=(_W, _H)) as writer: + for i in range(n_frames): + writer.write(_solid_frame((i * 20) % 255)) + return path + + +class TestVideoReader: + def test_roundtrip_frame_count_and_shape(self, tmp_path): + path = _write_video(tmp_path / "clip.mp4", n_frames=10) + with VideoReader(path) as reader: + frames = list(reader) + assert len(frames) == 10 + assert all(f.shape == (_H, _W, 3) for f in frames) + assert all(f.dtype == np.uint8 for f in frames) + + def test_metadata_properties(self, tmp_path): + path = _write_video(tmp_path / "clip.mp4", n_frames=10, fps=30.0) + with VideoReader(path) as reader: + assert reader.width == _W + assert reader.height == _H + assert reader.frame_count == 10 + assert reader.fps == pytest.approx(30.0, abs=1.0) + + def test_missing_file_raises(self, tmp_path): + with pytest.raises(FileNotFoundError, match="not found"): + VideoReader(tmp_path / "nope.mp4") + + def test_unreadable_file_raises(self, tmp_path): + bogus = tmp_path / "bogus.mp4" + bogus.write_text("this is not a video") + with pytest.raises(ValueError, match="could not open"): + VideoReader(bogus) + + +class TestVideoWriter: + def test_counts_written_frames(self, tmp_path): + with VideoWriter(tmp_path / "out.mp4", fps=30.0, frame_size=(_W, _H)) as writer: + writer.write(_solid_frame(0)) + writer.write(_solid_frame(100)) + assert writer.frames_written == 2 + + def test_wrong_frame_shape_raises(self, tmp_path): + with VideoWriter(tmp_path / "out.mp4", fps=30.0, frame_size=(_W, _H)) as writer: + bad = np.zeros((_H + 1, _W, 3), dtype=np.uint8) + with pytest.raises(ValueError, match="does not match"): + writer.write(bad) + + def test_wrong_frame_dtype_raises(self, tmp_path): + with VideoWriter(tmp_path / "out.mp4", fps=30.0, frame_size=(_W, _H)) as writer: + bad = np.zeros((_H, _W, 3), dtype=np.float32) + with pytest.raises(ValueError, match="dtype"): + writer.write(bad) + + def test_creates_parent_directories(self, tmp_path): + nested = tmp_path / "a" / "b" / "out.mp4" + with VideoWriter(nested, fps=30.0, frame_size=(_W, _H)) as writer: + writer.write(_solid_frame(50)) + assert nested.is_file()