app.py

import numpy as np
import cv2
from PIL import Image
import gradio as gr
import time
import os
import glob
from skimage.metrics import structural_similarity as ssim

class FaceImageMosaicGenerator:
    def __init__(self, faces_directory="Real_Images/"):
        self.faces_directory = faces_directory
        self.tile_size = 32
        self.face_tiles = {}
        self.load_status = self.load_face_tiles()
        
    def load_face_tiles(self):
        """Load face images from the Real_Images directory"""
        try:
            if not os.path.exists(self.faces_directory):
                return f"❌ Faces directory '{self.faces_directory}' not found!"
            
            # Load all image files from faces directory
            face_extensions = ['*.png', '*.jpg', '*.jpeg', '*.bmp', '*.tiff']
            face_files = []
            
            for extension in face_extensions:
                face_files.extend(glob.glob(os.path.join(self.faces_directory, extension)))
            
            if len(face_files) == 0:
                return f"❌ No face images found in {self.faces_directory}"
            
            print(f"Loading {len(face_files)} face images from {self.faces_directory}")
            
            for i, face_path in enumerate(face_files):
                try:
                    # Load face image
                    face_image = Image.open(face_path)
                    
                    # Convert to RGB if needed
                    if face_image.mode != 'RGB':
                        face_image = face_image.convert('RGB')
                    
                    # Resize to standard tile size while maintaining aspect ratio
                    face_image = self.resize_face_tile(face_image, self.tile_size)
                    
                    # Convert to numpy array
                    face_array = np.array(face_image)
                    
                    # Calculate average color for this face
                    avg_color = np.mean(face_array, axis=(0, 1))
                    
                    # Calculate brightness for sorting
                    brightness = np.mean(avg_color)
                    
                    # Extract filename for identification
                    face_name = os.path.splitext(os.path.basename(face_path))[0]
                    
                    # Store face tile with metadata
                    self.face_tiles[face_name] = {
                        'image': face_array,
                        'avg_color': avg_color,
                        'brightness': brightness,
                        'path': face_path,
                        'dominant_color': self.get_dominant_color(face_array)
                    }
                    
                    if i % 100 == 0 and i > 0:  # Progress indicator
                        print(f"Processed {i}/{len(face_files)} face images...")
                        
                except Exception as e:
                    print(f"Error loading face {face_path}: {e}")
                    continue
            
            if not self.face_tiles:
                return f"❌ No valid face images loaded from {self.faces_directory}"
            
            print(f"✅ Successfully loaded {len(self.face_tiles)} face tiles")
            return f"✅ Loaded {len(self.face_tiles)} face images"
            
        except Exception as e:
            return f"❌ Error loading faces: {str(e)}"
    
    def resize_face_tile(self, face_image, target_size):
        """Resize face image to tile size while maintaining aspect ratio"""
        # Calculate aspect ratio
        width, height = face_image.size
        aspect_ratio = width / height
        
        if aspect_ratio > 1:  # Wide image
            new_width = target_size
            new_height = int(target_size / aspect_ratio)
        else:  # Tall or square image
            new_height = target_size
            new_width = int(target_size * aspect_ratio)
        
        # Resize image
        resized = face_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
        
        # Create square canvas and paste resized image in center
        square_image = Image.new('RGB', (target_size, target_size), (128, 128, 128))
        paste_x = (target_size - new_width) // 2
        paste_y = (target_size - new_height) // 2
        square_image.paste(resized, (paste_x, paste_y))
        
        return square_image
    
    def get_dominant_color(self, image_array):
        """Get the dominant color of an image using uniform quantization"""
        # Apply uniform quantization to reduce color space
        quantized = self.quantize_colors(image_array, n_levels=8)
        
        # Flatten to get list of colors
        colors = quantized.reshape(-1, 3)
        
        # Find most frequent color
        unique_colors, counts = np.unique(colors, axis=0, return_counts=True)
        most_frequent_idx = np.argmax(counts)
        
        return unique_colors[most_frequent_idx].astype(float)
    
    def find_best_matching_face(self, target_color, matching_method='color_distance'):
        """Find the face tile that best matches the target color"""
        if not self.face_tiles:
            return None, "no_faces"
            
        best_face_name = None
        best_score = float('inf') if matching_method == 'color_distance' else float('-inf')
        
        for face_name, face_data in self.face_tiles.items():
            if matching_method == 'color_distance':
                # Euclidean distance in RGB space
                face_avg_color = face_data['avg_color']
                distance = np.sqrt(np.sum((target_color - face_avg_color) ** 2))
                score = distance
                is_better = score < best_score
                
            elif matching_method == 'brightness':
                # Match based on brightness similarity
                target_brightness = np.mean(target_color)
                face_brightness = face_data['brightness']
                distance = abs(target_brightness - face_brightness)
                score = distance
                is_better = score < best_score
                
            elif matching_method == 'dominant_color':
                # Match based on dominant color
                face_dominant = face_data['dominant_color']
                distance = np.sqrt(np.sum((target_color - face_dominant) ** 2))
                score = distance
                is_better = score < best_score
            
            if is_better:
                best_score = score
                best_face_name = face_name
        
        if best_face_name:
            return self.face_tiles[best_face_name]['image'], best_face_name
        else:
            return None, "no_match"
    
    def quantize_colors(self, image, n_levels=4):
        """Apply uniform color quantization to reduce color variations"""
        # Calculate quantization step size
        step_size = 256 // n_levels
        
        # Apply uniform quantization to each channel
        quantized_image = (image // step_size) * step_size
        
        # Add half step to center the quantized values
        quantized_image = quantized_image + step_size // 2
        
        # Ensure values stay within [0, 255] range
        quantized_image = np.clip(quantized_image, 0, 255).astype(np.uint8)
        
        return quantized_image
    
    def preprocess_image(self, image, target_size=(512, 512), quantize_colors=False, quantization_levels=4):
        """Preprocess the input image"""
        # Convert PIL to numpy array
        if isinstance(image, Image.Image):
            image = np.array(image)
        
        # Resize image
        image = cv2.resize(image, target_size)
        
        # Optional uniform color quantization
        if quantize_colors:
            image = self.quantize_colors(image, n_levels=quantization_levels)
        
        return image
    
    def create_grid_vectorized(self, image, grid_size):
        """Vectorized grid division and analysis"""
        h, w = image.shape[:2]
        cell_h, cell_w = h // grid_size, w // grid_size
        
        # Crop image to fit exact grid
        cropped_image = image[:cell_h * grid_size, :cell_w * grid_size]
        
        # Reshape image into grid cells using vectorized operations
        grid_cells = cropped_image.reshape(
            grid_size, cell_h, grid_size, cell_w, 3
        ).transpose(0, 2, 1, 3, 4)
        
        # Calculate average color for each cell
        avg_colors = np.mean(grid_cells, axis=(2, 3))
        
        return grid_cells, avg_colors, (cell_h, cell_w)
    
    def create_face_mosaic(self, image, grid_size=32, quantize=False, matching_method='color_distance', quantization_levels=4):
        """Create mosaic from image using FACE IMAGE TILES"""
        if not self.face_tiles:
            return None, None, 0, {"error": "No face tiles loaded"}
            
        start_time = time.time()
        
        try:
            # Preprocess image
            processed_image = self.preprocess_image(image, quantize_colors=quantize, 
                                                   quantization_levels=quantization_levels)
            
            # Create grid
            grid_cells, avg_colors, (cell_h, cell_w) = self.create_grid_vectorized(processed_image, grid_size)
            
            # Create mosaic using FACE IMAGE TILES
            mosaic_height = grid_size * self.tile_size
            mosaic_width = grid_size * self.tile_size
            mosaic = np.zeros((mosaic_height, mosaic_width, 3), dtype=np.uint8)
            
            # Classification visualization
            classified_image = np.zeros_like(processed_image[:grid_size*cell_h, :grid_size*cell_w])
            
            # Keep track of which faces were used
            used_faces = {}
            
            print(f"Creating {grid_size}x{grid_size} mosaic using face tiles...")
            
            for i in range(grid_size):
                for j in range(grid_size):
                    # Get target color for this cell
                    target_color = avg_colors[i, j]
                    
                    # Find the best matching FACE IMAGE TILE
                    best_face_tile, face_name = self.find_best_matching_face(
                        target_color, matching_method
                    )
                    
                    if best_face_tile is not None:
                        # Track face usage
                        if face_name in used_faces:
                            used_faces[face_name] += 1
                        else:
                            used_faces[face_name] = 1
                        
                        # Place the FACE IMAGE TILE in mosaic
                        y_start, y_end = i * self.tile_size, (i + 1) * self.tile_size
                        x_start, x_end = j * self.tile_size, (j + 1) * self.tile_size
                        mosaic[y_start:y_end, x_start:x_end] = best_face_tile
                    
                    # Fill classified image (for debugging)
                    cy_start, cy_end = i * cell_h, (i + 1) * cell_h
                    cx_start, cx_end = j * cell_w, (j + 1) * cell_w
                    classified_image[cy_start:cy_end, cx_start:cx_end] = target_color
                
                # Progress indicator
                if (i + 1) % 8 == 0 or i == grid_size - 1:
                    print(f"Completed {i + 1}/{grid_size} rows")
            
            processing_time = time.time() - start_time
            
            # Create face usage analysis
            face_analysis = {
                'used_faces': used_faces,
                'unique_faces': len(used_faces),
                'total_positions': grid_size * grid_size,
            }
            
            return mosaic, classified_image, processing_time, face_analysis
            
        except Exception as e:
            print(f"Error creating mosaic: {e}")
            return None, None, 0, {"error": str(e)}

def process_face_image(image, grid_size, use_quantization, matching_method, quantization_levels):
    """Main processing function for Gradio interface"""
    if image is None:
        return None, None, "❌ Please upload an image first."
    
    try:
        # Initialize generator (this will load faces once)
        if not hasattr(process_face_image, 'generator'):
            process_face_image.generator = FaceImageMosaicGenerator()
        
        generator = process_face_image.generator
        
        # Check if faces loaded successfully
        if not generator.face_tiles:
            return None, None, f"❌ {generator.load_status}\n\nPlease ensure:\n- Real_Images/ folder exists\n- It contains valid image files\n- Images are readable"
        
        # Create face mosaic
        mosaic, classified, processing_time, face_analysis = generator.create_face_mosaic(
            image, grid_size, use_quantization, matching_method, quantization_levels
        )
        
        if mosaic is None:
            error_msg = face_analysis.get('error', 'Unknown error occurred')
            return None, None, f"❌ Error creating mosaic: {error_msg}"
        
        # Calculate similarity metrics
        try:
            mse, ssim_score = calculate_similarity_metrics(np.array(image), mosaic)
        except Exception as e:
            print(f"Error calculating similarity metrics: {e}")
            mse, ssim_score = 0, 0
        
        # Create results text with safe division
        unique_faces = face_analysis.get('unique_faces', 0)
        total_positions = face_analysis.get('total_positions', 0)
        used_faces = face_analysis.get('used_faces', {})
        
        # Safe division for average reuse calculation
        if unique_faces > 0:
            average_reuse = total_positions / unique_faces
            utilization_percentage = 100 * unique_faces / len(generator.face_tiles)
        else:
            average_reuse = 0
            utilization_percentage = 0
        
        results_text = f"""
🎭 FACE MOSAIC RESULTS

Processing Time: {processing_time:.3f} seconds
Grid Size: {grid_size}×{grid_size} = {grid_size*grid_size} tiles
Color Quantization: {'Enabled' if use_quantization else 'Disabled'}
Quantization Levels: {quantization_levels if use_quantization else 'N/A'}
Matching Method: {matching_method.replace('_', ' ').title()}

📊 SIMILARITY METRICS
MSE: {mse:.2f}
SSIM: {ssim_score:.4f}

👥 FACE DATASET INFO
Total Face Images: {len(generator.face_tiles)}
Directory: {generator.faces_directory}

🎨 MOSAIC STATISTICS
Unique Faces Used: {unique_faces} / {len(generator.face_tiles)} ({utilization_percentage:.1f}%)
Total Tile Positions: {total_positions}
Average Reuse: {average_reuse:.1f} times per face

Most Used Faces:"""
        
        # Add most frequently used faces with safe checking
        if used_faces:
            sorted_faces = sorted(used_faces.items(), 
                                key=lambda x: x[1], reverse=True)
            
            for face_name, count in sorted_faces[:5]:  # Top 5 most used faces
                if total_positions > 0:
                    percentage = 100 * count / total_positions
                    results_text += f"\n- {face_name}: {count} times ({percentage:.1f}%)"
                else:
                    results_text += f"\n- {face_name}: {count} times"
        else:
            results_text += "\n- No faces were successfully matched"
        
        # Additional debug info for matching issues
        if unique_faces == 0:
            results_text += f"\n\n⚠️ DEBUG INFO:"
            results_text += f"\nMatching method: {matching_method}"
            results_text += f"\nTotal face tiles loaded: {len(generator.face_tiles)}"
            results_text += f"\nSample face data available: {'Yes' if generator.face_tiles else 'No'}"
        
        return (
            Image.fromarray(mosaic),
            Image.fromarray(classified),
            results_text
        )
        
    except Exception as e:
        import traceback
        error_msg = f"❌ ERROR: {str(e)}\n\nDebug info:\n{traceback.format_exc()}\n\nPlease check:\n- Real_Images/ directory exists\n- Directory contains valid image files (.jpg, .png, etc.)\n- Images are readable"
        return None, None, error_msg
    
def calculate_similarity_metrics(original, mosaic):
    """Calculate similarity metrics between original and mosaic"""
    try:
        # Resize original to match mosaic size for fair comparison
        original_resized = cv2.resize(original, (mosaic.shape[1], mosaic.shape[0]))
        
        # Convert to grayscale for SSIM
        orig_gray = cv2.cvtColor(original_resized, cv2.COLOR_RGB2GRAY)
        mosaic_gray = cv2.cvtColor(mosaic, cv2.COLOR_RGB2GRAY)
        
        # Calculate MSE
        mse = np.mean((original_resized.astype(float) - mosaic.astype(float)) ** 2)
        
        # Calculate SSIM
        ssim_score = ssim(orig_gray, mosaic_gray)
        
        return mse, ssim_score
    except:
        return 0, 0

# Create Gradio interface
def create_interface():
    with gr.Blocks(title="Face Image Mosaic Generator", theme=gr.themes.Soft()) as interface:
        gr.Markdown("# 🎭 Face Image Mosaic Generator")
        gr.Markdown("""
        Upload an image and convert it into an artistic mosaic using **real human face images**!
        
        Each grid cell in your input image will be replaced with the face image that best matches its color characteristics.
        
        ⚠️ **Important**: Make sure you have uploaded face images to the `Real_Images/` folder in this Space.
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(
                    type="pil", 
                    label="📷 Upload Image to Convert",
                    height=300
                )
                
                with gr.Accordion("⚙️ Mosaic Settings", open=True):
                    grid_size = gr.Slider(
                        minimum=8, maximum=64, value=16, step=8,
                        label="Grid Size (tiles per row/column)",
                        info="Higher = more detail, slower processing"
                    )
                    
                    matching_method = gr.Dropdown(
                        choices=[
                            ("Color Distance", "color_distance")
                        ],
                        value="color_distance",
                        label="Face Matching Method",
                        info="How to choose the best face for each tile"
                    )
                    
                    use_quantization = gr.Checkbox(
                        label="Uniform Color Quantization", 
                        value=False,
                        info="Reduce colors for cleaner matching"
                    )
                    
                    quantization_levels = gr.Slider(
                        minimum=2, maximum=16, value=4, step=1,
                        label="Quantization Levels",
                        info="Number of levels per RGB channel (only if quantization enabled)",
                        visible=False
                    )
                
                # Show/hide quantization levels based on checkbox
                use_quantization.change(
                    fn=lambda x: gr.update(visible=x),
                    inputs=[use_quantization],
                    outputs=[quantization_levels]
                )
                
                process_btn = gr.Button(
                    "🎨 Generate Face Mosaic", 
                    variant="primary", 
                    size="lg"
                )
            
            with gr.Column(scale=2):
                with gr.Tab("🎭 Mosaic Result"):
                    mosaic_output = gr.Image(
                        label="Face Mosaic Result",
                        height=400
                    )
                
                with gr.Tab("🎯 Color Analysis"):
                    classified_output = gr.Image(
                        label="Color Classification Grid",
                        height=400
                    )
                
                with gr.Tab("📊 Statistics"):
                    results_text = gr.Textbox(
                        label="Detailed Analysis",
                        lines=20,
                        max_lines=30
                    )
        
        # Process button click
        process_btn.click(
            fn=process_face_image,
            inputs=[image_input, grid_size, use_quantization, matching_method, quantization_levels],
            outputs=[mosaic_output, classified_output, results_text]
        )
        
        # Instructions
        gr.Markdown("""
        ### 📋 Instructions:
        1. **Upload face images** to the `Real_Images/` folder in this Space's Files tab
        2. **Upload an image** you want to convert into a mosaic
        3. **Adjust settings** (grid size, matching method, etc.)
        4. **Click "Generate Face Mosaic"**
        
        ### 🎯 Tips:
        - Start with smaller grid sizes (8×8, 16×16) for faster processing
        - Use "Color Distance" matching for most images
        - Enable quantization for noisy or complex images
        - More face images = better variety in the mosaic
        """)
    
    return interface

if __name__ == "__main__":
    # Create and launch interface
    interface = create_interface()
    interface.launch()