utils.py

import functools
import logging
import pickle
from typing import List, Tuple

import cv2 as cv
import numpy as np
import supervision as sv
from sklearn.cluster import KMeans

logger = logging.getLogger("UTILS")


@functools.cache
def _load_calibration_data() -> Tuple[np.ndarray, np.ndarray]:
    try:
        with open("camera-calibration/cameraMatrix.pkl", "rb") as f:
            camera_matrix = pickle.load(f)

        with open("camera-calibration/dist.pkl", "rb") as f:
            dist_coeffs = pickle.load(f)

        return camera_matrix, dist_coeffs
    except FileNotFoundError as e:
        raise FileNotFoundError(f"Camera calibration files not found: {e}")


def get_undistorted_remapped_frame(frame: np.ndarray) -> np.ndarray:
    camera_matrix, dist_coeffs = _load_calibration_data()

    h, w = frame.shape[:2]

    new_camera_matrix, roi = cv.getOptimalNewCameraMatrix(
        camera_matrix, dist_coeffs, (w, h), 1, (w, h)
    )

    mapx, mapy = cv.initUndistortRectifyMap(
        camera_matrix, dist_coeffs, None, new_camera_matrix, (w, h), 5
    )

    dst = cv.remap(frame, mapx, mapy, cv.INTER_LINEAR)

    x, y, crop_w, crop_h = roi
    dst = dst[y : y + crop_h, x : x + crop_w]

    return dst


def generate_labels(detections: sv.Detections) -> List[str]:
    labels = []
    for i in range(len(detections)):
        labels.append(
            f"#{detections.tracker_id[i] if detections.tracker_id is not None else 'NO ID'}"
        )
    return labels


def find_dominant_hsv(image):
    pixels = image.copy().reshape(-1, 3)

    if len(pixels) > 100000:
        step = len(pixels) // 100000
        pixels = pixels[::step]

    # 3 Clusters for H, S, V
    kmeans = KMeans(n_clusters=3, n_init="auto", random_state=42).fit(pixels)

    # Find cluster with most pixels
    counts = np.bincount(kmeans.labels_)
    dominant_idx = np.argmax(counts)
    dominant_hsv = kmeans.cluster_centers_[dominant_idx]

    return dominant_hsv