update app.py

app.py

@@ -8,369 +8,191 @@ import numpy as np
 import cv2
 import os
 
-# Workaround for Gradio API schema bug
-# Monkey-patch to handle the schema generation error gracefully
-try:
-    import gradio_client.utils as client_utils
-    original_get_type = client_utils.get_type
-    
-    def patched_get_type(schema):
-        if isinstance(schema, bool):
-            return "bool"
-        return original_get_type(schema)
-    
-    client_utils.get_type = patched_get_type
-except:
-    pass  # If patching fails, continue anyway
-
+# ======================
 # Model configuration
+# ======================
 MODEL_PATH = "robust_galaxy_model.pth"
-NUM_CLASSES = 2
 CLASS_NAMES = ["Elliptical", "Spiral"]
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Image preprocessing
+# ======================
+# Preprocessing
+# ======================
 preprocess = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    transforms.Normalize(
+        mean=[0.485, 0.456, 0.406],
+        std=[0.229, 0.224, 0.225]
+    )
 ])
 
-# Load model
+# ======================
+# Model loading
+# ======================
 def get_model(num_classes=2):
     model = models.resnet18(weights=None)
     model.fc = nn.Linear(model.fc.in_features, num_classes)
     return model
 
 def load_model():
-    model = get_model(NUM_CLASSES)
+    model = get_model()
     if os.path.exists(MODEL_PATH):
-        try:
-            state_dict = torch.load(MODEL_PATH, map_location=DEVICE)
-            model.load_state_dict(state_dict)
-            print(f"Model loaded successfully from {MODEL_PATH}")
-        except Exception as e:
-            print(f"Error loading model: {e}")
-            print("Using untrained model")
+        state_dict = torch.load(MODEL_PATH, map_location="cpu")
+        model.load_state_dict(state_dict)
+        print("✅ Model loaded")
     else:
-        print(f"Model file not found at {MODEL_PATH}. Using untrained model.")
+        print("⚠️ Model not found, using untrained model")
     model.to(DEVICE)
     model.eval()
     return model
 
-# Load model - handle errors gracefully
-model = None
-try:
-    model = load_model()
-    print("Model loaded successfully")
-except Exception as e:
-    print(f"Failed to load model: {e}")
-    import traceback
-    traceback.print_exc()
-    # Create a dummy model as fallback
-    model = get_model(NUM_CLASSES).to(DEVICE)
-    model.eval()
-    print("Using untrained model as fallback")
+model = load_model()
 
-# Grad-CAM implementation
+# ======================
+# Grad-CAM
+# ======================
 class GradCAM:
     def __init__(self, model, target_layer):
         self.model = model
         self.target_layer = target_layer
         self.gradients = None
         self.activations = None
-        self.hook_handles = []
-        
+
     def save_activation(self, module, input, output):
         self.activations = output.detach()
-    
+
     def save_gradient(self, module, grad_input, grad_output):
         self.gradients = grad_output[0].detach()
