fea: norfair demo nb (#3)

* fea: norfair demo nb

* fix: data folder + script

* fix: improve video generation

* fix: script

* fix: use .txt instead of pip tools

* fix: pre-commit

---------

Co-authored-by: TomDarmon <tom.darmon@artefact.com>

.gitignore

@@ -135,4 +135,9 @@ secrets/*
 !secrets/.gitkeep
 
 # Mac OS
-.DS_Store
+.DS_Store
+
+
+# Data ignore everythin data/detections and data/frames
+data/detections/*
+data/frames/*

.pre-commit-config.yaml

@@ -31,12 +31,6 @@ repos:
         types: [file]
         files: (.ipynb)$
         language: system
-      - id: python-bandit-vulnerability-check
-        name: Security check (bandit)
-        entry: bandit
-        types: [python]
-        args: ["--recursive", "lib/"]
-        language: system
       - id: pytest-check
         name: Tests (pytest)
         stages: [push]

Makefile

@@ -18,17 +18,10 @@ help:
 install:
 	@bash bin/$(INSTALL_SCRIPT)
 
-# help: install_pip_tools		- Install piptools and setuptools
-.PHONY: install_pip_tools
-install_pip_tools:
-	@echo "Installing pip-tools"
-	@pip install pip-tools==$(PIP_TOOLS_VERSION) setuptools==$(SETUPTOOLS_VERSION)
-
 # help: install_project_requirements		- Install prohect requirements
 .PHONY: install_project_requirements
 install_project_requirements: install_pip_tools
 	@pip install numpy==${NUMPY_VERSION}
-	@pip-compile requirements.in
 	@pip install -r requirements.txt
 
 # help: install_precommit			- Install pre-commit hooks

bin/download_sample_sequences.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+
+N_SEQUENCES=5
+
+sequences_detections=$(gsutil ls gs://data-track-reid/detections | head -$N_SEQUENCES)
+sequences_frames=$(gsutil ls gs://data-track-reid/frames | head -$N_SEQUENCES)
+
+# remove first sequence which is the bucket name
+sequences_detections=$(echo "$sequences_detections" | tail -n +2)
+sequences_frames=$(echo "$sequences_frames" | tail -n +2)
+
+
+# download the sequences to data/detections and data/frames
+for sequence in $sequences_detections; do
+    gsutil -m cp -r $sequence data/detections
+done
+
+for sequence in $sequences_frames; do
+    gsutil -m cp -r $sequence data/frames
+done

lib/bbox/utils.py

@@ -0,0 +1,25 @@
+import numpy as np
+
+
+def xy_center_to_xyxy(bbox: np.array) -> np.array:
+    """Convert bounding box from xy_center to xyxy format"""
+    return np.array(
+        [
+            bbox[0] - bbox[2] / 2,
+            bbox[1] - bbox[3] / 2,
+            bbox[0] + bbox[2] / 2,
+            bbox[1] + bbox[3] / 2,
+        ]
+    )
+
+
+def rescale_bbox(bbox: np.array, original_size: tuple) -> np.array:
+    """Rescale bounding box from current_size to original_size"""
+    return np.array(
+        [
+            bbox[0] * original_size[0],
+            bbox[1] * original_size[1],
+            bbox[2] * original_size[0],
+            bbox[3] * original_size[1],
+        ]
+    )

lib/norfair_helper/utils.py

@@ -0,0 +1,25 @@
+from typing import List
+
+import numpy as np
+from norfair import Detection
+
+from lib.bbox.utils import rescale_bbox, xy_center_to_xyxy
+
+
+def yolo_to_norfair_detection(
+    yolo_detections: np.array, original_img_size: tuple
+) -> List[Detection]:
+    """convert detections_as_xywh to norfair detections"""
+    norfair_detections: List[Detection] = []
+    for detection_output in yolo_detections:
+        bbox = np.array(
+            [
+                [detection_output[1].item(), detection_output[2].item()],
+                [detection_output[3].item(), detection_output[4].item()],
+            ]
+        )
+        bbox = xy_center_to_xyxy(bbox.flatten()).reshape(2, 2)
+        bbox = rescale_bbox(bbox.flatten(), original_img_size).reshape(2, 2)
+        scores = np.array([detection_output[5].item(), detection_output[5].item()])
+        norfair_detections.append(Detection(points=bbox, scores=scores, label=detection_output[0]))
+    return norfair_detections

lib/norfair_helper/video.py

@@ -0,0 +1,37 @@
+import cv2
+import numpy as np
+from norfair import Tracker, draw_boxes
+
+from lib.norfair_helper.utils import yolo_to_norfair_detection
+from lib.sequence import Sequence
+
+
+def generate_tracking_video(
+    sequence: Sequence, tracker: Tracker, frame_size: tuple, output_path: str
+) -> str:
+    """
+    Generate a video with the tracking results.
+
+    Args:
+        sequence: The sequence of frames and detections.
+        tracker: The tracker to use.
+        frame_size: The size of the frames.
+        output_path: The path to save the video to.
+
+    Returns:
+        The path to the video.
+
+    """
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # Changed codec to 'mp4v' for compatibility with Mac
+    out = cv2.VideoWriter(output_path, fourcc, 20.0, frame_size)  # Changed file extension to .mp4
+
+    for frame, detection in sequence:
+        detections_list = yolo_to_norfair_detection(detection, frame_size)
+        tracked_objects = tracker.update(detections=detections_list)
+        frame_detected = draw_boxes(
+            np.array(frame), tracked_objects, draw_ids=True, color="by_label"
+        )
+        frame_detected = cv2.cvtColor(frame_detected, cv2.COLOR_BGR2RGB)
+        out.write(frame_detected)
+    out.release()
+    return output_path

lib/sequence.py

@@ -0,0 +1,39 @@
+from dataclasses import dataclass
+from typing import List
+
+import numpy as np
+from PIL import Image
+
+
+@dataclass
+class Sequence:
+    """This class represents a sequence of frames and detections and make the data iterable."""
+
+    frame_paths: List[str]
+
+    def __post_init__(self):
+        self.detection_paths: List[str] = [
+            f.replace("frames", "detections").replace(".jpg", ".txt") for f in self.frame_paths
+        ]
+
+    def __repr__(self):
+        return (
+            f"Sequence(n_frames={len(self.frame_paths)}, n_detections={len(self.detection_paths)})"
+        )
+
+    def __iter__(self):
+        self.index = 0
+        return self
+
+    def __next__(self):
+        if self.index >= len(self.frame_paths):
+            raise StopIteration
+
+        frame = Image.open(self.frame_paths[self.index])
+        try:
+            detection = np.loadtxt(self.detection_paths[self.index], dtype="float")
+        except OSError:  # file doesn't exist not detection return empty file
+            detection = np.array([])
+
+        self.index += 1
+        return frame, detection