lighting_analysis_manager.py

import cv2
import numpy as np
import torch
import torch.nn as nn
from PIL import Image
from typing import Dict, Tuple
import torchvision.models as models
import torchvision.transforms as transforms

class LightingAnalysisManager:
    """Advanced lighting analysis using Places365 scene recognition + CV features"""

    def __init__(self):
        print("Initializing Lighting Analysis Manager with Places365...")

        # Places365 ResNet18
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self._load_places365_model()

        # CV feature weights (Places365 gets higher weight)
        self.feature_weights = {
            'places365': 0.50,     # Primary weight to Places365
            'brightness': 0.15,
            'color_temp': 0.15,
            'contrast': 0.08,
            'gradient': 0.05,      # Auxiliary features
            'laplacian': 0.04,
            'color_variation': 0.03
        }

        print("✓ Lighting Analysis Manager initialized with Places365 + advanced CV features")

    def _load_places365_model(self):
        """Load Places365 ResNet18 for scene attributes"""
        try:
            # Use ResNet18 pretrained on Places365
            model = models.resnet18(weights=None)
            model.fc = nn.Linear(model.fc.in_features, 365)

            # Load Places365 weights (if available, otherwise use ImageNet as fallback)
            try:
                import urllib
                checkpoint_url = 'http://places2.csail.mit.edu/models_places365/resnet18_places365.pth.tar'
                checkpoint = torch.hub.load_state_dict_from_url(
                    checkpoint_url,
                    map_location=self.device,
                    progress=False
                )
                state_dict = {str.replace(k, 'module.', ''): v for k, v in checkpoint['state_dict'].items()}
                model.load_state_dict(state_dict)
                print("      Loaded Places365 ResNet18 weights")
            except:
                print("      Using ImageNet pretrained ResNet18 (fallback)")
                model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)

            model = model.to(self.device)
            model.eval()
            self.places_model = model

            # Image preprocessing for Places365
            self.places_transform = transforms.Compose([
                transforms.Resize((224, 224)),
                transforms.ToTensor(),
                transforms.Normalize(
                    mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225]
                )
            ])

            # Scene categories related to lighting
            self.lighting_scenes = {
                'sunny': ['street', 'downtown', 'plaza', 'park', 'field'],
                'overcast': ['alley', 'covered_bridge', 'corridor'],
                'indoor': ['lobby', 'office', 'museum', 'restaurant'],
                'evening': ['street', 'downtown', 'plaza'],
                'natural': ['park', 'forest', 'mountain', 'coast']
            }

        except Exception as e:
            print(f"      Warning: Places365 loading failed ({e}), using CV-only mode")
            self.places_model = None

    def analyze_lighting(self, image: Image.Image) -> Dict:
        """Comprehensive lighting analysis using Places365 + CV"""

        # 1. CV-based physical features (including advanced features)
        cv_features = self._extract_cv_features(image)

        # 2. Places365 scene understanding (if available)
        scene_info = self._analyze_scene_places365(image)

        # 3. Determine lighting condition (adaptive with auxiliary features)
        lighting_condition, confidence = self._determine_lighting_adaptive(
            cv_features, scene_info
        )

        return {
            'lighting_type': lighting_condition,
            'confidence': confidence,
            'cv_features': cv_features,
            'scene_info': scene_info
        }

    def _extract_cv_features(self, image: Image.Image) -> Dict:
        """Extract CV-based features including advanced gradient and color analysis"""
        img_array = np.array(image)
        img_bgr = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)

        # Basic Features (Primary)
        # Brightness (LAB L-channel)
        lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)
        brightness = float(np.mean(lab[:, :, 0]))

        # Color temperature (R/B ratio)
        b_mean = np.mean(img_bgr[:, :, 0])
        r_mean = np.mean(img_bgr[:, :, 2])
        color_temp = float(r_mean / (b_mean + 1e-6))

        # Contrast (std of grayscale)
        gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
        contrast = float(np.std(gray))

        # Shadow ratio
        _, shadow_mask = cv2.threshold(gray, 80, 255, cv2.THRESH_BINARY_INV)
        shadow_ratio = float(np.sum(shadow_mask > 0) / shadow_mask.size)

        # Advanced Features
        # 1. First derivative: Sobel gradient magnitude (edge strength)
        # Strong gradients suggest directional lighting, weak suggest diffused
        sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
        sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
        gradient_magnitude = np.sqrt(sobelx**2 + sobely**2)
        gradient_strength = float(np.mean(gradient_magnitude))

        # 2. Second derivative: Laplacian variance (lighting change detection)
        # High variance indicates complex lighting with many transitions
        laplacian = cv2.Laplacian(gray, cv2.CV_64F)
        laplacian_var = float(np.var(laplacian))

        # 3. Color difference in LAB space (color uniformity)
        # Low variation suggests overcast/diffused, high suggests mixed lighting
        a_std = float(np.std(lab[:, :, 1]))  # a* channel (green-red)
        b_std = float(np.std(lab[:, :, 2]))  # b* channel (blue-yellow)
        color_variation = (a_std + b_std) / 2

        return {
            # Primary features
            'brightness': brightness,
            'color_temp': color_temp,
            'contrast': contrast,
            'shadow_ratio': shadow_ratio,
            # Advanced auxiliary features (to assist Places365)
            'gradient_strength': gradient_strength,
            'laplacian_variance': laplacian_var,
            'color_variation': color_variation
        }

    def _analyze_scene_places365(self, image: Image.Image) -> Dict:
        """Analyze scene using Places365"""
        if self.places_model is None:
            return {'scene_category': 'unknown', 'confidence': 0.0}

        try:
            with torch.no_grad():
                img_tensor = self.places_transform(image).unsqueeze(0).to(self.device)
                logits = self.places_model(img_tensor)
                probs = torch.nn.functional.softmax(logits, dim=1)

                # Get top prediction
                top_prob, top_idx = torch.max(probs, 1)

                # Simple scene categories
                # Using index ranges for common outdoor scenes
                is_outdoor = top_idx.item() < 200  # Rough heuristic

                return {
                    'scene_category': 'outdoor' if is_outdoor else 'indoor',
                    'confidence': float(top_prob.item()),
                    'scene_idx': int(top_idx.item())
                }
        except Exception as e:
            print(f"      Places365 inference failed: {e}")
            return {'scene_category': 'unknown', 'confidence': 0.0}

    def _detect_indoor_scene(self, cv_features: Dict, scene_info: Dict) -> bool:
        """
        Detect if scene is indoor or outdoor using multiple signals

        Args:
            cv_features: Computer vision features
            scene_info: Places365 scene information

        Returns:
            True if indoor, False if outdoor
        """
        indoor_score = 0.0

        # Signal 1: Places365 scene category (strongest signal)
        if scene_info.get('scene_category') == 'indoor':
            indoor_score += 0.5
        elif scene_info.get('scene_category') == 'outdoor':
            indoor_score -= 0.3

        # Signal 2: Brightness patterns
        # Indoor scenes typically have controlled brightness (not too bright, not too dark)
        brightness = cv_features['brightness']
        if 60 < brightness < 220:  # 放寬範圍，包含更多室內場景
            indoor_score += 0.15
        elif brightness > 230:  # Very bright suggests outdoor
            indoor_score -= 0.2

        # Signal 3: Low gradient suggests controlled/diffused indoor lighting
        gradient = cv_features['gradient_strength']
        if gradient < 20:  # 放寬閾值，更多室內場景符合
            indoor_score += 0.15

        # Signal 4: Low laplacian variance suggests smooth indoor lighting
        laplacian = cv_features['laplacian_variance']
        if laplacian < 400:  # 放寬閾值，包含更多室內場景
            indoor_score += 0.10

        # Signal 5: Shadow ratio - indoor scenes have less harsh shadows
        shadow_ratio = cv_features['shadow_ratio']
        if shadow_ratio < 0.25:  # 放寬閾值，包含更多室內場景
            indoor_score += 0.10
        elif shadow_ratio > 0.5:  # Strong shadows suggest outdoor sunlight
            indoor_score -= 0.15

        # Threshold: indoor if score > 0.15 (降低閾值，更容易判定為室內)
        return indoor_score > 0.15

    def _determine_indoor_lighting(self, cv_features: Dict) -> Tuple[str, float]:
        """
        Determine lighting type for indoor scenes

        Returns indoor-specific lighting types with confidence
        """
        brightness = cv_features['brightness']
        color_temp = cv_features['color_temp']
        contrast = cv_features['contrast']
        shadow_ratio = cv_features['shadow_ratio']
        gradient = cv_features['gradient_strength']
        laplacian = cv_features['laplacian_variance']

        # Normalize features
        brightness_norm = min(brightness / 255.0, 1.0)
        contrast_norm = min(contrast / 100.0, 1.0)
        gradient_norm = min(gradient / 50.0, 1.0)
        laplacian_norm = min(laplacian / 1000.0, 1.0)

        scores = {}

        # Studio/Product Lighting (工作室/產品攝影燈光)
        # Very controlled, bright, minimal shadows, low gradient
        studio_score = (
            0.35 * (1.0 if brightness_norm > 0.6 else 0.5) +  # Bright
            0.25 * (1.0 - shadow_ratio) +                      # Minimal shadows
            0.20 * (1.0 - gradient_norm) +                     # Smooth, even
            0.15 * (1.0 - laplacian_norm) +                   # Very smooth
            0.05 * (1.0 - abs(color_temp - 1.0))              # Neutral temp
        )
        scores['studio lighting'] = studio_score

        # Indoor Natural Light (室內自然光 - 窗光)
        # Medium-bright, some contrast, neutral to warm temp
        natural_indoor_score = (
            0.30 * (1.0 if 0.5 < brightness_norm < 0.8 else 0.5) +  # Medium-bright
            0.25 * min(contrast_norm, 0.6) +                         # Some contrast
            0.20 * (1.0 if color_temp > 0.95 else 0.5) +            # Neutral to warm
            0.15 * min(gradient_norm, 0.5) +                         # Some direction
            0.10 * (1.0 if shadow_ratio < 0.3 else 0.5)             # Some shadows
        )
        scores['indoor natural light'] = natural_indoor_score

        # Warm Artificial Lighting (溫暖人工照明)
        # Warm color temp, medium brightness, soft
        warm_artificial_score = (
            0.35 * (1.0 if color_temp > 1.1 else 0.3) +              # Warm temp
            0.25 * (1.0 - abs(brightness_norm - 0.5)) +              # Medium brightness
            0.20 * (1.0 - gradient_norm) +                           # Soft
            0.15 * (1.0 - shadow_ratio) +                            # Minimal shadows
            0.05 * (1.0 - laplacian_norm)                            # Smooth
        )
        scores['warm artificial lighting'] = warm_artificial_score

        # Cool Artificial Lighting (冷色人工照明)
        # Cool/neutral temp, medium-bright
        cool_artificial_score = (
            0.35 * (1.0 if color_temp < 1.05 else 0.4) +            # Cool/neutral temp
            0.25 * (1.0 if brightness_norm > 0.5 else 0.5) +        # Medium-bright
            0.20 * (1.0 - gradient_norm) +                           # Smooth
            0.15 * (1.0 - shadow_ratio) +                            # Minimal shadows
            0.05 * (1.0 - laplacian_norm)                            # Even
        )
        scores['cool artificial lighting'] = cool_artificial_score

        # Soft Indoor Lighting (柔和室內光線)
        # Low contrast, diffused, medium brightness
        soft_indoor_score = (
            0.30 * (1.0 - abs(brightness_norm - 0.5)) +             # Medium brightness
            0.30 * (1.0 - contrast_norm) +                           # Low contrast
            0.20 * (1.0 - gradient_norm) +                           # Very soft
            0.15 * (1.0 - shadow_ratio) +                            # Minimal shadows
            0.05 * (1.0 - laplacian_norm)                            # Smooth
        )
        scores['soft indoor lighting'] = soft_indoor_score

        # Dramatic Indoor Lighting (戲劇性室內光線)
        # High contrast, directional, some shadows
        dramatic_score = (
            0.35 * contrast_norm +                                   # High contrast
            0.25 * gradient_norm +                                   # Directional
            0.20 * shadow_ratio +                                    # Shadows present
            0.15 * laplacian_norm +                                  # Sharp transitions
            0.05 * (1.0 if brightness_norm < 0.6 else 0.5)          # Can be darker
        )
        scores['dramatic indoor lighting'] = dramatic_score

        # Get best match
        best_condition = max(scores.items(), key=lambda x: x[1])

        # Calculate confidence
        sorted_scores = sorted(scores.values(), reverse=True)
        if len(sorted_scores) > 1:
            score_gap = sorted_scores[0] - sorted_scores[1]
            confidence = min(0.7 + score_gap * 0.3, 0.95)
        else:
            confidence = 0.7

        return best_condition[0], confidence

    def _determine_lighting_adaptive(self, cv_features: Dict, scene_info: Dict) -> Tuple[str, float]:
        """Determine lighting using adaptive thresholds with indoor/outdoor detection"""

        # Extract all features
        brightness = cv_features['brightness']
        color_temp = cv_features['color_temp']
        contrast = cv_features['contrast']
        shadow = cv_features['shadow_ratio']
        gradient = cv_features['gradient_strength']
        laplacian = cv_features['laplacian_variance']
        color_var = cv_features['color_variation']

        # NEW: Detect indoor vs outdoor
        is_indoor = self._detect_indoor_scene(cv_features, scene_info)
        if is_indoor:
            # 室內場景優先使用室內光線類型
            return self._determine_indoor_lighting(cv_features)
        # 否則使用原有邏輯

        # Normalize features to 0-1 scale
        brightness_norm = min(brightness / 255.0, 1.0)
        contrast_norm = min(contrast / 100.0, 1.0)
        gradient_norm = min(gradient / 50.0, 1.0)  # Typical range 0-50
        laplacian_norm = min(laplacian / 1000.0, 1.0)  # Typical range 0-1000
        color_var_norm = min(color_var / 50.0, 1.0)  # Typical range 0-50

        # Adaptive scoring (Places365 dominant, CV features assist)
        scores = {}

        # Soft diffused light (柔和漫射光)
        # Characteristics: medium brightness, low contrast, neutral temp
        # Auxiliary: low gradient (no strong edges), low laplacian (smooth transitions)
        diffuse_score = (
            0.40 * (1.0 - abs(brightness_norm - 0.5)) +  # Medium brightness
            0.25 * (1.0 - contrast_norm) +                # Low contrast
            0.20 * (1.0 - abs(color_temp - 1.0)) +       # Neutral temp
            0.08 * (1.0 - gradient_norm) +                # Weak edges (diffused)
            0.05 * (1.0 - laplacian_norm) +              # Smooth transitions
            0.02 * (1.0 - color_var_norm)                # Uniform color
        )
        scores['soft diffused light'] = diffuse_score

        # Natural daylight (自然光)
        # Characteristics: bright, moderate contrast
        # Auxiliary: moderate gradient, moderate color variation
        daylight_score = (
            0.40 * brightness_norm +                      # Bright
            0.25 * min(contrast_norm, 0.7) +             # Moderate contrast
            0.20 * (1.0 - abs(color_temp - 1.0)) +       # Neutral temp
            0.08 * min(gradient_norm, 0.6) +             # Moderate edges
            0.05 * min(laplacian_norm, 0.6) +            # Some detail
            0.02 * min(color_var_norm, 0.5)              # Some color variation
        )
        scores['natural daylight'] = daylight_score

        # Overcast atmosphere (陰天氛圍)
        # Characteristics: medium-low brightness, very low contrast, cool temp, minimal shadow
        # Auxiliary: very low gradient (flat), low laplacian, low color variation
        overcast_score = (
            0.35 * (1.0 - abs(brightness_norm - 0.45)) + # Medium-low brightness
            0.25 * (1.0 - contrast_norm) +                # Very low contrast
            0.15 * (1.0 if color_temp < 1.05 else 0.5) + # Cool temp
            0.10 * (1.0 - shadow) +                       # Minimal shadows
            0.08 * (1.0 - gradient_norm) +                # Flat appearance
            0.05 * (1.0 - laplacian_norm) +              # Smooth lighting
            0.02 * (1.0 - color_var_norm)                # Uniform color
        )
        scores['overcast atmosphere'] = overcast_score

        # Warm ambient light (溫暖環境光)
        # Characteristics: medium brightness, warm temp
        # Auxiliary: moderate gradient, warm color bias
        warm_score = (
            0.40 * (1.0 - abs(brightness_norm - 0.5)) +  # Medium brightness
            0.30 * (1.0 if color_temp > 1.1 else 0.5) +  # Warm temp
            0.15 * min(contrast_norm, 0.6) +             # Moderate contrast
            0.08 * min(gradient_norm, 0.5) +             # Soft edges
            0.05 * min(laplacian_norm, 0.5) +            # Soft transitions
            0.02 * color_var_norm                        # Some color variation (warmth)
        )
        scores['warm ambient light'] = warm_score

        # Evening light (傍晚光線)
        # Characteristics: medium-low brightness, warm temp, medium contrast
        # Auxiliary: moderate gradient (directional), some color variation
        evening_score = (
            0.35 * (1.0 if brightness_norm < 0.6 else 0.5) +  # Lower brightness
            0.30 * (1.0 if color_temp > 1.05 else 0.5) +      # Slightly warm
            0.20 * contrast_norm +                             # Some contrast
            0.08 * min(gradient_norm, 0.7) +                   # Directional light
            0.05 * laplacian_norm +                            # Detail present
            0.02 * color_var_norm                              # Color variation
        )
        scores['evening light'] = evening_score

        # Bright sunlight (明亮陽光)
        # Characteristics: high brightness, high contrast, strong shadows
        # Auxiliary: high gradient (strong edges), high laplacian (sharp transitions)
        sunlight_score = (
            0.40 * (1.0 if brightness_norm > 0.7 else 0.3) +  # High brightness
            0.25 * contrast_norm +                             # High contrast
            0.15 * shadow +                                    # Strong shadows
            0.10 * gradient_norm +                             # Strong edges
            0.08 * laplacian_norm +                            # Sharp detail
            0.02 * color_var_norm                              # Color variation
        )
        scores['bright sunlight'] = sunlight_score

        # Get top scoring condition
        best_condition = max(scores.items(), key=lambda x: x[1])

        # Calculate confidence based on score separation
        sorted_scores = sorted(scores.values(), reverse=True)
        if len(sorted_scores) > 1:
            score_gap = sorted_scores[0] - sorted_scores[1]
            confidence = min(0.7 + score_gap * 0.3, 0.95)
        else:
            confidence = 0.7

        return best_condition[0], confidence

print("✓ LightingAnalysisManager (with Places365 + advanced CV features) defined")