AI-Video/pixelle_video/services/quality/objective_metrics.py

# Copyright (C) 2025 AIDC-AI
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
ObjectiveMetrics - Technical image quality metrics

Provides fast, local computation of:
1. Sharpness/clarity detection
2. Color distribution analysis
3. Exposure detection (over/under)

Only depends on PIL and numpy (no heavy ML dependencies).
"""

from dataclasses import dataclass, field
from typing import List, Tuple

import numpy as np
from loguru import logger


@dataclass
class TechnicalMetrics:
    """Technical quality metrics for an image"""

    # Sharpness (0.0-1.0, higher = sharper)
    sharpness_score: float = 0.0

    # Color metrics
    color_variance: float = 0.0
    saturation_score: float = 0.0

    # Exposure metrics
    brightness_score: float = 0.0
    is_overexposed: bool = False
    is_underexposed: bool = False

    # Overall technical score
    overall_technical: float = 0.0

    # Issues detected
    issues: List[str] = field(default_factory=list)

    def to_dict(self) -> dict:
        """Convert to dictionary"""
        return {
            "sharpness_score": self.sharpness_score,
            "color_variance": self.color_variance,
            "saturation_score": self.saturation_score,
            "brightness_score": self.brightness_score,
            "is_overexposed": self.is_overexposed,
            "is_underexposed": self.is_underexposed,
            "overall_technical": self.overall_technical,
            "issues": self.issues,
        }


class ObjectiveMetricsCalculator:
    """
    Calculate objective technical quality metrics

    All computations are local and fast (no external API calls).
    Uses PIL and numpy only.

    Example:
        >>> calculator = ObjectiveMetricsCalculator()
        >>> metrics = calculator.analyze_image("frame_001.png")
        >>> print(f"Sharpness: {metrics.sharpness_score:.2f}")
    """

    def __init__(
        self,
        sharpness_threshold: float = 0.3,
        overexposure_threshold: float = 0.85,
        underexposure_threshold: float = 0.15,
    ):
        self.sharpness_threshold = sharpness_threshold
        self.overexposure_threshold = overexposure_threshold
        self.underexposure_threshold = underexposure_threshold

    def analyze_image(self, image_path: str) -> TechnicalMetrics:
        """
        Analyze image and return technical metrics

        Args:
            image_path: Path to image file

        Returns:
            TechnicalMetrics with all computed values
        """
        try:
            from PIL import Image

            with Image.open(image_path) as img:
                img_rgb = img.convert("RGB")
                img_array = np.array(img_rgb)

                # Calculate individual metrics
                sharpness = self._calculate_sharpness(img_rgb)
                color_var, saturation = self._calculate_color_metrics(img_array)
                brightness, overexp, underexp = self._calculate_exposure(img_array)

                # Detect issues
                issues = []
                if sharpness < self.sharpness_threshold:
                    issues.append("Image appears blurry")
                if overexp:
                    issues.append("Image is overexposed")
                if underexp:
                    issues.append("Image is underexposed")
                if color_var < 0.1:
                    issues.append("Low color diversity")

                # Calculate overall technical score
                overall = self._calculate_overall(
                    sharpness, color_var, saturation, brightness
                )

                return TechnicalMetrics(
                    sharpness_score=sharpness,
                    color_variance=color_var,
                    saturation_score=saturation,
                    brightness_score=brightness,
                    is_overexposed=overexp,
                    is_underexposed=underexp,
                    overall_technical=overall,
                    issues=issues
                )

        except Exception as e:
            logger.warning(f"Failed to analyze image: {e}")
            return TechnicalMetrics(issues=[f"Analysis failed: {str(e)}"])

    def _calculate_sharpness(self, img) -> float:
        """Calculate sharpness using edge detection"""
        from PIL import ImageFilter

        # Convert to grayscale
        gray = img.convert("L")

        # Apply edge detection
        edges = gray.filter(ImageFilter.FIND_EDGES)
        edge_array = np.array(edges).astype(np.float32)

        # Variance of edge response
        variance = np.var(edge_array)

        # Normalize to 0-1 (empirical scaling)
        sharpness = min(1.0, variance / 2000)

        return float(sharpness)

    def _calculate_color_metrics(
        self,
        img_array: np.ndarray
    ) -> Tuple[float, float]:
        """Calculate color variance and saturation"""
        r, g, b = img_array[:, :, 0], img_array[:, :, 1], img_array[:, :, 2]

        max_rgb = np.maximum(np.maximum(r, g), b)
        min_rgb = np.minimum(np.minimum(r, g), b)

        # Saturation
        delta = max_rgb - min_rgb
        saturation = np.where(max_rgb > 0, delta / (max_rgb + 1e-7), 0)
        avg_saturation = np.mean(saturation)

        # Color variance (diversity)
        color_std = np.std(img_array.reshape(-1, 3), axis=0)
        color_variance = np.mean(color_std) / 128

        return float(color_variance), float(avg_saturation)

    def _calculate_exposure(
        self,
        img_array: np.ndarray
    ) -> Tuple[float, bool, bool]:
        """Calculate brightness and detect exposure issues"""
        # Calculate luminance
        luminance = (
            0.299 * img_array[:, :, 0] +
            0.587 * img_array[:, :, 1] +
            0.114 * img_array[:, :, 2]
        ) / 255.0

        avg_brightness = float(np.mean(luminance))

        # Check for over/under exposure
        overexposed_pixels = np.mean(luminance > 0.95)
        underexposed_pixels = np.mean(luminance < 0.05)

        is_overexposed = (
            avg_brightness > self.overexposure_threshold or
            overexposed_pixels > 0.3
        )
        is_underexposed = (
            avg_brightness < self.underexposure_threshold or
            underexposed_pixels > 0.3
        )

        return avg_brightness, is_overexposed, is_underexposed

    def _calculate_overall(
        self,
        sharpness: float,
        color_var: float,
        saturation: float,
        brightness: float
    ) -> float:
        """Calculate overall technical score"""
        # Brightness penalty (ideal is 0.5)
        brightness_score = 1.0 - abs(brightness - 0.5) * 2
        brightness_score = max(0, brightness_score)

        # Weighted combination
        overall = (
            sharpness * 0.4 +
            min(1.0, color_var * 2) * 0.2 +
            saturation * 0.2 +
            brightness_score * 0.2
        )

        return float(overall)