Update app.py

app.py

@@ -1,6 +1,6 @@
 """
 Hugging Face Spaces App for Kolam AI Generator
-Interactive web interface for generating Kolam designs
+Improved version: more beautiful & diverse Kolams
 """
 
 import gradio as gr
@@ -8,10 +8,9 @@ import torch
 import numpy as np
 import matplotlib.pyplot as plt
 from PIL import Image
-import io
-import base64
 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))
@@ -21,6 +20,7 @@ from models.gan_generator import KolamGenerator
 from models.gan_discriminator import KolamDiscriminator
 from utils.metrics import KolamDesignMetrics
 
+
 class KolamAIGenerator:
     def __init__(self):
         """Initialize the Kolam AI Generator."""
@@ -45,100 +45,105 @@ class KolamAIGenerator:
                 image_size=64
             )
             
-            # Set to evaluation mode
+            # 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.")
+            
+            # Set eval mode
             self.generator.eval()
             self.discriminator.eval()
             
-            print("✅ Models loaded successfully!")
-            
         except Exception as e:
             print(f"❌ Error loading models: {e}")
-            # Create dummy models for demo
             self.generator = KolamGenerator(100, 128, 1, 64)
             self.generator.eval()
     
-    def generate_kolam(self, complexity, symmetry, seed=None):
+    def generate_kolam(self, complexity, symmetry, seed=None, use_color=True):
         """Generate a Kolam design with specified parameters."""
         try:
-            # Set random seed for reproducibility
+            # Set random seed
             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)
             
             # Generate random noise
-            noise = torch.randn(1, 10000)
+            noise = torch.randn(1, 100)
             
             # Adjust noise based on complexity
             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
+                noise = noise * 1.5 + torch.randn_like(noise) * 0.5
             
             # Generate Kolam
             with torch.no_grad():
                 generated_kolam = self.generator(noise)
             
-            # Convert to numpy
+            # Convert to numpy [0,1]
             kolam_image = generated_kolam.squeeze().cpu().numpy()
-            kolam_image = (kolam_image + 1) / 2  # Convert to [0,1]
+            kolam_image = (kolam_image + 1) / 2
             kolam_image = np.clip(kolam_image, 0, 1)
             
-            # Apply symmetry if requested
+            # Apply symmetry
             if symmetry == "High":
-                # Enhance symmetry
                 kolam_image = self.enhance_symmetry(kolam_image)
             
             # Convert to PIL Image
-            kolam_pil = Image.fromarray((kolam_image * 255).astype(np.uint8), mode='L')
+            if use_color:
+                kolam_colored = cm.viridis(kolam_image)[:, :, :3]  # RGB colormap
+                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 a simple pattern as fallback
             return self.create_fallback_pattern()
     
     def enhance_symmetry(self, image):
