Update app.py

app.py

@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 """
-HuggingFace Spaces App for ImageNet ResNet50 Classifier
-Trained from scratch to 78%+ Top-1 accuracy
+HuggingFace App for ImageNet ResNet50 Classifier - 77.09% Accuracy
 """
 
 import gradio as gr
@@ -83,48 +82,34 @@ class ResNet50(nn.Module):
 
 
 # ============================================================================
-# MODEL LOADING
+# LOAD MODEL
 # ============================================================================
 
 def load_model():
-    """Load the trained model (CPU-optimized for HuggingFace)"""
     model = ResNet50(num_classes=1000)
-    
     try:
-        # Try to load checkpoint
-        checkpoint_path = "best_model_final.pth"  # Will be uploaded separately
-        checkpoint = torch.load(checkpoint_path, map_location='cpu')
-        
-        # Handle different checkpoint formats
+        checkpoint = torch.load("best_model_final.pth", map_location='cpu')
         if isinstance(checkpoint, dict):
-            if 'model' in checkpoint:
-                state_dict = checkpoint['model']
-            elif 'state_dict' in checkpoint:
-                state_dict = checkpoint['state_dict']
-            else:
-                state_dict = checkpoint
+            state_dict = checkpoint.get('model', checkpoint.get('state_dict', checkpoint))
         else:
             state_dict = checkpoint
         
-        # Remove 'module.' prefix if present (from DataParallel)
         new_state_dict = {}
         for k, v in state_dict.items():
             name = k.replace('module.', '') if k.startswith('module.') else k
             new_state_dict[name] = v
         
         model.load_state_dict(new_state_dict)
-        print(f"✅ Model loaded successfully from {checkpoint_path}")
-        
+        print("✅ Model loaded successfully")
     except Exception as e:
         print(f"⚠️ Could not load checkpoint: {e}")
-        print("Using randomly initialized model for demo purposes")
     
     model.eval()
     return model
 
 
 # ============================================================================
