detector.py

import os
import time

import cv2
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np
import onnxruntime as ort
import pandas as pd
from typing import Tuple
from huggingface_hub import hf_hub_download

from constants import REPO_ID, FILENAME, MODEL_DIR, MODEL_PATH
from metrics_storage import MetricsStorage


def download_model():
    """Download the model using Hugging Face Hub"""
    # Ensure model directory exists
    os.makedirs(MODEL_DIR, exist_ok=True)

    try:
        print(f"Downloading model from {REPO_ID}...")
        # Download the model file from Hugging Face Hub
        model_path = hf_hub_download(
            repo_id=REPO_ID,
            filename=FILENAME,
            local_dir=MODEL_DIR,
            force_download=True,
            cache_dir=None,
        )

        # Move the file to the correct location if it's not there already
        if os.path.exists(model_path) and model_path != MODEL_PATH:
            os.rename(model_path, MODEL_PATH)

            # Remove empty directories if they exist
            empty_dir = os.path.join(MODEL_DIR, "tune")
            if os.path.exists(empty_dir):
                import shutil

                shutil.rmtree(empty_dir)

        print("Model downloaded successfully!")
        return MODEL_PATH

    except Exception as e:
        print(f"Error downloading model: {e}")
        raise e


class SignatureDetector:
    def __init__(self, model_path: str = MODEL_PATH):
        self.model_path = model_path
        self.classes = ["signature"]
        self.input_width = 640
        self.input_height = 640

        # Initialize ONNX Runtime session
        options = ort.SessionOptions()
        options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
        self.session = ort.InferenceSession(self.model_path, options)
        self.session.set_providers(
            ["OpenVINOExecutionProvider"], [{"device_type": "CPU"}]
        )

        self.metrics_storage = MetricsStorage()

    def update_metrics(self, inference_time: float) -> None:
        """
        Updates metrics in persistent storage.

        Args:
            inference_time (float): The time taken for inference in milliseconds.
        """
        self.metrics_storage.add_metric(inference_time)

    def get_metrics(self) -> dict:
        """
        Retrieves current metrics from storage.

        Returns:
            dict: A dictionary containing times, total inferences, average time, and start index.
        """
        times = self.metrics_storage.get_recent_metrics()
        total = self.metrics_storage.get_total_inferences()
        avg = self.metrics_storage.get_average_time()

        start_index = max(0, total - len(times))

        return {
            "times": times,
            "total_inferences": total,
            "avg_time": avg,
            "start_index": start_index,
        }

    def load_initial_metrics(
        self,
    ) -> Tuple[None, str, plt.Figure, plt.Figure, str, str]:
        """
        Loads initial metrics for display.

        Returns:
            tuple: A tuple containing None, total inferences, histogram figure, line figure, average time, and last time.
        """
        metrics = self.get_metrics()

        if not metrics["times"]:
            return None, None, None, None, None, None

        hist_data = pd.DataFrame({"Time (ms)": metrics["times"]})
        indices = range(
            metrics["start_index"], metrics["start_index"] + len(metrics["times"])
        )

        line_data = pd.DataFrame(
            {
                "Inference": indices,
                "Time (ms)": metrics["times"],
                "Mean": [metrics["avg_time"]] * len(metrics["times"]),
            }
        )

        hist_fig, line_fig = self.create_plots(hist_data, line_data)

        return (
            None,
            f"{metrics['total_inferences']}",
            hist_fig,
            line_fig,
            f"{metrics['avg_time']:.2f}",
            f"{metrics['times'][-1]:.2f}",
        )

    def create_plots(
        self, hist_data: pd.DataFrame, line_data: pd.DataFrame
    ) -> Tuple[plt.Figure, plt.Figure]:
        """
        Helper method to create plots.

        Args:
            hist_data (pd.DataFrame): Data for histogram plot.
            line_data (pd.DataFrame): Data for line plot.

        Returns:
            tuple: A tuple containing histogram figure and line figure.
        """
        plt.style.use("dark_background")

        # Histogram plot
        hist_fig, hist_ax = plt.subplots(figsize=(8, 4), facecolor="#f0f0f5")
        hist_ax.set_facecolor("#f0f0f5")
        hist_data.hist(
            bins=20, ax=hist_ax, color="#4F46E5", alpha=0.7, edgecolor="white"
        )
        hist_ax.set_title(
            "Distribution of Inference Times",
            pad=15,
            fontsize=12,
            color="#1f2937",
        )
        hist_ax.set_xlabel("Time (ms)", color="#374151")
        hist_ax.set_ylabel("Frequency", color="#374151")
        hist_ax.tick_params(colors="#4b5563")
        hist_ax.grid(True, linestyle="--", alpha=0.3)

        # Line plot
        line_fig, line_ax = plt.subplots(figsize=(8, 4), facecolor="#f0f0f5")
        line_ax.set_facecolor("#f0f0f5")
        line_data.plot(
            x="Inference",
            y="Time (ms)",
            ax=line_ax,
            color="#4F46E5",
            alpha=0.7,
            label="Time",
        )
        line_data.plot(
            x="Inference",
            y="Mean",
            ax=line_ax,
            color="#DC2626",
            linestyle="--",
            label="Mean",
        )
        line_ax.set_title(
            "Inference Time per Execution", pad=15, fontsize=12, color="#1f2937"
        )
        line_ax.set_xlabel("Inference Number", color="#374151")
        line_ax.set_ylabel("Time (ms)", color="#374151")
        line_ax.tick_params(colors="#4b5563")
        line_ax.grid(True, linestyle="--", alpha=0.3)
        line_ax.legend(
            frameon=True, facecolor="#f0f0f5", edgecolor="white", labelcolor="black"
        )

        hist_fig.tight_layout()
        line_fig.tight_layout()

        plt.close(hist_fig)
        plt.close(line_fig)

        return hist_fig, line_fig

    def preprocess(self, img: Image.Image) -> Tuple[np.ndarray, np.ndarray]:
        """
        Preprocesses the image for inference.

        Args:
            img: The image to process.

        Returns:
            tuple: A tuple containing the processed image data and the original image.
        """
        # Convert PIL Image to cv2 format
        img_cv2 = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

        self.img_height, self.img_width = img_cv2.shape[:2]

        # Convert back to RGB for processing
        img_rgb = cv2.cvtColor(img_cv2, cv2.COLOR_BGR2RGB)

        # Resize
        img_resized = cv2.resize(img_rgb, (self.input_width, self.input_height))

        # Normalize and transpose
        image_data = np.array(img_resized) / 255.0
        image_data = np.transpose(image_data, (2, 0, 1))
        image_data = np.expand_dims(image_data, axis=0).astype(np.float32)

        return image_data, img_cv2

    def draw_detections(
        self, img: np.ndarray, box: list, score: float, class_id: int
    ) -> None:
        """
        Draws the detections on the image.

        Args:
            img: The image to draw on.
            box (list): The bounding box coordinates.
            score (float): The confidence score.
            class_id (int): The class ID.
        """
        x1, y1, w, h = box
        self.color_palette = np.random.uniform(0, 255, size=(len(self.classes), 3))
        color = self.color_palette[class_id]

        cv2.rectangle(img, (int(x1), int(y1)), (int(x1 + w), int(y1 + h)), color, 2)

        label = f"{self.classes[class_id]}: {score:.2f}"
        (label_width, label_height), _ = cv2.getTextSize(
            label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1
        )

        label_x = x1
        label_y = y1 - 10 if y1 - 10 > label_height else y1 + 10

        cv2.rectangle(
            img,
            (int(label_x), int(label_y - label_height)),
            (int(label_x + label_width), int(label_y + label_height)),
            color,
            cv2.FILLED,
        )

        cv2.putText(
            img,
            label,
            (int(label_x), int(label_y)),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (0, 0, 0),
            1,
            cv2.LINE_AA,
        )

    def postprocess(
        self,
        input_image: np.ndarray,
        output: np.ndarray,
        conf_thres: float,
        iou_thres: float,
    ) -> np.ndarray:
        """
        Postprocesses the output from inference.

        Args:
            input_image: The input image.
            output: The output from inference.
            conf_thres (float): Confidence threshold for detection.
            iou_thres (float): Intersection over Union threshold for detection.

        Returns:
            np.ndarray: The output image with detections drawn
        """
        outputs = np.transpose(np.squeeze(output[0]))
        rows = outputs.shape[0]

        boxes = []
        scores = []
        class_ids = []

        x_factor = self.img_width / self.input_width
        y_factor = self.img_height / self.input_height

        for i in range(rows):
            classes_scores = outputs[i][4:]
            max_score = np.amax(classes_scores)

            if max_score >= conf_thres:
                class_id = np.argmax(classes_scores)
                x, y, w, h = outputs[i][0], outputs[i][1], outputs[i][2], outputs[i][3]

                left = int((x - w / 2) * x_factor)
                top = int((y - h / 2) * y_factor)
                width = int(w * x_factor)
                height = int(h * y_factor)

                class_ids.append(class_id)
                scores.append(max_score)
                boxes.append([left, top, width, height])

        indices = cv2.dnn.NMSBoxes(boxes, scores, conf_thres, iou_thres)

        for i in indices:
            box = boxes[i]
            score = scores[i]
            class_id = class_ids[i]
            self.draw_detections(input_image, box, score, class_id)

        return cv2.cvtColor(input_image, cv2.COLOR_BGR2RGB)

    def detect(
        self, image: Image.Image, conf_thres: float = 0.25, iou_thres: float = 0.5
    ) -> Tuple[Image.Image, dict]:
        """
        Detects signatures in the given image.

        Args:
            image: The image to process.
            conf_thres (float): Confidence threshold for detection.
            iou_thres (float): Intersection over Union threshold for detection.

        Returns:
            tuple: A tuple containing the output image and metrics.
        """
        # Preprocess the image
        img_data, original_image = self.preprocess(image)

        # Run inference
        start_time = time.time()
        outputs = self.session.run(None, {self.session.get_inputs()[0].name: img_data})
        inference_time = (time.time() - start_time) * 1000  # Convert to milliseconds

        # Postprocess the results
        output_image = self.postprocess(original_image, outputs, conf_thres, iou_thres)

        self.update_metrics(inference_time)

        return output_image, self.get_metrics()

    def detect_example(
        self, image: Image.Image, conf_thres: float = 0.25, iou_thres: float = 0.5
    ) -> Image.Image:
        """
        Wrapper method for examples that returns only the image.

        Args:
            image: The image to process.
            conf_thres (float): Confidence threshold for detection.
            iou_thres (float): Intersection over Union threshold for detection.

        Returns:
            The output image.
        """
        output_image, _ = self.detect(image, conf_thres, iou_thres)
        return output_image