From 54f5be978b52a66d727de1986e3ced35f496fd42 Mon Sep 17 00:00:00 2001 From: Omkar Kabde Date: Thu, 25 Jun 2026 16:57:36 +0530 Subject: [PATCH 1/9] feat: add OpenCV video reader and writer --- .../getitrack/src/getitrack/io/__init__.py | 7 + libraries/getitrack/src/getitrack/io/video.py | 162 ++++++++++++++++++ .../getitrack/tests/unit/test_video_io.py | 71 ++++++++ 3 files changed, 240 insertions(+) create mode 100644 libraries/getitrack/src/getitrack/io/__init__.py create mode 100644 libraries/getitrack/src/getitrack/io/video.py create mode 100644 libraries/getitrack/tests/unit/test_video_io.py 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..569c8b103fe --- /dev/null +++ b/libraries/getitrack/src/getitrack/io/video.py @@ -0,0 +1,162 @@ +# 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(): + 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(): + 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 shape does not match ``frame_size``. + """ + 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/tests/unit/test_video_io.py b/libraries/getitrack/tests/unit/test_video_io.py new file mode 100644 index 00000000000..042336fa84b --- /dev/null +++ b/libraries/getitrack/tests/unit/test_video_io.py @@ -0,0 +1,71 @@ +# 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_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() From 38af3a5a326099332a5bf30aca0e885c503d4e2c Mon Sep 17 00:00:00 2001 From: Omkar Kabde Date: Thu, 25 Jun 2026 16:57:36 +0530 Subject: [PATCH 2/9] feat: add track and video annotators --- .../src/getitrack/visualization/__init__.py | 7 + .../src/getitrack/visualization/annotator.py | 207 ++++++++++++++++++ .../getitrack/tests/unit/test_annotator.py | 114 ++++++++++ 3 files changed, 328 insertions(+) create mode 100644 libraries/getitrack/src/getitrack/visualization/__init__.py create mode 100644 libraries/getitrack/src/getitrack/visualization/annotator.py create mode 100644 libraries/getitrack/tests/unit/test_annotator.py 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..68f118e14a3 --- /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_annotator.py b/libraries/getitrack/tests/unit/test_annotator.py new file mode 100644 index 00000000000..b1b29f6e89b --- /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" From 2d96f8fefa0a2597aa9097946f8503e68b0f8ab3 Mon Sep 17 00:00:00 2001 From: Omkar Kabde Date: Thu, 25 Jun 2026 16:57:37 +0530 Subject: [PATCH 3/9] feat: add detector adapter interface and getitune adapter --- .../src/getitrack/adapters/__init__.py | 13 ++ .../getitrack/src/getitrack/adapters/base.py | 48 ++++++ .../getitrack/src/getitrack/adapters/geti.py | 158 +++++++++++++++++ .../getitrack/tests/unit/test_adapters.py | 159 ++++++++++++++++++ 4 files changed, 378 insertions(+) create mode 100644 libraries/getitrack/src/getitrack/adapters/__init__.py create mode 100644 libraries/getitrack/src/getitrack/adapters/base.py create mode 100644 libraries/getitrack/src/getitrack/adapters/geti.py create mode 100644 libraries/getitrack/tests/unit/test_adapters.py 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..9395cd599f2 --- /dev/null +++ b/libraries/getitrack/src/getitrack/adapters/geti.py @@ -0,0 +1,158 @@ +# 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. `GetiAdapter.to_detections` is +duck-typed and works on any prediction-batch-shaped object. +""" + +from __future__ import annotations + +import re +from typing import Any + +import cv2 +import numpy as np + +from getitrack.adapters.base import DetectionAdapter +from getitrack.core.detection import Detections + +# Matches getitune placeholder class names like "label_0", "label_1". +_PLACEHOLDER_NAME = re.compile(r"^label_\d+$") + + +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: Any, device: str = "cpu") -> None: # noqa: ANN401 + """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 + + @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) -> Any: # noqa: ANN401 + """Preprocess a BGR frame into a getitune ``SampleBatch``. + + Applies the resize and normalization described by the model's + ``data_input_params`` 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 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 + + inp_h, inp_w = self.model.data_input_params.input_size + mean = self.model.data_input_params.mean + std = self.model.data_input_params.std + ori_h, ori_w = frame_bgr.shape[:2] + + rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) + tensor = torch.from_numpy(cv2.resize(rgb, (inp_w, inp_h))).permute(2, 0, 1).float() + # Mean values below 1.0 indicate the model expects 0-1 normalized pixels + # (e.g. RF-DETR); otherwise it expects the raw 0-255 range (e.g. YOLOX). + if all(m < 1.0 for m in mean): + tensor = tensor / 255.0 + tensor = (tensor - torch.tensor(mean).view(3, 1, 1)) / torch.tensor(std).view(3, 1, 1) + + 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: Any, frame_id: int, image_index: int = 0) -> Detections: # noqa: ANN401 + """Convert one image's predictions from a getitune ``PredictionBatch``. + + Duck-typed: works without getitune or torch installed, on any + object with per-image ``bboxes``, ``scores``, and ``labels`` lists. + + Args: + batch: A getitune ``PredictionBatch`` with per-image ``bboxes``, + ``scores``, and ``labels`` lists (torch tensors or arrays). + 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=_to_numpy(batch.bboxes[image_index]).reshape(-1, 4).astype(np.float32), + scores=_to_numpy(batch.scores[image_index]).reshape(-1).astype(np.float32), + class_ids=_to_numpy(batch.labels[image_index]).reshape(-1).astype(np.int64), + frame_id=frame_id, + ) + + +def _to_numpy(value: Any) -> np.ndarray: # noqa: ANN401 + """Convert a torch tensor (possibly on an accelerator) or array-like to numpy.""" + if hasattr(value, "cpu"): + value = value.cpu() + if hasattr(value, "numpy"): + value = value.numpy() + return np.asarray(value) diff --git a/libraries/getitrack/tests/unit/test_adapters.py b/libraries/getitrack/tests/unit/test_adapters.py new file mode 100644 index 00000000000..f7294a23ca4 --- /dev/null +++ b/libraries/getitrack/tests/unit/test_adapters.py @@ -0,0 +1,159 @@ +# 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", +) + + +class _FakeTensor: + """Mimics a torch tensor on an accelerator: needs .cpu().numpy().""" + + def __init__(self, data): + self._data = np.asarray(data) + self.moved = False + + def cpu(self) -> "_FakeTensor": + self.moved = True + return self + + def numpy(self) -> np.ndarray: + assert self.moved, "numpy() requires cpu() first, like a real cuda/mps tensor" + return self._data + + +def _batch(n=2, wrap=_FakeTensor) -> SimpleNamespace: + return SimpleNamespace( + bboxes=[wrap([[10.0, 10.0, 50.0, 50.0], [60.0, 60.0, 90.0, 90.0]][:n])], + scores=[wrap([0.9, 0.4][:n])], + labels=[wrap([3, 8][:n])], + ) + + +class TestToDetections: + def test_converts_fake_tensors(self): + d = GetiAdapter.to_detections(_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_converts_plain_numpy(self): + d = GetiAdapter.to_detections(_batch(wrap=np.asarray), frame_id=0) + assert len(d) == 2 + assert d.scores.tolist() == pytest.approx([0.9, 0.4]) + + def test_empty_image(self): + batch = SimpleNamespace( + bboxes=[np.empty((0, 4))], + scores=[np.empty((0,))], + labels=[np.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): + batch = SimpleNamespace( + bboxes=[np.zeros((1, 4)), np.array([[5.0, 5.0, 9.0, 9.0]])], + scores=[np.array([0.5]), np.array([0.7])], + labels=[np.array([0]), np.array([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)) -> GetiAdapter: + model = SimpleNamespace(data_input_params=SimpleNamespace(input_size=input_size, mean=mean, std=std)) + return GetiAdapter(model) + + 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_raw_pixel_range_for_large_means(self): + batch = self._adapter(mean=(123.0, 117.0, 104.0), std=(58.0, 57.0, 57.0)).preprocess(self._frame()) + # Raw 0-255 pixels minus a ~100 mean over a ~57 std exceeds 1.5. + assert float(batch.images.max()) > 1.5 + + 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) + # mean 0.0 < 1.0 triggers the 0-1 branch, so blue 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)) + + def predict_step(self, batch, batch_idx) -> SimpleNamespace: + assert tuple(batch.images.shape) == (1, 3, 32, 32) + return SimpleNamespace( + bboxes=[np.array([[1.0, 2.0, 3.0, 4.0]])], + scores=[np.array([0.9])], + labels=[np.array([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] From 40c143c87f1ad09a5330047072ec4dae3c1b0ca1 Mon Sep 17 00:00:00 2001 From: Omkar Kabde Date: Thu, 2 Jul 2026 12:09:23 +0530 Subject: [PATCH 4/9] make annotator private --- .../getitrack/src/getitrack/visualization/annotator.py | 4 ++-- libraries/getitrack/tests/unit/test_annotator.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/libraries/getitrack/src/getitrack/visualization/annotator.py b/libraries/getitrack/src/getitrack/visualization/annotator.py index 68f118e14a3..3cf5ea4fbbd 100644 --- a/libraries/getitrack/src/getitrack/visualization/annotator.py +++ b/libraries/getitrack/src/getitrack/visualization/annotator.py @@ -169,7 +169,7 @@ def __init__( show_score: Append the detection score to each label when true. class_names: Optional class-id-to-name lookup for labels. """ - self.annotator = TrackAnnotator( + self._annotator = TrackAnnotator( thickness=thickness, font_scale=font_scale, show_score=show_score, @@ -189,7 +189,7 @@ def frames_written(self) -> int: 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)) + self._writer.write(self._annotator.annotate(frame, tracked)) def close(self) -> None: """Finalise the container and release the writer handle.""" diff --git a/libraries/getitrack/tests/unit/test_annotator.py b/libraries/getitrack/tests/unit/test_annotator.py index b1b29f6e89b..ab33037ed55 100644 --- a/libraries/getitrack/tests/unit/test_annotator.py +++ b/libraries/getitrack/tests/unit/test_annotator.py @@ -104,9 +104,9 @@ def test_style_params_forwarded(self, tmp_path): 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 + 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): From 3d9109e3a37f0745d79bcd704e64849ddbb26009 Mon Sep 17 00:00:00 2001 From: Omkar Kabde Date: Thu, 2 Jul 2026 12:11:10 +0530 Subject: [PATCH 5/9] fix: release VideoCapture handle when VideoReader fails to open --- libraries/getitrack/src/getitrack/io/video.py | 1 + 1 file changed, 1 insertion(+) diff --git a/libraries/getitrack/src/getitrack/io/video.py b/libraries/getitrack/src/getitrack/io/video.py index 569c8b103fe..0c7dc674637 100644 --- a/libraries/getitrack/src/getitrack/io/video.py +++ b/libraries/getitrack/src/getitrack/io/video.py @@ -43,6 +43,7 @@ def __init__(self, path: str | Path) -> None: 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) From a3c27e30a0eb45d5731d2b011c959d9760262da3 Mon Sep 17 00:00:00 2001 From: Omkar Kabde Date: Thu, 2 Jul 2026 12:23:17 +0530 Subject: [PATCH 6/9] fix: release VideoWriter handle when it fails to open. --- libraries/getitrack/src/getitrack/io/video.py | 1 + 1 file changed, 1 insertion(+) diff --git a/libraries/getitrack/src/getitrack/io/video.py b/libraries/getitrack/src/getitrack/io/video.py index 0c7dc674637..7bf343e61c5 100644 --- a/libraries/getitrack/src/getitrack/io/video.py +++ b/libraries/getitrack/src/getitrack/io/video.py @@ -126,6 +126,7 @@ def __init__( 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 From 5805500ed1683e9d53a4eafb896c6225d6973493 Mon Sep 17 00:00:00 2001 From: Omkar Kabde Date: Thu, 2 Jul 2026 12:30:51 +0530 Subject: [PATCH 7/9] improve type hints for geti Adapter --- .../getitrack/src/getitrack/adapters/geti.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/libraries/getitrack/src/getitrack/adapters/geti.py b/libraries/getitrack/src/getitrack/adapters/geti.py index 9395cd599f2..f5fcbfe836e 100644 --- a/libraries/getitrack/src/getitrack/adapters/geti.py +++ b/libraries/getitrack/src/getitrack/adapters/geti.py @@ -10,7 +10,7 @@ from __future__ import annotations import re -from typing import Any +from typing import TYPE_CHECKING, Literal import cv2 import numpy as np @@ -18,6 +18,11 @@ 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.data.entity import PredictionBatch, SampleBatch + # Matches getitune placeholder class names like "label_0", "label_1". _PLACEHOLDER_NAME = re.compile(r"^label_\d+$") @@ -34,7 +39,7 @@ class GetiAdapter(DetectionAdapter): >>> detections = adapter.detect(frame, frame_id=0) """ - def __init__(self, model: Any, device: str = "cpu") -> None: # noqa: ANN401 + def __init__(self, model: LightningDetectionModel, device: Literal["cpu", "cuda", "mps"] = "cpu") -> None: """Wrap a getitune detection model. Args: @@ -79,7 +84,7 @@ def detect(self, frame_bgr: np.ndarray, frame_id: int) -> Detections: 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) -> Any: # noqa: ANN401 + def preprocess(self, frame_bgr: np.ndarray) -> SampleBatch: """Preprocess a BGR frame into a getitune ``SampleBatch``. Applies the resize and normalization described by the model's @@ -119,7 +124,7 @@ def preprocess(self, frame_bgr: np.ndarray) -> Any: # noqa: ANN401 return SampleBatch(images=tensor.unsqueeze(0).to(self.device), imgs_info=[img_info]) @staticmethod - def to_detections(batch: Any, frame_id: int, image_index: int = 0) -> Detections: # noqa: ANN401 + def to_detections(batch: PredictionBatch, frame_id: int, image_index: int = 0) -> Detections: """Convert one image's predictions from a getitune ``PredictionBatch``. Duck-typed: works without getitune or torch installed, on any @@ -149,8 +154,8 @@ class ids. ) -def _to_numpy(value: Any) -> np.ndarray: # noqa: ANN401 - """Convert a torch tensor (possibly on an accelerator) or array-like to numpy.""" +def _to_numpy(value: torch.Tensor | np.ndarray) -> np.ndarray: + """Convert a torch tensor (possibly on an accelerator) or array to numpy.""" if hasattr(value, "cpu"): value = value.cpu() if hasattr(value, "numpy"): From 244f823bfa928aa6bd71d49ccd9d961882c9a66d Mon Sep 17 00:00:00 2001 From: Omkar Kabde Date: Thu, 2 Jul 2026 17:04:32 +0530 Subject: [PATCH 8/9] validate uint8 dtype in VideoWriter.write and preprocess caching --- .../getitrack/src/getitrack/adapters/geti.py | 55 ++++++++++++++----- libraries/getitrack/src/getitrack/io/video.py | 6 +- .../getitrack/tests/unit/test_adapters.py | 26 ++++++--- .../getitrack/tests/unit/test_video_io.py | 6 ++ 4 files changed, 68 insertions(+), 25 deletions(-) diff --git a/libraries/getitrack/src/getitrack/adapters/geti.py b/libraries/getitrack/src/getitrack/adapters/geti.py index f5fcbfe836e..321bd0831ea 100644 --- a/libraries/getitrack/src/getitrack/adapters/geti.py +++ b/libraries/getitrack/src/getitrack/adapters/geti.py @@ -21,11 +21,32 @@ 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. @@ -49,6 +70,9 @@ def __init__(self, model: LightningDetectionModel, device: Literal["cpu", "cuda" """ 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: @@ -87,14 +111,14 @@ def detect(self, frame_bgr: np.ndarray, frame_id: int) -> Detections: def preprocess(self, frame_bgr: np.ndarray) -> SampleBatch: """Preprocess a BGR frame into a getitune ``SampleBatch``. - Applies the resize and normalization described by the model's - ``data_input_params`` and sets ``scale_factor`` so predicted - boxes map back to the original frame coordinates. Requires - getitune and torch. + 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 frame in BGR order, e.g. from - ``cv2.VideoCapture``. + 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``. @@ -102,18 +126,19 @@ def preprocess(self, frame_bgr: np.ndarray) -> SampleBatch: import torch from getitune.data.entity import ImageInfo, SampleBatch - inp_h, inp_w = self.model.data_input_params.input_size - mean = self.model.data_input_params.mean - std = self.model.data_input_params.std + 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) - tensor = torch.from_numpy(cv2.resize(rgb, (inp_w, inp_h))).permute(2, 0, 1).float() - # Mean values below 1.0 indicate the model expects 0-1 normalized pixels - # (e.g. RF-DETR); otherwise it expects the raw 0-255 range (e.g. YOLOX). - if all(m < 1.0 for m in mean): - tensor = tensor / 255.0 - tensor = (tensor - torch.tensor(mean).view(3, 1, 1)) / torch.tensor(std).view(3, 1, 1) + 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, diff --git a/libraries/getitrack/src/getitrack/io/video.py b/libraries/getitrack/src/getitrack/io/video.py index 7bf343e61c5..c5213200eee 100644 --- a/libraries/getitrack/src/getitrack/io/video.py +++ b/libraries/getitrack/src/getitrack/io/video.py @@ -139,8 +139,12 @@ def write(self, frame: np.ndarray) -> None: ``frame_size`` given at construction. Raises: - ValueError: If the frame shape does not match ``frame_size``. + 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}" diff --git a/libraries/getitrack/tests/unit/test_adapters.py b/libraries/getitrack/tests/unit/test_adapters.py index f7294a23ca4..79860526147 100644 --- a/libraries/getitrack/tests/unit/test_adapters.py +++ b/libraries/getitrack/tests/unit/test_adapters.py @@ -104,9 +104,11 @@ def test_empty_names_return_none(self): 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)) -> GetiAdapter: - model = SimpleNamespace(data_input_params=SimpleNamespace(input_size=input_size, mean=mean, std=std)) - return GetiAdapter(model) + 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) @@ -124,16 +126,20 @@ def test_unit_range_normalization_for_small_means(self): assert float(batch.images.min()) >= -1.0 assert float(batch.images.max()) <= 1.0 - def test_raw_pixel_range_for_large_means(self): - batch = self._adapter(mean=(123.0, 117.0, 104.0), std=(58.0, 57.0, 57.0)).preprocess(self._frame()) - # Raw 0-255 pixels minus a ~100 mean over a ~57 std exceeds 1.5. - assert float(batch.images.max()) > 1.5 + 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) - # mean 0.0 < 1.0 triggers the 0-1 branch, so blue becomes 1.0 in channel 2. + # 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 @@ -142,7 +148,9 @@ def test_bgr_to_rgb_conversion(self): 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)) + 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: assert tuple(batch.images.shape) == (1, 3, 32, 32) diff --git a/libraries/getitrack/tests/unit/test_video_io.py b/libraries/getitrack/tests/unit/test_video_io.py index 042336fa84b..cdab070b63b 100644 --- a/libraries/getitrack/tests/unit/test_video_io.py +++ b/libraries/getitrack/tests/unit/test_video_io.py @@ -64,6 +64,12 @@ def test_wrong_frame_shape_raises(self, tmp_path): 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: From 489da2ea69b9aee2c07c1aa6982c6253eb66ea2e Mon Sep 17 00:00:00 2001 From: Omkar Kabde Date: Sun, 5 Jul 2026 01:17:18 +0530 Subject: [PATCH 9/9] remove numpy helper --- .../getitrack/src/getitrack/adapters/geti.py | 23 ++----- .../getitrack/tests/unit/test_adapters.py | 65 ++++++++----------- 2 files changed, 32 insertions(+), 56 deletions(-) diff --git a/libraries/getitrack/src/getitrack/adapters/geti.py b/libraries/getitrack/src/getitrack/adapters/geti.py index 321bd0831ea..e20760fcf83 100644 --- a/libraries/getitrack/src/getitrack/adapters/geti.py +++ b/libraries/getitrack/src/getitrack/adapters/geti.py @@ -3,8 +3,7 @@ """Adapter for getitune detection models. getitune and torch are imported lazily inside the methods that need them, so -getitrack imports without either installed. `GetiAdapter.to_detections` is -duck-typed and works on any prediction-batch-shaped object. +getitrack imports without either installed. """ from __future__ import annotations @@ -152,12 +151,9 @@ def preprocess(self, frame_bgr: np.ndarray) -> SampleBatch: def to_detections(batch: PredictionBatch, frame_id: int, image_index: int = 0) -> Detections: """Convert one image's predictions from a getitune ``PredictionBatch``. - Duck-typed: works without getitune or torch installed, on any - object with per-image ``bboxes``, ``scores``, and ``labels`` lists. - Args: batch: A getitune ``PredictionBatch`` with per-image ``bboxes``, - ``scores``, and ``labels`` lists (torch tensors or arrays). + ``scores``, and ``labels`` torch tensors. frame_id: Frame index to stamp on the returned `Detections`. image_index: Which image of the batch to convert. @@ -172,17 +168,8 @@ class ids. msg = "PredictionBatch must carry bboxes, scores, and labels" raise ValueError(msg) return Detections( - bboxes=_to_numpy(batch.bboxes[image_index]).reshape(-1, 4).astype(np.float32), - scores=_to_numpy(batch.scores[image_index]).reshape(-1).astype(np.float32), - class_ids=_to_numpy(batch.labels[image_index]).reshape(-1).astype(np.int64), + 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, ) - - -def _to_numpy(value: torch.Tensor | np.ndarray) -> np.ndarray: - """Convert a torch tensor (possibly on an accelerator) or array to numpy.""" - if hasattr(value, "cpu"): - value = value.cpu() - if hasattr(value, "numpy"): - value = value.numpy() - return np.asarray(value) diff --git a/libraries/getitrack/tests/unit/test_adapters.py b/libraries/getitrack/tests/unit/test_adapters.py index 79860526147..6d2af79432b 100644 --- a/libraries/getitrack/tests/unit/test_adapters.py +++ b/libraries/getitrack/tests/unit/test_adapters.py @@ -16,33 +16,21 @@ ) -class _FakeTensor: - """Mimics a torch tensor on an accelerator: needs .cpu().numpy().""" - - def __init__(self, data): - self._data = np.asarray(data) - self.moved = False - - def cpu(self) -> "_FakeTensor": - self.moved = True - return self - - def numpy(self) -> np.ndarray: - assert self.moved, "numpy() requires cpu() first, like a real cuda/mps tensor" - return self._data - +@_NEEDS_GETITUNE +class TestToDetections: + """Requires torch; ``to_detections`` converts ``PredictionBatch`` tensors.""" -def _batch(n=2, wrap=_FakeTensor) -> SimpleNamespace: - return SimpleNamespace( - bboxes=[wrap([[10.0, 10.0, 50.0, 50.0], [60.0, 60.0, 90.0, 90.0]][:n])], - scores=[wrap([0.9, 0.4][:n])], - labels=[wrap([3, 8][:n])], - ) + 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])], + ) -class TestToDetections: - def test_converts_fake_tensors(self): - d = GetiAdapter.to_detections(_batch(), frame_id=7) + 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 @@ -50,16 +38,13 @@ def test_converts_fake_tensors(self): assert d.bboxes.shape == (2, 4) assert d.class_ids.tolist() == [3, 8] - def test_converts_plain_numpy(self): - d = GetiAdapter.to_detections(_batch(wrap=np.asarray), frame_id=0) - assert len(d) == 2 - assert d.scores.tolist() == pytest.approx([0.9, 0.4]) - def test_empty_image(self): + import torch + batch = SimpleNamespace( - bboxes=[np.empty((0, 4))], - scores=[np.empty((0,))], - labels=[np.empty((0,))], + 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 @@ -70,10 +55,12 @@ def test_missing_fields_raise(self): GetiAdapter.to_detections(batch, frame_id=0) def test_image_index_selects_batch_element(self): + import torch + batch = SimpleNamespace( - bboxes=[np.zeros((1, 4)), np.array([[5.0, 5.0, 9.0, 9.0]])], - scores=[np.array([0.5]), np.array([0.7])], - labels=[np.array([0]), np.array([2])], + 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] @@ -153,11 +140,13 @@ class _FakeModel: ) def predict_step(self, batch, batch_idx) -> SimpleNamespace: + import torch + assert tuple(batch.images.shape) == (1, 3, 32, 32) return SimpleNamespace( - bboxes=[np.array([[1.0, 2.0, 3.0, 4.0]])], - scores=[np.array([0.9])], - labels=[np.array([1])], + 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)