feat: Adicionar anotações de tipo e documentação para métodos na classe SignatureDetector

detector.py

@@ -3,9 +3,11 @@ 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
@@ -48,7 +50,7 @@ def download_model():
 
 
 class SignatureDetector:
-    def __init__(self, model_path):
+    def __init__(self, model_path: str = MODEL_PATH):
         self.model_path = model_path
         self.classes = ["signature"]
         self.input_width = 640
@@ -57,19 +59,29 @@ class SignatureDetector:
         # Initialize ONNX Runtime session
         options = ort.SessionOptions()
         options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
-        self.session = ort.InferenceSession(MODEL_PATH, options)
+        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):
-        """Update metrics in persistent storage"""
+    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):
-        """Get current metrics from storage"""
+    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()
@@ -80,17 +92,23 @@ class SignatureDetector:
             "times": times,
             "total_inferences": total,
             "avg_time": avg,
-            "start_index": start_index,  # Adicionar índice inicial
+            "start_index": start_index,
         }
 
-    def load_initial_metrics(self):
-        """Load initial metrics for display"""
+    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"]:  # Se não houver dados
+        if not metrics["times"]:
             return None, None, None, None, None, None
 
-        # Criar plots data
         hist_data = pd.DataFrame({"Tempo (ms)": metrics["times"]})
         indices = range(
             metrics["start_index"], metrics["start_index"] + len(metrics["times"])
@@ -104,7 +122,6 @@ class SignatureDetector:
             }
         )
 
-        # Criar plots
         hist_fig, line_fig = self.create_plots(hist_data, line_data)
 
         return (
@@ -116,11 +133,22 @@ class SignatureDetector:
             f"{metrics['times'][-1]:.2f}",
         )
 
-    def create_plots(self, hist_data, line_data):
-        """Helper method to create plots"""
+    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")
 
-        # Histograma
+        # Histogram plot
         hist_fig, hist_ax = plt.subplots(figsize=(8, 4), facecolor="#f0f0f5")
         hist_ax.set_facecolor("#f0f0f5")
         hist_data.hist(
@@ -137,7 +165,7 @@ class SignatureDetector:
         hist_ax.tick_params(colors="#4b5563")
         hist_ax.grid(True, linestyle="--", alpha=0.3)
 
-        # Gráfico de linha
+        # Line plot
         line_fig, line_ax = plt.subplots(figsize=(8, 4), facecolor="#f0f0f5")
         line_ax.set_facecolor("#f0f0f5")
         line_data.plot(
@@ -168,17 +196,24 @@ class SignatureDetector:
         hist_fig.tight_layout()
         line_fig.tight_layout()
 
-        # Fechar as figuras para liberar memória
         plt.close(hist_fig)
         plt.close(line_fig)
 
         return hist_fig, line_fig
 
-    def preprocess(self, img):
+    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)
 
-        # Get image dimensions
         self.img_height, self.img_width = img_cv2.shape[:2]
 
         # Convert back to RGB for processing
@@ -194,7 +229,18 @@ class SignatureDetector:
 
         return image_data, img_cv2
 
-    def draw_detections(self, img, box, score, class_id):
+    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]
@@ -228,7 +274,25 @@ class SignatureDetector:
             cv2.LINE_AA,
         )
 
-    def postprocess(self, input_image, output, conf_thres, iou_thres):
+    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]
 
@@ -266,7 +330,20 @@ class SignatureDetector:
 
         return cv2.cvtColor(input_image, cv2.COLOR_BGR2RGB)
 
-    def detect(self, image, conf_thres=0.25, iou_thres=0.5):
+    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)
 
@@ -282,7 +359,19 @@ class SignatureDetector:
 
         return output_image, self.get_metrics()
 
-    def detect_example(self, image, conf_thres=0.25, iou_thres=0.5):
-        """Wrapper method for examples that returns only the image"""
+    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