logos/image_analyzer.py

"""
image_analyzer.py - LOGOS Image Analysis Pipeline
Batch-process architectural diagrams and UI screenshots.
"""

import os
import cv2
import numpy as np
from PIL import Image
from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass
from collections import Counter


@dataclass
class ImageAnalysis:
    """Result of analyzing a single image."""
    filename: str
    path: str
    width: int
    height: int
    aspect_ratio: float
    classification: str  # "diagram", "ui", "photo", "other"
    dominant_colors: List[Tuple[int, int, int]]
    edge_density: float
    text_region_ratio: float
    thumbnail: Optional[np.ndarray] = None


def classify_image(edge_density: float, color_variance: float, aspect: float) -> str:
    """
    Simple heuristic classification.
    - Diagrams: High edge density, low color variance, often wide aspect.
    - UI: Medium edge density, structured colors, standard aspect.
    - Photos: Low edge density, high color variance.
    """
    if edge_density > 0.15 and color_variance < 50:
        return "diagram"
    elif 0.05 < edge_density < 0.20 and 0.5 < aspect < 2.0:
        return "ui"
    elif color_variance > 80:
        return "photo"
    else:
        return "other"


def get_dominant_colors(image: np.ndarray, k: int = 3) -> List[Tuple[int, int, int]]:
    """Extract k dominant colors using k-means clustering."""
    pixels = image.reshape(-1, 3).astype(np.float32)
    
    # Subsample for speed
    if len(pixels) > 10000:
        indices = np.random.choice(len(pixels), 10000, replace=False)
        pixels = pixels[indices]
    
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
    _, labels, centers = cv2.kmeans(pixels, k, None, criteria, 3, cv2.KMEANS_PP_CENTERS)
    
    # Sort by frequency
    counts = Counter(labels.flatten())
    sorted_centers = [centers[i] for i, _ in counts.most_common(k)]
    
    return [(int(c[2]), int(c[1]), int(c[0])) for c in sorted_centers]  # BGR -> RGB


def calculate_edge_density(gray: np.ndarray) -> float:
    """Calculate edge density using Canny edge detection."""
    edges = cv2.Canny(gray, 50, 150)
    return np.count_nonzero(edges) / edges.size


def estimate_text_regions(gray: np.ndarray) -> float:
    """
    Estimate ratio of image containing text-like regions.
    Uses morphological operations to find text blocks.
    """
    # Threshold
    _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    
    # Dilate to connect text characters
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 3))
    dilated = cv2.dilate(binary, kernel, iterations=2)
    
    # Find contours
    contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # Filter by aspect ratio (text regions are usually wide)
    text_area = 0
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        if w > h * 2 and w > 20:  # Wide and reasonably sized
            text_area += w * h
    
    return text_area / (gray.shape[0] * gray.shape[1])


def analyze_image(path: str, thumbnail_size: int = 128) -> ImageAnalysis:
    """
    Analyze a single image and return structured metadata.
    """
    img = cv2.imread(path)
    if img is None:
        raise ValueError(f"Could not load image: {path}")
    
    height, width = img.shape[:2]
    aspect = width / height
    
    # Convert to RGB for color analysis
    rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # Metrics
    dominant_colors = get_dominant_colors(rgb)
    edge_density = calculate_edge_density(gray)
    text_ratio = estimate_text_regions(gray)
    color_variance = np.std(rgb)
    
    # Classification
    classification = classify_image(edge_density, color_variance, aspect)
    
    # Thumbnail
    scale = thumbnail_size / max(width, height)
    thumb = cv2.resize(rgb, (int(width * scale), int(height * scale)))
    
    return ImageAnalysis(
        filename=os.path.basename(path),
        path=path,
        width=width,
        height=height,
        aspect_ratio=round(aspect, 2),
        classification=classification,
        dominant_colors=dominant_colors,
        edge_density=round(edge_density, 4),
        text_region_ratio=round(text_ratio, 4),
        thumbnail=thumb
    )


def batch_analyze(folder: str, extensions: List[str] = None) -> List[ImageAnalysis]:
    """
    Analyze all images in a folder.
    
    Args:
        folder: Path to folder containing images.
        extensions: List of valid extensions (default: ['.png', '.jpg', '.jpeg']).
    
    Returns:
        List of ImageAnalysis results.
    """
    if extensions is None:
        extensions = ['.png', '.jpg', '.jpeg', '.bmp', '.webp']
    
    results = []
    
    for filename in os.listdir(folder):
        ext = os.path.splitext(filename)[1].lower()
        if ext in extensions:
            path = os.path.join(folder, filename)
            try:
                analysis = analyze_image(path)
                results.append(analysis)
            except Exception as e:
                print(f"[ANALYZER] Error processing {filename}: {e}")
    
    return results


def summarize_analysis(results: List[ImageAnalysis]) -> Dict:
    """Generate summary statistics from batch analysis."""
    if not results:
        return {"count": 0}
    
    classifications = Counter(r.classification for r in results)
    avg_edge = sum(r.edge_density for r in results) / len(results)
    avg_text = sum(r.text_region_ratio for r in results) / len(results)
    
    return {
        "count": len(results),
        "classifications": dict(classifications),
        "avg_edge_density": round(avg_edge, 4),
        "avg_text_ratio": round(avg_text, 4),
        "total_size_mb": round(sum(r.width * r.height * 3 for r in results) / (1024 * 1024), 2)
    }


# CLI for standalone testing
if __name__ == "__main__":
    import sys
    
    if len(sys.argv) < 2:
        print("Usage: python image_analyzer.py <folder_path>")
        sys.exit(1)
    
    folder = sys.argv[1]
    print(f"[ANALYZER] Processing folder: {folder}")
    
    results = batch_analyze(folder)
    summary = summarize_analysis(results)
    
    print(f"\n[SUMMARY]")
    print(f"  Total Images: {summary['count']}")
    print(f"  Classifications: {summary['classifications']}")
    print(f"  Avg Edge Density: {summary['avg_edge_density']}")
    print(f"  Avg Text Ratio: {summary['avg_text_ratio']}")