-    
-    def generate_cam(self, input_image, target_class=None):
-        # Register hooks
-        forward_handle = self.target_layer.register_forward_hook(self.save_activation)
-        backward_handle = self.target_layer.register_full_backward_hook(self.save_gradient)
-        
-        try:
-            # Forward pass
-            model_output = self.model(input_image)
-            
-            if target_class is None:
-                target_class = model_output.argmax(dim=1).item()
-            
-            # Backward pass
-            self.model.zero_grad()
-            class_score = model_output[0, target_class]
-            class_score.backward(retain_graph=False)
-            
-            if self.gradients is None or self.activations is None:
-                return np.zeros((7, 7))  # Default size for ResNet layer4
-            
-            gradients = self.gradients[0]
-            activations = self.activations[0]
-            
-            # Global average pooling of gradients
-            weights = gradients.mean(dim=(1, 2), keepdim=True)
-            cam = (weights * activations).sum(dim=0)
-            
-            # Apply ReLU and normalize
-            cam = F.relu(cam)
-            cam = cam - cam.min()
-            if cam.max() > 0:
-                cam = cam / cam.max()
-            
-            return cam.detach().cpu().numpy()
-        finally:
-            # Remove hooks
-            forward_handle.remove()
-            backward_handle.remove()
-            self.gradients = None
-            self.activations = None
-
-def overlay_heatmap(image, heatmap, alpha=0.4):
-    """Overlay heatmap on original image"""
-    heatmap_resized = cv2.resize(heatmap, (image.shape[1], image.shape[0]))
-    heatmap_colored = cv2.applyColorMap(np.uint8(255 * heatmap_resized), cv2.COLORMAP_JET)
-    output = cv2.addWeighted(image, 1 - alpha, heatmap_colored, alpha, 0)
-    return output
-
-def predict_galaxy(image):
-    """Predict galaxy morphology and generate Grad-CAM"""
+
+    def generate(self, x, class_idx):
+        h1 = self.target_layer.register_forward_hook(self.save_activation)
+        h2 = self.target_layer.register_full_backward_hook(self.save_gradient)
+
+        out = self.model(x)
+        score = out[0, class_idx]
+        self.model.zero_grad()
+        score.backward()
+
+        weights = self.gradients.mean(dim=(2, 3), keepdim=True)
+        cam = (weights * self.activations).sum(dim=1)
+        cam = F.relu(cam)
+        cam = cam - cam.min()
+        cam = cam / cam.max()
+
+        h1.remove()
+        h2.remove()
+
+        return cam[0].cpu().numpy()
+
+def overlay_heatmap(img, cam):
+    cam = cv2.resize(cam, (img.shape[1], img.shape[0]))
+    heatmap = cv2.applyColorMap(np.uint8(255 * cam), cv2.COLORMAP_JET)
+    return cv2.addWeighted(img, 0.6, heatmap, 0.4, 0)
+
+# ======================
+# Prediction function
+# ======================
+def predict(image):
     if image is None:
