Strengthen the types

color_correction/core/card_detection/det_yv8_onnx.py

@@ -43,11 +43,11 @@ class YOLOv8CardDetector(BaseCardDetector):
         Flag indicating whether to use GPU for inference
     session : onnxruntime.InferenceSession
         ONNX Runtime session for model inference
-    input_names : list
+    input_names : list[str]
         Names of model input nodes
-    output_names : list
+    output_names : list[str]
         Names of model output nodes
-    input_shape : tuple
+    input_shape : tuple[int, ...]
         Shape of the input tensor
     input_height : int
         Height of the input image required by the model

color_correction/core/correction/_factory.py

@@ -1,21 +1,53 @@
+from typing import Any
+
 from color_correction.core.correction.affine_reg import AffineRegression
 from color_correction.core.correction.least_squares import (
     LeastSquaresRegression,
 )
 from color_correction.core.correction.linear_reg import LinearRegression
 from color_correction.core.correction.polynomial import Polynomial
+from color_correction.schemas.custom_types import LiteralModelCorrection
+
+# Type alias for correction models
+CorrectionModel = LeastSquaresRegression | Polynomial | LinearRegression | AffineRegression
 
 
 class CorrectionModelFactory:
+    """Factory class for creating color correction models."""
+
     @staticmethod
     def create(
-        model_name: str,
-        **kwargs: dict,
-    ) -> LeastSquaresRegression | Polynomial | LinearRegression | AffineRegression:
-        model_registry = {
+        model_name: LiteralModelCorrection,
+        **kwargs: Any,  # noqa: ANN401
+    ) -> CorrectionModel:
+        """
+        Create a correction model instance based on the model name.
+
+        Parameters
+        ----------
+        model_name : LiteralModelCorrection
+            Name of the correction model to create.
+        **kwargs : Any
+            Additional parameters passed to the model constructor.
+
+        Returns
+        -------
+        CorrectionModel
+            An instance of the requested correction model.
+
+        Raises
+        ------
+        KeyError
+            If model_name is not a valid correction model.
+        """
+        model_registry: dict[str, CorrectionModel] = {
             "least_squares": LeastSquaresRegression(),
             "polynomial": Polynomial(**kwargs),
             "linear_reg": LinearRegression(),
             "affine_reg": AffineRegression(),
         }
-        return model_registry.get(model_name)
+        model = model_registry.get(model_name)
+        if model is None:
+            valid_models = list(model_registry.keys())
+            raise KeyError(f"Unknown model '{model_name}'. Valid options: {valid_models}")
+        return model

color_correction/core/correction/polynomial.py

@@ -1,8 +1,9 @@
 import time
+from typing import Any
 
 import numpy as np
 from sklearn.linear_model import LinearRegression
-from sklearn.pipeline import make_pipeline
+from sklearn.pipeline import Pipeline, make_pipeline
 from sklearn.preprocessing import PolynomialFeatures
 
 from color_correction.core.correction.base import BaseComputeCorrection
@@ -18,20 +19,20 @@ class Polynomial(BaseComputeCorrection):
 
     Parameters
     ----------
-    **kwargs : dict, optional
+    **kwargs : Any
         Keyword arguments. Recognized keyword:
 
         - `degree` : int, optional, default 2
              Degree of the polynomial.
     """
 
-    def __init__(self, **kwargs: dict) -> None:
+    def __init__(self, **kwargs: Any) -> None:  # noqa: ANN401
         """
         Initialize the Polynomial correction model.
 
         Parameters
         ----------
-        **kwargs : dict
+        **kwargs : Any
             Keyword arguments for initialization.
 
         Other Parameters
@@ -41,15 +42,15 @@ def __init__(self, **kwargs: dict) -> None:
             The more complex the polynomial, the more flexible the model.
             But it may also lead to overfitting.
         """
-        self.model = None
-        self.degree = kwargs.get("degree", 2)
+        self.model: Pipeline | None = None
+        self.degree: int = kwargs.get("degree", 2)
 
     def fit(
         self,
         x_patches: np.ndarray,  # input patches
         y_patches: np.ndarray,  # reference patches
-        **kwargs: dict,
-    ) -> np.ndarray:
+        **kwargs: Any,  # noqa: ANN401
+    ) -> Pipeline:
         """
         Fit the polynomial regression model.
 
