Numba_Scripts/main.py

"""
main.py - OPTIMIZED
Complete mosaic generator with Numba + parallel optimizations.
"""

import numpy as np
import cv2
from pathlib import Path
from typing import Tuple, Dict, Optional
import time
import matplotlib.pyplot as plt

from ImagePreprocessor import ImagePreprocessor
from contextual_Mosaic_Builder import ContextualMosaicGenerator, ImageContext
from Performance_metrics import PerformanceEvaluator


def create_advanced_mosaic(image_path: str,
                          tile_folder: str = "extracted_images", 
                          grid_size: Tuple[int, int] = (64, 64),
                          tile_size: Tuple[int, int] = (32, 32),
                          diversity_factor: float = 0.15,
                          enable_rotation: bool = True,
                          enable_face_detection: bool = False,
                          colour_bins: int = 8,
                          apply_quantization: bool = False,
                          n_colors: int = 12,
                          use_numba: bool = True,
                          use_parallel: bool = True,
                          evaluate_quality: bool = False,
                          show_visualizations: bool = False,
                          verbose: bool = False) -> Tuple[np.ndarray, Optional[Dict], Optional[ImageContext]]:
    """
    Create advanced mosaic - FULLY OPTIMIZED.
    
    Args:
        image_path: Path to input image
        tile_folder: Path to tile images folder
        grid_size: Grid dimensions (rows, cols)
        tile_size: Individual tile size (width, height)
        diversity_factor: Tile diversity factor (0.0-0.5)
        enable_rotation: Whether to enable 4-way tile rotation
        enable_face_detection: Whether to enable face detection
        colour_bins: Number of color bins for subgrouping
        apply_quantization: Whether to apply color quantization
        n_colors: Number of colors for quantization
        use_numba: Use Numba JIT compilation (recommended)
        use_parallel: Use parallel processing (recommended)
        evaluate_quality: Whether to calculate quality metrics
        show_visualizations: Whether to display analysis
        verbose: Enable detailed logging
        
    Returns:
        Tuple of (mosaic_image, quality_metrics, context_analysis)
    """
    
    if verbose:
        print(f"Advanced Mosaic Generator (Numba={'ON' if use_numba else 'OFF'}, Parallel={'ON' if use_parallel else 'OFF'})")
        print(f"Image: {image_path}")
        print(f"Grid: {grid_size[0]}x{grid_size[1]} = {np.prod(grid_size):,} tiles")
        print(f"Tile size: {tile_size[0]}x{tile_size[1]}")
    
    total_start = time.time()
    
    # Initialize generator
    cache_filename = f"cache_{tile_size[0]}x{tile_size[1]}_bins{colour_bins}{'_rot' if enable_rotation else ''}.pkl"
    
    generator = ContextualMosaicGenerator(
        cache_file=cache_filename,
        tile_folder=tile_folder,
        tile_size=tile_size,
        colour_bins=colour_bins,
        enable_rotation=enable_rotation,
        enable_face_detection=enable_face_detection,
        verbose=verbose,
        use_numba=use_numba
    )
    
    # Preprocess image
    target_width = grid_size[1] * tile_size[0]
    target_height = grid_size[0] * tile_size[1]
    
    preprocessor = ImagePreprocessor(
        target_resolution=(target_width, target_height),
        grid_size=grid_size,
        verbose=verbose
    )
    
    processed_image = preprocessor.load_and_preprocess_image(
        image_path,
        apply_quantization=apply_quantization,
        n_colors=n_colors
    )
    
    if processed_image is None:
        raise ValueError("Failed to preprocess image")
    
    # Analyze context (optional)
    context = None
    if enable_face_detection or show_visualizations:
        context = generator.analyze_image_context(processed_image)
    
    # Generate mosaic
    mosaic, context = generator.create_contextual_mosaic(
        processed_image, 
        grid_size, 
        diversity_factor,
        context=context,
        use_parallel=use_parallel
    )
    
    # Visualizations (optional)
    if show_visualizations and context is not None:
        generator.visualize_context_analysis(processed_image, context)
    
    # Quality evaluation (optional)
    metrics = None
    if evaluate_quality:
        evaluator = PerformanceEvaluator()
        metrics = evaluator.evaluate_mosaic_quality(processed_image, mosaic, image_path)
        
        if show_visualizations:
            evaluator.visualize_quality_comparison(processed_image, mosaic, metrics)
    
    total_time = time.time() - total_start
    
    if verbose:
        print(f"\n✅ Total time: {total_time:.2f}s")
        if metrics:
            print(f"Quality score: {metrics['overall_quality']:.1f}/100")
    
    return mosaic, metrics, context


def create_mosaic_for_gradio(image_array: np.ndarray,
                           tile_folder: str = "extracted_images",
                           grid_size: int = 64,
                           tile_size: int = 32,
                           diversity_factor: float = 0.15,
                           enable_rotation: bool = True) -> Tuple[np.ndarray, str]:
    """Gradio-optimized mosaic creation."""
    try:
        grid_tuple = (grid_size, grid_size)
        tile_tuple = (tile_size, tile_size)
        
        temp_path = "temp_gradio_input.jpg"
        cv2.imwrite(temp_path, cv2.cvtColor(image_array, cv2.COLOR_RGB2BGR))
        
        mosaic, metrics, context = create_advanced_mosaic(
            image_path=temp_path,
            tile_folder=tile_folder,
            grid_size=grid_tuple,
            tile_size=tile_tuple,
            diversity_factor=diversity_factor,
            enable_rotation=enable_rotation,
            enable_face_detection=False,
            use_numba=True,
            use_parallel=True,
            evaluate_quality=True,
            show_visualizations=False,
            verbose=False
        )
        
        Path(temp_path).unlink(missing_ok=True)
        
        status = f"""Mosaic Generated Successfully!

Configuration:
• Grid: {grid_size}×{grid_size} = {grid_size**2:,} tiles
• Tile size: {tile_size}×{tile_size} pixels
• Optimizations: Numba JIT + Parallel
• Quality score: {metrics['overall_quality']:.1f}/100"""
        
        return mosaic, status
        
    except Exception as e:
        return None, f"Error: {str(e)}"


if __name__ == "__main__":
    IMAGE_PATH = "Images/EmmaPotrait.jpg"
    TILE_FOLDER = "extracted_images"
    
    if not Path(TILE_FOLDER).exists():
        print(f"Tile folder not found: {TILE_FOLDER}")
        exit(1)
    
    if not Path(IMAGE_PATH).exists():
        print(f"Image not found: {IMAGE_PATH}")
        exit(1)
    
    print("Creating optimized mosaic (Numba + Parallel)...")
    
    mosaic, metrics, context = create_advanced_mosaic(
        image_path=IMAGE_PATH,
        tile_folder=TILE_FOLDER,
        grid_size=(64, 64),
        tile_size=(64, 64),
        diversity_factor=0.15,
        enable_rotation=True,
        enable_face_detection=False,
        apply_quantization=False,
        use_numba=True,
        use_parallel=True,
        evaluate_quality=True,
        show_visualizations=False,
        verbose=True
    )
    
    if mosaic is not None:
        print("✅ Mosaic generation completed!")
        if metrics:
            print(f"Quality score: {metrics['overall_quality']:.1f}/100")
        
        cv2.imwrite("output_mosaic.jpg", cv2.cvtColor(mosaic, cv2.COLOR_RGB2BGR))
        print("Saved: output_mosaic.jpg")