-        return None, "Please upload an image."
-    
-    if model is None:
-        return None, "Error: Model not loaded. Please check the logs."
-    
-    try:
-        # Ensure model is in eval mode
-        model.eval()
-        
-        # Convert image to PIL if it's not already
-        if isinstance(image, np.ndarray):
-            image = Image.fromarray(image.astype('uint8'))
-        elif not isinstance(image, Image.Image):
-            image = Image.open(image) if hasattr(image, 'read') else Image.fromarray(np.array(image))
-        
-        # Ensure image is RGB
-        if image.mode != 'RGB':
-            image = image.convert('RGB')
-        
-        # Preprocess image
-        img_tensor = preprocess(image).unsqueeze(0).to(DEVICE)
-        img_tensor.requires_grad = True
-        
-        # Get prediction
-        with torch.set_grad_enabled(True):
-            outputs = model(img_tensor)
-            probs = F.softmax(outputs, dim=1)
-            pred_class = outputs.argmax(dim=1).item()
-            confidence = probs[0][pred_class].item()
-            
-            # Generate Grad-CAM
-            try:
-                gradcam = GradCAM(model, model.layer4)
-                cam = gradcam.generate_cam(img_tensor, pred_class)
-            except Exception as cam_error:
-                print(f"Grad-CAM error: {cam_error}")
-                import traceback
-                traceback.print_exc()
-                # If Grad-CAM fails, just return the original image
-                cam = None
-        
-        # Prepare original image for overlay
-        img_np = np.array(image)
-        img_resized = cv2.resize(img_np, (224, 224))
-        
-        # Create overlay if Grad-CAM succeeded
-        if cam is not None:
-            try:
-                overlay = overlay_heatmap(img_resized, cam)
-                overlay_rgb = cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB)
-                overlay_pil = Image.fromarray(overlay_rgb)
-            except Exception as overlay_error:
-                print(f"Overlay error: {overlay_error}")
-                overlay_pil = image.resize((224, 224))
-        else:
-            overlay_pil = image.resize((224, 224))
-        
-        # Format results
-        result_text = f"Predicted Class: {CLASS_NAMES[pred_class]}\nConfidence: {confidence:.2%}"
-        
-        # Ensure we return PIL Image
-        if not isinstance(overlay_pil, Image.Image):
-            overlay_pil = Image.fromarray(np.array(overlay_pil))
-        
-        return overlay_pil, str(result_text)
-    except Exception as e:
-        import traceback
-        error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
-        print(error_msg)  # Print for debugging
-        return None, f"Error: {str(e)}"
-
-# Custom CSS for black background and white text
+        return None, "⚠️ Please upload an image."
+
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+
+    image = image.convert("RGB")
+    x = preprocess(image).unsqueeze(0).to(DEVICE)
+    x.requires_grad_(True)
+
+    outputs = model(x)
+    probs = F.softmax(outputs, dim=1)[0]
+    pred_idx = probs.argmax().item()
+    confidence = probs[pred_idx].item()
+
+    cam = GradCAM(model, model.layer4).generate(x, pred_idx)
+
+    img_np = np.array(image.resize((224, 224)))
+    overlay = overlay_heatmap(img_np, cam)
+    overlay = Image.fromarray(overlay)
+
+    result_md = f"""
+### 🌌 Prediction
+**Class:** `{CLASS_NAMES[pred_idx]}`  
+**Confidence:** `{confidence*100:.2f}%`
+
+**Class Probabilities**
+- Elliptical: `{probs[0]*100:.2f}%`
+- Spiral: `{probs[1]*100:.2f}%`
+"""
+
+    return overlay, result_md
+
+# ======================
+# Clean Dark UI CSS
+# ======================
 custom_css = """
-    .gradio-container {
-        background-color: #000000 !important;
-        color: #ffffff !important;
-    }
-    body {
-        background-color: #000000 !important;
-        color: #ffffff !important;
-    }
-    .gradio-container * {
-        color: #ffffff !important;
-    }
-    h1, h2, h3, h4, p, label, span, div {
-        color: #ffffff !important;
-    }
-    .gr-markdown, .gr-markdown * {
-        color: #ffffff !important;
-    }
-    .gr-button {
-        background-color: #333333 !important;
-        color: #ffffff !important;
-        border: 1px solid #555555 !important;
-    }
-    .gr-button:hover {
-        background-color: #555555 !important;
-    }
-    .gr-textbox, .gr-textbox input, .gr-textbox textarea {
-        background-color: #1a1a1a !important;
-        color: #ffffff !important;
-        border: 1px solid #555555 !important;
-    }
-    .gr-image {
-        background-color: #000000 !important;
-        border: none !important;
-        padding: 0 !important;
-        margin: 0 !important;
-    }
-    .gr-image img {
-        border: none !important;
-        box-shadow: none !important;
-        background-color: #000000 !important;
-    }
-    .gr-image-container, .image-container, .image-wrapper {
-        border: none !important;
-        background-color: #000000 !important;
-        padding: 0 !important;
-        margin: 0 !important;
-    }
-    .gr-image .toolbar, .gr-image .image-controls {
-        display: none !important;
-    }
-    .gr-image label, .gr-image .label-wrap {
-        display: none !important;
-    }
-    .gr-box {
-        border: none !important;
-        background-color: #000000 !important;
-    }
-    .panel, .panel-header {
-        background-color: #000000 !important;
-        border: none !important;
-    }
+body, .gradio-container {
+    background-color: #000000;
+    color: #ffffff;
+}
+
+.gr-image, .gr-image img {
+    background: #000000 !important;
+    border: none !important;
+}
+
+.gr-button {
+    background-color: #222 !important;
+    color: white !important;
+    border-radius: 8px;
+}
+
+.gr-button:hover {
+    background-color: #444 !important;
+}
+
+.gr-markdown {
+    color: white !important;
+}
+
+footer { display: none !important; }
 """
 
-# Create Gradio interface
-# Note: There's a known Gradio bug with API schema generation that causes errors
-# The app will still work for classification, but API endpoints may fail
+# ======================
+# UI
+# ======================
 with gr.Blocks(css=custom_css) as demo:
