From 280adc6d240a9073ef188184638acbd597c6238a Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 4 Jan 2026 01:52:07 +0000 Subject: [PATCH] feat: strengthen type annotations across the codebase - Replace Any types with TypedDict in device_info.py for device specs - Add ColorDiffMetrics and ColorDiffResult TypedDicts for calc_color_diff_patches - Add DetectionParams and CorrectionParams TypedDicts for analyzer - Update loose dict and list types with proper type parameters - Add CorrectionModel type alias in _factory.py - Update docstrings to reflect new type annotations - Add noqa comments for intentional Any usage in **kwargs --- .../core/card_detection/det_yv8_onnx.py | 6 +- color_correction/core/correction/_factory.py | 42 +++++++++++-- .../core/correction/polynomial.py | 23 ++++---- color_correction/services/color_correction.py | 44 +++++++++----- .../services/correction_analyzer.py | 59 +++++++++++-------- color_correction/utils/device_info.py | 34 +++++++---- color_correction/utils/report_generator.py | 9 +-- 7 files changed, 144 insertions(+), 73 deletions(-) diff --git a/color_correction/core/card_detection/det_yv8_onnx.py b/color_correction/core/card_detection/det_yv8_onnx.py index 19de9bb..c494feb 100644 --- a/color_correction/core/card_detection/det_yv8_onnx.py +++ b/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 diff --git a/color_correction/core/correction/_factory.py b/color_correction/core/correction/_factory.py index 11ed2c6..ed4327a 100644 --- a/color_correction/core/correction/_factory.py +++ b/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 diff --git a/color_correction/core/correction/polynomial.py b/color_correction/core/correction/polynomial.py index 0e0a997..f0052ae 100644 --- a/color_correction/core/correction/polynomial.py +++ b/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,7 +60,7 @@ def fit( Input image patches. y_patches : np.ndarray Reference image patches. - **kwargs : dict + **kwargs : Any Additional keyword arguments. Recognized keyword: - `degree` : int, optional @@ -67,8 +68,8 @@ def fit( Returns ------- - np.ndarray - Fitted model pipeline. + Pipeline + Fitted sklearn pipeline with polynomial features and linear regression. """ start_time = time.perf_counter() diff --git a/color_correction/services/color_correction.py b/color_correction/services/color_correction.py index f6aa240..ad696bb 100644 --- a/color_correction/services/color_correction.py +++ b/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}, " diff --git a/color_correction/services/correction_analyzer.py b/color_correction/services/correction_analyzer.py index 1d226d9..3b388fd 100644 --- a/color_correction/services/correction_analyzer.py +++ b/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,8 +50,8 @@ 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: """ @@ -46,9 +59,9 @@ def __init__( 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,13 +74,13 @@ 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. @@ -75,22 +88,22 @@ def _run_single_exp( ---------- 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 diff --git a/color_correction/utils/device_info.py b/color_correction/utils/device_info.py index 868f95a..fa858c3 100644 --- a/color_correction/utils/device_info.py +++ b/color_correction/utils/device_info.py @@ -1,6 +1,7 @@ import platform import subprocess -from typing import Any +from collections.abc import Callable +from typing import TypedDict from color_correction.schemas.device import ( CPUArchitecture, @@ -9,18 +10,27 @@ ) -def detect_darwin(specs: dict[str, Any]) -> dict[str, Any]: +class DeviceSpecsDict(TypedDict, total=False): + """TypedDict for device specifications before validation.""" + + os_name: str + cpu_arch: CPUArchitecture + gpu_type: GPUType + is_apple_silicon: bool + + +def detect_darwin(specs: DeviceSpecsDict) -> DeviceSpecsDict: """ Detect hardware specifications on macOS, including CPU and GPU details. Parameters ---------- - specs : dict + specs : DeviceSpecsDict Initial dictionary containing OS information. Returns ------- - dict + DeviceSpecsDict Updated dictionary with CPU architecture and GPU type for macOS. """ try: @@ -52,18 +62,18 @@ def detect_darwin(specs: dict[str, Any]) -> dict[str, Any]: return specs -def detect_linux(specs: dict[str, Any]) -> dict[str, Any]: +def detect_linux(specs: DeviceSpecsDict) -> DeviceSpecsDict: """ Detect hardware specifications on Linux systems. Parameters ---------- - specs : dict + specs : DeviceSpecsDict Initial dictionary with OS information. Returns ------- - dict + DeviceSpecsDict Updated dictionary with CPU architecture and GPU type for Linux. """ try: @@ -99,18 +109,18 @@ def detect_linux(specs: dict[str, Any]) -> dict[str, Any]: return specs -def detect_windows(specs: dict[str, Any]) -> dict[str, Any]: +def detect_windows(specs: DeviceSpecsDict) -> DeviceSpecsDict: """ Detect hardware specifications on Windows systems. Parameters ---------- - specs : dict + specs : DeviceSpecsDict Initial dictionary with OS information. Returns ------- - dict + DeviceSpecsDict Updated dictionary with CPU architecture and GPU type for Windows. """ proc = platform.processor().lower() @@ -137,14 +147,14 @@ def get_device_specs() -> DeviceSpecs: An object containing OS name, CPU architecture, GPU type, and Apple Silicon flag. """ - specs = { + specs: DeviceSpecsDict = { "os_name": platform.system(), "cpu_arch": CPUArchitecture.UNKNOWN, "gpu_type": GPUType.UNKNOWN, "is_apple_silicon": False, } - detector_map = { + detector_map: dict[str, Callable[[DeviceSpecsDict], DeviceSpecsDict]] = { "Darwin": detect_darwin, "Linux": detect_linux, "Windows": detect_windows, diff --git a/color_correction/utils/report_generator.py b/color_correction/utils/report_generator.py index 2f95f5e..ae5009d 100644 --- a/color_correction/utils/report_generator.py +++ b/color_correction/utils/report_generator.py @@ -1,6 +1,7 @@ # report_generator.py from datetime import datetime from importlib import resources +from typing import Any import pandas as pd @@ -60,15 +61,15 @@ def generate_report(self, body_report: str) -> str: ) return final_html - def generate_table(self, headers: list, rows: list) -> str: + def generate_table(self, headers: list[Any], rows: list[str]) -> str: """ Generate an HTML table from headers and row data. Parameters ---------- - headers : list - List of table headers. - rows : list + headers : list[Any] + List of table headers (typically strings or column names). + rows : list[str] List of rows where each row is a string of HTML table cells. Returns