test.py

#!/usr/bin/env python3
"""
Comprehensive test script for Interactive Image Mosaic Generator
Tests multiple parameter combinations and saves results with performance metrics
"""

import os
import json
import csv
from pathlib import Path
from itertools import product
import tempfile
from tqdm import tqdm
import pandas as pd

from simple_mosaic import SimpleMosaicImage
from tile_library import build_cifar10_tile_library, build_cifar100_tile_library
from performance import calculate_all_metrics
from PIL import Image

# Test parameters
TEST_PARAMS = {
    'min_size': [4, 16, 32],
    'start_size': [32, 64, 128],  # max size
    'threshold': [0.1, 5.0],      # subdivision threshold
    'color_quantization': [0, 16, 32],
    'tile_library': ['none', 'cifar-10', 'cifar-100']
}

# Input and output directories
SAMPLES_DIR = Path('./samples')
OUTPUT_DIR = Path('./output/test_result')
MOSAIC_DIR = OUTPUT_DIR / 'mosaic_images'
METRICS_DIR = OUTPUT_DIR / 'metrics'

# Tile cache to avoid reloading
tile_cache = {}

def setup_directories():
    """Create necessary output directories"""
    for dir_path in [OUTPUT_DIR, MOSAIC_DIR, METRICS_DIR]:
        dir_path.mkdir(parents=True, exist_ok=True)
    print(f"[INFO] Created output directories under {OUTPUT_DIR}")

def get_tile_library(tile_type, max_per_class=500):
    """Get cached tile library or create new one"""
    if tile_type == 'none':
        return None, None, None

    cache_key = f"{tile_type}_{max_per_class}"

    if cache_key not in tile_cache:
        print(f"[INFO] Loading {tile_type} tile library...")
        if tile_type == "cifar-10":
            tiles, means, labels = build_cifar10_tile_library(max_per_class=max_per_class)
        elif tile_type == "cifar-100":
            tiles, means, labels = build_cifar100_tile_library(max_per_class=max_per_class)
        else:
            return None, None, None

        tile_cache[cache_key] = (tiles, means, labels)
        print(f"[INFO] Cached {len(tiles)} tiles for {tile_type}")

    return tile_cache[cache_key]

def generate_filename(image_name, min_size, start_size, threshold, color_quantization, tile_library):
    """Generate standardized filename for results"""
    # Clean image name (remove extension)
    base_name = Path(image_name).stem

    # Format threshold to avoid decimal issues
    threshold_str = f"{threshold:g}".replace('.', 'p')

    # Create filename
    filename = f"{base_name}-min{min_size}-max{start_size}-sub{threshold_str}-quant{color_quantization}-{tile_library}"

    return filename

def process_single_combination(image_path, params, progress_bar=None):
    """Process a single parameter combination and return results"""
    image_name = image_path.name
    min_size, start_size, threshold, color_quantization, tile_library = params

    try:
        # Load original image
        original_image = Image.open(image_path).convert("RGB")

        # Create temporary file for processing
        with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp_file:
            original_image.save(tmp_file.name)

            # Load and process image
            loader = SimpleMosaicImage(tmp_file.name)

            # Apply resize if needed (use existing MAX_IMAGE_SIZE logic)
            MAX_IMAGE_SIZE = 4500
            if max(loader.width, loader.height) > MAX_IMAGE_SIZE:
                loader.resize(MAX_IMAGE_SIZE)

            # Apply color quantization if requested
            if color_quantization > 0:
                loader.quantize_colors(color_quantization)

            # Smart boundary handling
            loader.crop_to_grid(2)

            # Build adaptive cells
            cells = loader.build_adaptive_cells(
                start_size=start_size,
                min_size=min_size,
                threshold=threshold
            )

            # Generate mosaic based on tile type
            if tile_library == "none":
                result_img = loader.mosaic_average_color_adaptive(cells)
                processing_info = f"Average colors with {len(cells)} adaptive cells"
            else:
                tiles, tile_means, _ = get_tile_library(tile_library, max_per_class=500)
                if tiles is None:
                    raise ValueError(f"Failed to load {tile_library} tile library")

                result_img = loader.mosaic_with_tiles_adaptive(cells, tiles, tile_means)
                processing_info = f"{tile_library} with {len(cells)} cells using {len(tiles)} tiles"

            # Calculate metrics
            metrics = calculate_all_metrics(original_image, result_img)

            # Generate filenames
            base_filename = generate_filename(
                image_name, min_size, start_size, threshold, color_quantization, tile_library
            )

            # Save mosaic image
            mosaic_path = MOSAIC_DIR / f"{base_filename}.jpg"
            result_img.save(mosaic_path, quality=95)

            # Save metrics
            metrics_path = METRICS_DIR / f"{base_filename}_metrics.json"
            with open(metrics_path, 'w') as f:
                json.dump(metrics, f, indent=2)

            # Clean up temporary file
            os.unlink(tmp_file.name)

            # Update progress bar
            if progress_bar:
                progress_bar.set_postfix({
                    'image': image_name[:15],
                    'params': f"min{min_size}_max{start_size}_sub{threshold:g}_q{color_quantization}_{tile_library[:6]}"
                })
                progress_bar.update(1)

            return {
                'image_name': image_name,
                'min_size': min_size,
                'start_size': start_size,
                'threshold': threshold,
                'color_quantization': color_quantization,
                'tile_library': tile_library,
                'num_cells': len(cells),
                'processing_info': processing_info,
                'mosaic_path': str(mosaic_path),
                'metrics_path': str(metrics_path),
                'success': True,
                **metrics
            }

    except Exception as e:
        error_msg = f"Error processing {image_name} with params {params}: {str(e)}"
        print(f"[ERROR] {error_msg}")

        if progress_bar:
            progress_bar.update(1)

        return {
            'image_name': image_name,
            'min_size': min_size,
            'start_size': start_size,
            'threshold': threshold,
            'color_quantization': color_quantization,
            'tile_library': tile_library,
            'success': False,
            'error': str(e),
            'MSE': float('nan'),
            'SSIM': float('nan'),
            'Histogram_Correlation': float('nan'),
            'Edge_Similarity': float('nan'),
            'Overall_Quality': float('nan')
        }