-    # Landing Section
-    with gr.Column():
-        landing_img = gr.Image(value="landing.jpg", height=500, show_label=False, container=False)
-        landing_text = gr.Markdown("""
-        <div style="text-align: center; padding: 30px; color: white; background-color: #000000; width: 100%; display: flex; flex-direction: column; align-items: center; justify-content: center;">
-            <h1 style="font-size: 96px; font-weight: bold; margin: 0 auto 30px auto; text-align: center; width: 100%;">Galaxy Morphology AI</h1>
-            <p style="font-size: 56px; font-weight: normal; margin: 0 auto; text-align: center; width: 100%;">Classify galaxies with state-of-the-art deep learning</p>
-        </div>
-        """)
-    
-    # Spacing between sections
-    gr.Markdown("<div style='height: 60px;'></div>")
-    
-    # How Astrophysicists Use This Section
-    with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("""
-            # How Astrophysicists Use This
-            
-            Galaxy morphology classification is a fundamental tool in modern astrophysics. 
-            By automatically identifying whether a galaxy is elliptical or spiral, researchers 
-            can analyze large datasets from telescopes like the Hubble Space Telescope and 
-            the James Webb Space Telescope. This classification helps understand galaxy 
-            formation, evolution, and the distribution of matter in the universe.
-            
-            The deep learning model uses convolutional neural networks to identify key 
-            features in galaxy images, such as spiral arms, central bulges, and overall 
-            structure. This automated classification enables astronomers to process millions 
-            of galaxy images efficiently, accelerating discoveries in cosmology and 
-            extragalactic astronomy.
-            """)
-        with gr.Column(scale=1):
-            astro_img = gr.Image(value="astro.jpg", show_label=False, container=False, height=400)
-            gr.Markdown("<p style='text-align: center; color: white; margin-top: 10px;'>Astrophysics Research</p>")
-    
-    # Spacing between sections
-    gr.Markdown("<div style='height: 60px;'></div>")
-    
-    # Classification Section
-    gr.Markdown("# Galaxy Morphology Classification")
-    gr.Markdown("Upload a galaxy image to classify its morphology and visualize the model's attention using Grad-CAM.")
-    
+    gr.Markdown(
+        "# 🌌 Galaxy Morphology Classification\n"
+        "Upload a galaxy image to classify its morphology and visualize model attention."
+    )
+
     with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(label="Upload Galaxy Image")
-            classify_btn = gr.Button("Classify Galaxy")
-        
-        with gr.Column():
-            output_image = gr.Image(label="Grad-CAM Visualization")
-            result_text = gr.Textbox(label="Classification Result")
-    
-    # Register the classification function
-    # Disable API to avoid Gradio schema generation bug
+        input_img = gr.Image(
+            label=None,
+            type="pil",
+            show_label=False,
+            container=False
+        )
+
+        output_img = gr.Image(
+            label=None,
+            show_label=False,
+            container=False
+        )
+
+    result_md = gr.Markdown()
+
+    classify_btn = gr.Button("Classify Galaxy")
+
     classify_btn.click(
-        fn=predict_galaxy,
-        inputs=[input_image],
-        outputs=[output_image, result_text],
-        api_name=False
+        fn=predict,
+        inputs=input_img,
+        outputs=[output_img, result_md]
     )
-    
-    # Spacing between sections
-    gr.Markdown("<div style='height: 60px;'></div>")
-    
-    # Dark Energy Section
-    gr.Markdown("""
-    # Understanding Dark Energy Through Galaxy Morphology
-    
-    Galaxy morphology classification plays a crucial role in understanding dark energy, 
-    one of the most profound mysteries in modern cosmology. Dark energy is the 
-    mysterious force driving the accelerated expansion of the universe, and its nature 
-    remains one of the biggest questions in physics.
-    
-    By classifying large numbers of galaxies and mapping their distribution across 
-    cosmic time, astronomers can trace the expansion history of the universe. 
-    Different galaxy types (elliptical vs spiral) form and evolve differently, and 
-    their relative abundances at different redshifts provide clues about the universe's 
-    evolution. The distribution and clustering of these galaxies help measure the 
-    large-scale structure of the universe, which is directly influenced by dark energy.
-    
-    Automated classification systems like this one enable the analysis of millions of 
-    galaxies from current and future surveys, such as the Vera C. Rubin Observatory's 
-    Legacy Survey of Space and Time (LSST). These massive datasets will provide 
-    unprecedented precision in measuring dark energy's properties and understanding 
-    its role in the fate of the universe.
-    """)
-
-# Launch the demo
-# For Hugging Face Spaces, Gradio will automatically detect and launch the demo
-# The API error is a known Gradio bug - the app will still work for classification
+
+# ======================
+# Launch
+# ======================
 if __name__ == "__main__":
-    try:
-        demo.launch(show_api=False)
-    except Exception as e:
-        # If launch fails, try without API
-        print(f"Launch error (non-critical): {e}")
-        demo.launch()
+    demo.launch()