-        """Enhance the symmetry of the generated image."""
-        # Create horizontal symmetry
-        left_half = image[:, :image.shape[1]//2]
-        right_half = np.fliplr(left_half)
-        image[:, image.shape[1]//2:] = right_half
+        """Enhance symmetry with mirror + rotation."""
+        # Horizontal + vertical flip
+        img_sym = (image + np.fliplr(image)) / 2
+        img_sym = (img_sym + np.flipud(img_sym)) / 2
         
-        # Create vertical symmetry
-        top_half = image[:image.shape[0]//2, :]
-        bottom_half = np.flipud(top_half)
-        image[image.shape[0]//2:, :] = bottom_half
+        # Add rotational symmetry (90 degrees)
+        rotated = np.rot90(image)
+        img_sym = (img_sym + rotated) / 2
         
-        return image
+        return np.clip(img_sym, 0, 1)
     
     def create_fallback_pattern(self):
-        """Create a simple fallback pattern."""
-        # Create a simple geometric pattern
+        """Fallback: simple geometric pattern."""
         size = 64
         pattern = np.zeros((size, size), dtype=np.float32)
-        
-        # Draw concentric circles
         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 generated Kolam."""
         try:
-            # Convert PIL to numpy
             if isinstance(image, Image.Image):
                 image_array = np.array(image) / 255.0
             else:
                 image_array = image
-            
-            # Calculate quality metrics
             quality = self.metrics.calculate_overall_quality(image_array)
-            
             return {
                 "Overall Quality": f"{quality['overall_quality']:.3f}",
                 "Horizontal Symmetry": f"{quality['horizontal']:.3f}",
@@ -147,173 +152,60 @@ class KolamAIGenerator:
                 "Balance": f"{quality['balance']:.3f}",
                 "Rhythm": f"{quality['rhythm']:.3f}"
             }
-            
         except Exception as e:
             print(f"❌ Error analyzing quality: {e}")
-            return {
-                "Overall Quality": "N/A",
-                "Horizontal Symmetry": "N/A",
-                "Vertical Symmetry": "N/A",
-                "Complexity": "N/A",
-                "Balance": "N/A",
-                "Rhythm": "N/A"
-            }
+            return {k: "N/A" for k in [
+                "Overall Quality", "Horizontal Symmetry", 
+                "Vertical Symmetry", "Complexity", 
+                "Balance", "Rhythm"
+            ]}
+
 
 # Initialize the generator
 kolam_ai = KolamAIGenerator()
 
 def generate_and_analyze(complexity, symmetry, seed):
-    """Generate Kolam and analyze its quality."""
-    # Generate Kolam
-    kolam_image = kolam_ai.generate_kolam(complexity, symmetry, seed)
-    
-    # Analyze quality
+    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():
-    """Create the Gradio interface."""
-    
-    # Custom CSS for better styling
     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;
-    }
+    .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:
-        
-        # Header
         gr.HTML("""
         <div class="title">
             <h1>🎨 Kolam AI Generator</h1>
-            <p class="description">
-                Generate beautiful traditional Indian Kolam designs using AI
-            </p>
+            <p class="description">Generate beautiful traditional Indian Kolam designs using AI</p>
         </div>
         """)
         
         with gr.Row():
             with gr.Column(scale=1):
-                # Input controls
-                gr.Markdown("### 🎛️ Generation Controls")
-                
-                complexity = gr.Dropdown(
-                    choices=["Simple", "Medium", "Complex"],
-                    value="Medium",
-                    label="Pattern Complexity",
-                    info="Controls the intricacy of the design"
-                )
-                
-                symmetry = gr.Dropdown(
-                    choices=["Low", "Medium", "High"],
-                    value="Medium",
-                    label="Symmetry Level",
-                    info="Controls the geometric balance"
-                )
-                
-                seed = gr.Number(
-                    value=None,
-                    label="Random Seed (Optional)",
-                    info="Set a seed for reproducible results"
-                )
-                
-                generate_btn = gr.Button(
-                    "🎨 Generate Kolam",
-                    variant="primary",
-                    size="lg"
-                )
+                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):
-                # Output
                 gr.Markdown("### 🖼️ Generated Kolam")
-                
-                output_image = gr.Image(
-                    label="Generated Design",
-                    type="pil",
-                    height=400
-                )
-                
-                # Quality metrics
+                output_image = gr.Image(label="Generated Design", type="pil", height=400)
                 gr.Markdown("### 📊 Quality Analysis")
-                
-                quality_output = gr.JSON(
-                    label="Design Quality Metrics",
-                    value={}
-                )
-        
-        # Examples
-        gr.Markdown("### 🌟 Example Generations")
-        
-        examples = gr.Examples(
-            examples=[
-                ["Simple", "High", 42],
-                ["Medium", "Medium", 123],
-                ["Complex", "Low", 456],
-                ["Complex", "High", 789]
-            ],
-            inputs=[complexity, symmetry, seed],
-            label="Try these examples:"
-        )
-        
-        # Information section
-        with gr.Accordion("ℹ️ About Kolam AI Generator", open=False):
-            gr.Markdown("""
-            ## What is Kolam?
-            Kolam is a traditional Indian floor art created using rice flour, chalk, or colored powders. 
-            It features intricate geometric designs and is used in daily rituals and festivals.
-            
-            ## How it Works
-            - **Generator Network**: Creates Kolam designs from random noise
-            - **Discriminator Network**: Ensures realistic pattern generation
-            - **Quality Metrics**: Analyzes symmetry, complexity, and balance
-            
-            ## Technical Details
-            - **Framework**: PyTorch
-            - **Architecture**: Generative Adversarial Network (GAN)
-            - **Input**: Random noise (100 dimensions)
-            - **Output**: 64x64 grayscale Kolam image
-            - **Generation Time**: <1 second
-            
-            ## Quality Metrics
-            - **Symmetry**: Horizontal and vertical balance
-            - **Complexity**: Pattern intricacy and density
-            - **Balance**: Visual weight distribution
-            - **Rhythm**: Pattern repetition and flow
-            """)
+                quality_output = gr.JSON(label="Design Quality Metrics", value={})
         
-        # Event handlers
-        generate_btn.click(
-            fn=generate_and_analyze,
-            inputs=[complexity, symmetry, seed],
-            outputs=[output_image, quality_output]
-        )
+        generate_btn.click(fn=generate_and_analyze, inputs=[complexity, symmetry, seed], outputs=[output_image, quality_output])
         
-        # Auto-generate on load
-        interface.load(
-            fn=lambda: generate_and_analyze("Medium", "Medium", None),
-            outputs=[output_image, quality_output]
-        )
+        interface.load(fn=lambda: generate_and_analyze("Medium", "Medium", None), outputs=[output_image, quality_output])
     
     return interface
 
-# Create and launch the interface
+
 if __name__ == "__main__":
     interface = create_interface()
-    interface.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=True,
-        show_error=True
-    )
+    interface.launch(server_name="0.0.0.0", server_port=7860, share=True, show_error=True)