app.py

"""
Hugging Face Spaces App for Kolam AI Generator
Enhanced with StyleConditionedGenerator for more variety
"""

import gradio as gr
import torch
import numpy as np
from PIL import Image
from pathlib import Path
import sys
import matplotlib.cm as cm  # for color mapping

# Add project paths
sys.path.insert(0, str(Path(__file__).parent.parent))
sys.path.insert(0, str(Path(__file__).parent.parent / 'models'))

from models.gan_generator import StyleConditionedGenerator
from models.gan_discriminator import KolamDiscriminator
from utils.metrics import KolamDesignMetrics


class KolamAIGenerator:
    def __init__(self):
        """Initialize the Kolam AI Generator."""
        self.generator = None
        self.discriminator = None
        self.metrics = KolamDesignMetrics()
        self.load_models()
    
    def load_models(self):
        """Load the pre-trained models."""
        try:
            # Use StyleConditionedGenerator
            self.generator = StyleConditionedGenerator(
                noise_dim=100,
                feature_dim=128,
                style_dim=32,
                output_channels=1,
                image_size=64
            )
            
            self.discriminator = KolamDiscriminator(
                input_channels=1,
                image_size=64
            )
            
            # Try to load pretrained weights
            weights_path = Path("models/generator.pth")
            if weights_path.exists():
                self.generator.load_state_dict(
                    torch.load(weights_path, map_location="cpu")
                )
                print("✅ Loaded pretrained generator weights!")
            else:
                print("⚠️ No pretrained weights found, using untrained model.")
            
            self.generator.eval()
            self.discriminator.eval()
        
        except Exception as e:
            print(f"❌ Error loading models: {e}")
            self.generator = StyleConditionedGenerator()
            self.generator.eval()
    
    def generate_kolam(self, complexity, symmetry, seed=None, use_color=True):
        """Generate a Kolam design with specified parameters."""
        try:
            # Seed control
            if seed is not None:
                torch.manual_seed(int(seed))
                np.random.seed(int(seed))
            else:
                seed = np.random.randint(0, 100000)
                torch.manual_seed(seed)
                np.random.seed(seed)
            
            # Random noise
            noise = torch.randn(1, 100)
            
            # Complexity tuning
            if complexity == "Simple":
                noise = noise * 0.5
            elif complexity == "Medium":
                noise = noise * 1.0 + torch.randn_like(noise) * 0.3
            elif complexity == "Complex":
                noise = noise * 1.5 + torch.randn_like(noise) * 0.5
            
            # Random features & style vector for variety
            features = torch.randn(1, 128)
            style = torch.randn(1, 32)
            
            # Generate image
            with torch.no_grad():
                generated_kolam = self.generator(noise, features, style)
            
            # Normalize to [0,1]
            kolam_image = generated_kolam.squeeze().cpu().numpy()
            kolam_image = (kolam_image + 1) / 2
            kolam_image = np.clip(kolam_image, 0, 1)
            
            # Apply symmetry
            if symmetry == "High":
                kolam_image = self.enhance_symmetry(kolam_image)
            
            # Convert to color
            if use_color:
                kolam_colored = cm.viridis(kolam_image)[:, :, :3]
                kolam_pil = Image.fromarray((kolam_colored * 255).astype(np.uint8))
            else:
                kolam_pil = Image.fromarray((kolam_image * 255).astype(np.uint8), mode='L')
            
            return kolam_pil
        
        except Exception as e:
            print(f"❌ Error generating Kolam: {e}")
            return self.create_fallback_pattern()
    
    def enhance_symmetry(self, image):
        """Enhance symmetry with mirroring + rotation."""
        img_sym = (image + np.fliplr(image)) / 2
        img_sym = (img_sym + np.flipud(img_sym)) / 2
        rotated = np.rot90(image)
        img_sym = (img_sym + rotated) / 2
        return np.clip(img_sym, 0, 1)
    
    def create_fallback_pattern(self):
        """Fallback geometric pattern."""
        size = 64
        pattern = np.zeros((size, size), dtype=np.float32)
        center = size // 2
        for radius in range(5, center, 8):
            y, x = np.ogrid[:size, :size]
            mask = (x - center) ** 2 + (y - center) ** 2 <= radius ** 2
            pattern[mask] = 1.0
        return Image.fromarray((pattern * 255).astype(np.uint8), mode='L')
    
    def analyze_quality(self, image):
        """Analyze the quality of the Kolam."""
        try:
            if isinstance(image, Image.Image):
                image_array = np.array(image) / 255.0
            else:
                image_array = image
            quality = self.metrics.calculate_overall_quality(image_array)
            return {
                "Overall Quality": f"{quality['overall_quality']:.3f}",
                "Horizontal Symmetry": f"{quality['horizontal']:.3f}",
                "Vertical Symmetry": f"{quality['vertical']:.3f}",
                "Complexity": f"{quality['complexity']:.3f}",
                "Balance": f"{quality['balance']:.3f}",
                "Rhythm": f"{quality['rhythm']:.3f}"
            }
        except Exception as e:
            print(f"❌ Error analyzing quality: {e}")
            return {k: "N/A" for k in [
                "Overall Quality", "Horizontal Symmetry",
                "Vertical Symmetry", "Complexity",
                "Balance", "Rhythm"
            ]}


# -------------------------
# Interface setup
# -------------------------
kolam_ai = KolamAIGenerator()

def generate_and_analyze(complexity, symmetry, seed):
    kolam_image = kolam_ai.generate_kolam(complexity, symmetry, seed, use_color=True)
    quality_metrics = kolam_ai.analyze_quality(kolam_image)
    return kolam_image, quality_metrics


def create_interface():
    css = """
    .gradio-container { font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif; }
    .title { text-align: center; color: #2E86AB; margin-bottom: 20px; }
    .description { text-align: center; color: #666; margin-bottom: 30px; }
    """
    
    with gr.Blocks(css=css, title="Kolam AI Generator") as interface:
        gr.HTML("""
        <div class="title">
            <h1>🎨 Kolam AI Generator</h1>
            <p class="description">Generate beautiful, diverse Kolam designs using AI</p>
        </div>
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 🎛️ Controls")
                complexity = gr.Dropdown(["Simple", "Medium", "Complex"], value="Medium", label="Pattern Complexity")
                symmetry = gr.Dropdown(["Low", "Medium", "High"], value="Medium", label="Symmetry Level")
                seed = gr.Number(value=None, label="Random Seed (Optional)")
                generate_btn = gr.Button("🎨 Generate Kolam", variant="primary", size="lg")
            
            with gr.Column(scale=2):
                gr.Markdown("### 🖼️ Generated Kolam")
                output_image = gr.Image(label="Generated Design", type="pil", height=400)
                gr.Markdown("### 📊 Quality Analysis")
                quality_output = gr.JSON(label="Design Quality Metrics", value={})
        
        generate_btn.click(fn=generate_and_analyze, inputs=[complexity, symmetry, seed], outputs=[output_image, quality_output])
        interface.load(fn=lambda: generate_and_analyze("Medium", "Medium", None), outputs=[output_image, quality_output])
    
    return interface


if __name__ == "__main__":
    interface = create_interface()
    interface.launch(server_name="0.0.0.0", server_port=7860, share=True, show_error=True)