@@ -59,16 +60,16 @@ def fit(
             Input image patches.
         y_patches : np.ndarray
             Reference image patches.
-        **kwargs : dict
+        **kwargs : Any
             Additional keyword arguments. Recognized keyword:
 
             - `degree` : int, optional
                  Degree of the polynomial.
 
         Returns
         -------
-        np.ndarray
-            Fitted model pipeline.
+        Pipeline
+            Fitted sklearn pipeline with polynomial features and linear regression.
 
         """
         start_time = time.perf_counter()

color_correction/services/color_correction.py

@@ -1,4 +1,5 @@
 import os
+from typing import Any, TypedDict
 
 import cv2
 import numpy as np
@@ -31,6 +32,23 @@
 )
 
 
+class ColorDiffMetrics(TypedDict):
+    """Color difference metrics from CIE 2000 calculation."""
+
+    min: float
+    max: float
+    mean: float
+    std: float
+
+
+class ColorDiffResult(TypedDict):
+    """Result structure from calc_color_diff_patches."""
+
+    initial: ColorDiffMetrics
+    corrected: ColorDiffMetrics
+    delta: ColorDiffMetrics
+
+
 class ColorCorrection:
     """Color correction handler using color `card_detection` and `correction_models`.
     This class handles the complete workflow of color correction, including:
@@ -54,7 +72,7 @@ class ColorCorrection:
         If None, uses standard D50 values.
     use_gpu : bool, default=False
         True to use GPU for card detection. False will use CPU.
-    **kwargs : dict
+    **kwargs : Any
         Additional parameters for the correction model.
 
     Other parameters
@@ -80,7 +98,7 @@ def __init__(
         correction_model: LiteralModelCorrection = "least_squares",
         reference_image: ImageBGR | None = None,
         use_gpu: bool = False,
-        **kwargs: dict,
+        **kwargs: Any,  # noqa: ANN401
     ) -> None:
         # Validate reference_image if provided
         if reference_image is not None:
@@ -494,7 +512,7 @@ def predict(
 
         return corrected_image
 
-    def calc_color_diff_patches(self) -> dict:
+    def calc_color_diff_patches(self) -> ColorDiffResult:
         """
         Calculate color difference metrics for image patches using the dE CIE 2000 metric.
 
@@ -513,19 +531,15 @@ def calc_color_diff_patches(self) -> dict:
 
         Returns
         -------
-        dict
-            A dictionary with the following keys:
-
-            - `initial`: dict containing the color difference metrics for the initial patches versus the reference.
-            - `corrected`: dict containing the color difference metrics for the corrected patches versus the reference.
-            - `delta`: dict with metrics representing the difference between the initial and corrected color differences.
-                      Each metric is computed as:
-                      ```python
-                        metric_delta = metric_initial - metric_corrected,
-                      ```
-                      where metrics include `min`, `max`, `mean`, and `std`.
+        ColorDiffResult
+            A TypedDict with the following keys:
 
-        """  # noqa: E501
+            - `initial`: ColorDiffMetrics for the initial patches versus the reference.
+            - `corrected`: ColorDiffMetrics for the corrected patches versus the reference.
+            - `delta`: ColorDiffMetrics representing the difference (initial - corrected).
+                      Each metric includes `min`, `max`, `mean`, and `std`.
+
+        """
         # check input_grid_image, reference_grid_image, corrected_grid_image
         print(
             f"input_grid_image: {self.input_grid_image.shape}, "

color_correction/services/correction_analyzer.py

@@ -1,9 +1,10 @@
 import os
+from typing import Any, TypedDict
 
-import numpy as np
 import pandas as pd
 
 from color_correction.schemas.custom_types import (
+    ImageBGR,
     LiteralModelCorrection,
     LiteralModelDetection,
 )
@@ -15,6 +16,18 @@
 from color_correction.utils.report_generator import ReportGenerator
 
 
+class DetectionParams(TypedDict, total=False):
+    """Parameters for detection methods."""
+
+    detection_conf_th: float
+
+
+class CorrectionParams(TypedDict, total=False):
+    """Parameters for correction methods."""
+
+    degree: int
+
+
 class ColorCorrectionAnalyzer:
     """
     Analyzer for benchmarking color correction methods.
@@ -25,10 +38,10 @@ class ColorCorrectionAnalyzer:
 
     Parameters
     ----------
-    list_correction_methods : list of tuple[LiteralModelCorrection, dict]
+    list_correction_methods : list[tuple[LiteralModelCorrection, CorrectionParams]]
         A list of tuples, where each tuple contains a correction method identifier
         and its parameters.
-    list_detection_methods : list of tuple[LiteralModelDetection, dict]
+    list_detection_methods : list[tuple[LiteralModelDetection, DetectionParams]]
         A list of tuples, where each tuple contains a detection method identifier
         and its parameters.
     use_gpu : bool, optional
@@ -37,18 +50,18 @@ class ColorCorrectionAnalyzer:
 
     def __init__(
         self,
-        list_correction_methods: list[tuple[LiteralModelCorrection, dict]],
-        list_detection_methods: list[tuple[LiteralModelDetection, dict]],
+        list_correction_methods: list[tuple[LiteralModelCorrection, CorrectionParams]],
+        list_detection_methods: list[tuple[LiteralModelDetection, DetectionParams]],
         use_gpu: bool = False,
     ) -> None:
         """
         Initialize the ColorCorrectionAnalyzer.
 
         Parameters
         ----------
-        list_correction_methods : list of tuple[LiteralModelCorrection, dict]]
+        list_correction_methods : list[tuple[LiteralModelCorrection, CorrectionParams]]
             List of correction methods and their parameters.
-        list_detection_methods : list of tuple[LiteralModelDetection, dict]]
+        list_detection_methods : list[tuple[LiteralModelDetection, DetectionParams]]
             List of detection methods and their parameters.
         use_gpu : bool, optional
             Whether to use GPU acceleration, by default True.
@@ -61,36 +74,36 @@ def __init__(
     def _run_single_exp(
         self,
         idx: int,
-        input_image: np.ndarray,
+        input_image: ImageBGR,
         det_method: LiteralModelDetection,
-        det_params: dict,
+        det_params: DetectionParams,
         cc_method: LiteralModelCorrection,
-        cc_params: dict,
-        reference_image: np.ndarray | None = None,
-    ) -> dict:
+        cc_params: CorrectionParams,
+        reference_image: ImageBGR | None = None,
+    ) -> dict[str, Any]:
         """
         Run a single experiment for a given detection and correction method.
 
         Parameters
         ----------
         idx : int
             Index of the experiment.
-        input_image : np.ndarray
-            The input image array.
+        input_image : ImageBGR
+            The input image array in BGR format.
         det_method : LiteralModelDetection
             The detection method identifier.
-        det_params : dict
+        det_params : DetectionParams
             Parameters for the detection method.
         cc_method : LiteralModelCorrection
             The correction method identifier.
-        cc_params : dict
+        cc_params : CorrectionParams
             Parameters for the correction method.
-        reference_image : np.ndarray, optional
+        reference_image : ImageBGR | None, optional
             The reference image, by default None.
 
         Returns
         -------
-        dict
+        dict[str, Any]
             A dictionary containing evaluation data and results of the experiment.
         """
         cc = ColorCorrection(
@@ -152,20 +165,20 @@ def _run_single_exp(
 
     def run(
         self,
-        input_image: np.ndarray,
+        input_image: ImageBGR,
         output_dir: str = "benchmark_debug",
-        reference_image: np.ndarray | None = None,
+        reference_image: ImageBGR | None = None,
     ) -> pd.DataFrame:
         """
         Run the full benchmark for color correction and generate reports.
 
         Parameters
         ----------
-        input_image : np.ndarray
-            The image to be processed.
+        input_image : ImageBGR
+            The image to be processed in BGR format.
         output_dir : str, optional
             The directory to save reports, by default `benchmark_debug`.
-        reference_image : np.ndarray, optional
+        reference_image : ImageBGR | None, optional
             Optional reference image used for evaluation, by default None.
 
         Returns