def run_comprehensive_test():
    """Run comprehensive test across all parameter combinations"""
    setup_directories()

    # Find all image files in samples directory
    image_files = []
    for ext in ['*.jpg', '*.jpeg', '*.JPG', '*.JPEG', '*.png', '*.PNG']:
        image_files.extend(list(SAMPLES_DIR.glob(ext)))

    if not image_files:
        print(f"[ERROR] No image files found in {SAMPLES_DIR}")
        return

    print(f"[INFO] Found {len(image_files)} images to process")
    print(f"[INFO] Images: {[f.name for f in image_files]}")

    # Generate all parameter combinations
    param_combinations = list(product(
        TEST_PARAMS['min_size'],
        TEST_PARAMS['start_size'],
        TEST_PARAMS['threshold'],
        TEST_PARAMS['color_quantization'],
        TEST_PARAMS['tile_library']
    ))

    # Filter out invalid combinations (min_size >= start_size)
    valid_combinations = [
        combo for combo in param_combinations
        if combo[0] < combo[1]  # min_size < start_size
    ]

    total_tests = len(image_files) * len(valid_combinations)
    print(f"[INFO] Total test combinations: {total_tests}")
    print(f"[INFO] Parameter combinations per image: {len(valid_combinations)}")

    # Results storage
    all_results = []

    # Create progress bar
    with tqdm(total=total_tests, desc="Processing mosaics", unit="test") as pbar:
        for image_path in image_files:
            print(f"\n[INFO] Processing image: {image_path.name}")

            for params in valid_combinations:
                result = process_single_combination(image_path, params, pbar)
                all_results.append(result)

    # Save comprehensive results to CSV
    summary_path = OUTPUT_DIR / 'summary.csv'
    df = pd.DataFrame(all_results)
    df.to_csv(summary_path, index=False)

    # Print summary statistics
    successful_tests = df[df['success'] == True]
    failed_tests = df[df['success'] == False]

    print(f"\n[SUMMARY]")
    print(f"Total tests run: {len(all_results)}")
    print(f"Successful: {len(successful_tests)}")
    print(f"Failed: {len(failed_tests)}")

    if len(successful_tests) > 0:
        print(f"\nPerformance Statistics (successful tests):")
        print(f"Average MSE: {successful_tests['MSE'].mean():.2f}")
        print(f"Average SSIM: {successful_tests['SSIM'].mean():.4f}")
        print(f"Average Histogram Correlation: {successful_tests['Histogram_Correlation'].mean():.4f}")
        print(f"Average Edge Similarity: {successful_tests['Edge_Similarity'].mean():.4f}")
        print(f"Average Overall Quality: {successful_tests['Overall_Quality'].mean():.4f}")

    print(f"\nResults saved to:")
    print(f"- Summary CSV: {summary_path}")
    print(f"- Mosaic images: {MOSAIC_DIR}")
    print(f"- Individual metrics: {METRICS_DIR}")

    return all_results

if __name__ == "__main__":
    print("=" * 60)
    print("Interactive Image Mosaic Generator - Comprehensive Test")
    print("=" * 60)

    results = run_comprehensive_test()

    print("\n" + "=" * 60)
    print("Test completed successfully!")
    print("=" * 60)