-# IMAGE PREPROCESSING
+# PREPROCESSING
 # ============================================================================
 
 transform = transforms.Compose([
@@ -136,31 +121,17 @@ transform = transforms.Compose([
 
 
 # ============================================================================
-# IMAGENET CLASS LABELS
+# IMAGENET CLASSES
 # ============================================================================
 
-# Top 20 most common ImageNet classes for demo
-IMAGENET_CLASSES = {
-    0: "tench", 1: "goldfish", 2: "great white shark", 3: "tiger shark",
-    4: "hammerhead", 5: "electric ray", 6: "stingray", 7: "cock",
-    8: "hen", 9: "ostrich", 10: "brambling", 11: "goldfinch",
-    12: "house finch", 13: "junco", 14: "indigo bunting", 15: "robin",
-    151: "Chihuahua", 207: "golden retriever", 281: "tabby cat",
-    282: "tiger cat", 283: "Persian cat", 285: "Egyptian cat",
-    291: "lion", 292: "tiger", 293: "jaguar", 294: "leopard",
-    404: "airliner", 407: "container ship", 468: "cab",
-    511: "convertible", 609: "jeep", 627: "limousine",
-    817: "sports car", 751: "racer", 779: "school bus",
-    555: "fire engine", 569: "garbage truck", 717: "pickup",
-    # Add more as needed
-}
-
-# Load full class names if available
+IMAGENET_CLASSES = {}
 try:
     with open('imagenet_classes.json', 'r') as f:
         IMAGENET_CLASSES = json.load(f)
 except:
-    pass  # Use default subset
+    # Fallback - create basic class mapping
+    IMAGENET_CLASSES = {str(i): f"Class {i}" for i in range(1000)}
+    print("⚠️ Using default class indices")
 
 
 # ============================================================================
@@ -168,21 +139,21 @@ except:
 # ============================================================================
 
 def predict(image):
-    """
-    Predict ImageNet class for input image
+    """Predict ImageNet class for input image"""
     
-    Args:
-        image: PIL Image
-        
-    Returns:
-        dict: Top-5 predictions with confidence scores
-    """
     if image is None:
-        return {"error": "Please upload an image"}
+        # Return dummy predictions for error case
+        return {
+            "Error - No Image": 1.0,
+            "Please upload an image": 0.0,
+            "": 0.0,
+            " ": 0.0,
+            "  ": 0.0
+        }
     
     try:
         # Preprocess
-        img_tensor = transform(image).unsqueeze(0)  # Add batch dimension
+        img_tensor = transform(image).unsqueeze(0)
         
         # Inference
         with torch.no_grad():
@@ -192,67 +163,76 @@ def predict(image):
         # Get top 5 predictions
         top5_prob, top5_indices = torch.topk(probabilities, 5)
         
-        # Format results
+        # Format results - dict with string keys and float values
         results = {}
         for i in range(5):
             idx = top5_indices[i].item()
             prob = top5_prob[i].item()
-            class_name = IMAGENET_CLASSES.get(str(idx), f"Class {idx}")
-            results[f"{class_name}"] = float(prob)
+            
+            # CRITICAL: Convert idx to string for JSON lookup
+            class_name = IMAGENET_CLASSES.get(str(idx), f"Class_{idx}")
+            
+            # Ensure float probability
+            results[class_name] = float(prob)
         
         return results
         
     except Exception as e:
-        return {"error": f"Prediction failed: {str(e)}"}
+        # Return error in valid format
+        error_msg = str(e)[:80]
+        return {
+            f"Error: {error_msg}": 0.5,
+            "Please try another image": 0.3,
+            "Check console for details": 0.2,
+            "": 0.0,
+            " ": 0.0
+        }
 
 
 # ============================================================================
 # GRADIO INTERFACE
 # ============================================================================
 
-# Load model globally
 print("Loading model...")
 model = load_model()
-print("Model loaded successfully!")
+print("Model ready!")
 
-# Create Gradio interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
     # 🔥 ImageNet ResNet50 Classifier
     
-    Upload any image and get top-5 predictions with confidence scores.
+    **77.09% Top-1 Accuracy** - Trained from scratch on ImageNet (1.2M images, 1000 classes)
+    
+    Upload an image to get top-5 predictions with confidence scores.
     """)
     
     with gr.Row():
         with gr.Column():
             image_input = gr.Image(type="pil", label="Upload Image")
-            predict_btn = gr.Button("Classify Image", variant="primary")
+            predict_btn = gr.Button("Classify Image", variant="primary", size="lg")
+            
             
         
         with gr.Column():
             output = gr.Label(num_top_classes=5, label="Top-5 Predictions")
             
             gr.Markdown("""
-            ### 🎯 Model Info:
+            ### 📊 Model Info:
             - **Architecture:** ResNet50 (25.5M params)
-            - **Training:** From scratch (no pretrained weights)
-            - **Dataset:** ImageNet (1.2M images, 1000 classes)
-            - **Accuracy:** 77.09% Top-1 validation
-            
-            ### 🔗 Links:
-            - [GitHub Repository](https://github.com/Shwethaamrutha/TSAI-S8)
+            - **Training:** From scratch (no pretrained)
+            - **Accuracy:** 77.09% Top-1
+            - **Hardware:** 8× A100 GPUs
             """)
-
     
-    # Connect button
     predict_btn.click(fn=predict, inputs=image_input, outputs=output)
     
+    gr.Markdown("""
+    ---
+    **Links:** [GitHub Code](https://github.com/Shwethaamrutha/TSAI-S8) | [Training Details](https://github.com/Shwethaamrutha/TSAI-S8/blob/main/README.md)
     
-    
-    
-
+    Built with PyTorch • 
+    """)
 
-# Launch
 if __name__ == "__main__":
     demo.launch()