Update app.py

app.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 """
-HuggingFace App for ImageNet ResNet50 Classifier - 77.09% Accuracy
+DEBUG VERSION - HuggingFace App for ImageNet ResNet50
 """
 
 import gradio as gr
@@ -86,14 +86,25 @@ class ResNet50(nn.Module):
 # ============================================================================
 
 def load_model():
+    print("="*70)
+    print("LOADING MODEL")
+    print("="*70)
+    
     model = ResNet50(num_classes=1000)
+    
     try:
-        checkpoint = torch.load("best_model_final.pth", map_location='cpu')
+        checkpoint = torch.load("best_model_final.pth", map_location='cpu', weights_only=False)
+        print(f"Checkpoint type: {type(checkpoint)}")
+        print(f"Checkpoint keys: {list(checkpoint.keys())[:5] if isinstance(checkpoint, dict) else 'Not a dict'}")
+        
         if isinstance(checkpoint, dict):
             state_dict = checkpoint.get('model', checkpoint.get('state_dict', checkpoint))
         else:
             state_dict = checkpoint
         
+        print(f"State dict type: {type(state_dict)}")
+        print(f"State dict keys (first 5): {list(state_dict.keys())[:5]}")
+        
         new_state_dict = {}
         for k, v in state_dict.items():
             name = k.replace('module.', '') if k.startswith('module.') else k
@@ -101,10 +112,21 @@ def load_model():
         
         model.load_state_dict(new_state_dict)
         print("✅ Model loaded successfully")
+        
+        # Test forward pass
+        test_input = torch.randn(1, 3, 224, 224)
+        with torch.no_grad():
+            test_output = model(test_input)
+        print(f"✅ Model output shape: {test_output.shape}")
+        print(f"✅ Model output range: [{test_output.min():.2f}, {test_output.max():.2f}]")
+        
     except Exception as e:
-        print(f"⚠️ Could not load checkpoint: {e}")
+        print(f"❌ Error loading checkpoint: {e}")
+        import traceback
+        traceback.print_exc()
     
     model.eval()
+    print("="*70)
     return model
 
 
@@ -129,16 +151,21 @@ try:
     with open('imagenet_classes.json', 'r') as f:
         data = json.load(f)
         
+        print(f"JSON data type: {type(data)}")
+        
         # Handle both dict and list formats
         if isinstance(data, dict):
             IMAGENET_CLASSES = data
+            print(f"✅ Loaded as dict with {len(IMAGENET_CLASSES)} classes")
         elif isinstance(data, list):
             # Convert list to dict with string indices
             IMAGENET_CLASSES = {str(i): data[i] for i in range(len(data))}
+            print(f"✅ Converted list to dict with {len(IMAGENET_CLASSES)} classes")
         else:
-            raise ValueError("Unexpected JSON format")
+            raise ValueError(f"Unexpected JSON format: {type(data)}")
+        
+        print(f"Sample classes: {list(IMAGENET_CLASSES.items())[:3]}")
         
-        print(f"✅ Loaded {len(IMAGENET_CLASSES)} ImageNet classes")
 except Exception as e:
     # Fallback - create basic class mapping
     IMAGENET_CLASSES = {str(i): f"Class_{i}" for i in range(1000)}
@@ -150,12 +177,11 @@ except Exception as e:
 # ============================================================================
 
 def predict(image):
-    """Predict ImageNet class for input image"""
+    """Predict ImageNet class for input image - WITH DEBUG INFO"""
     
     if image is None:
-        # Return dummy predictions for error case
         return {
-            "Error - No Image": 1.0,
+            "No Image Uploaded": 1.0,
             "Please upload an image": 0.0,
             "": 0.0,
             " ": 0.0,
@@ -163,38 +189,58 @@ def predict(image):
         }
     
     try:
+        print(f"\n{'='*70}")
+        print(f"PREDICTION DEBUG")
+        print(f"{'='*70}")
+        print(f"Image type: {type(image)}")
+        print(f"Image size: {image.size}")
+        print(f"Image mode: {image.mode}")
+        
         # Preprocess
         img_tensor = transform(image).unsqueeze(0)
+        print(f"Tensor shape: {img_tensor.shape}")
+        print(f"Tensor range: [{img_tensor.min():.3f}, {img_tensor.max():.3f}]")
         
         # Inference
         with torch.no_grad():
             outputs = model(img_tensor)
+            print(f"Raw outputs shape: {outputs.shape}")
+            print(f"Raw outputs range: [{outputs.min():.2f}, {outputs.max():.2f}]")
+            
             probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
+            print(f"Probabilities sum: {probabilities.sum():.4f}")
         
         # Get top 5 predictions
         top5_prob, top5_indices = torch.topk(probabilities, 5)
         
-        # Format results - dict with string keys and float values
+        print(f"\nTop-5 Predictions:")
+        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}")
+            print(f"  {idx}: {class_name} = {prob:.4f}")
+        
+        print(f"{'='*70}\n")
+        
+        # Format results
         results = {}
         for i in range(5):
             idx = top5_indices[i].item()
             prob = top5_prob[i].item()
-            
-            # 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 in valid format
-        error_msg = str(e)[:80]
+        print(f"❌ Prediction error: {e}")
+        import traceback
+        traceback.print_exc()
+        
         return {
-            f"Error: {error_msg}": 0.5,
-            "Please try another image": 0.3,
-            "Check console for details": 0.2,
+            f"Error {str(e)[:50]}": 0.5,
+            "Check logs": 0.3,
+            "Try another image": 0.2,
             "": 0.0,
             " ": 0.0
         }
@@ -210,43 +256,27 @@ print("Model ready!")
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
-    # 🔥 ImageNet ResNet50 Classifier
+    # 🔥 ImageNet ResNet50 Classifier (DEBUG VERSION)
     
-    **77.09% Top-1 Accuracy** - Trained from scratch on ImageNet (1.2M images, 1000 classes)
+    **77.09% Top-1 Accuracy** - From scratch training
     
-    Upload an image to get top-5 predictions with confidence scores.
+    Upload an image to test. Check console for debug output.
     """)
     
     with gr.Row():
         with gr.Column():
             image_input = gr.Image(type="pil", label="Upload Image")
-            predict_btn = gr.Button("Classify Image", variant="primary", size="lg")
-            
-            gr.Markdown("""
-            ### 💡 Tips:
-            - Works best with clear, centered objects
-            - Supports 1000 ImageNet classes
-            - Try different images!
-            """)
+            predict_btn = gr.Button("Classify", variant="primary")
         
         with gr.Column():
-            output = gr.Label(num_top_classes=5, label="Top-5 Predictions")
-            
-            gr.Markdown("""
-            ### 📊 Model Info:
-            - **Architecture:** ResNet50 (25.5M params)
-            - **Training:** From scratch (no pretrained)
-            - **Accuracy:** 77.09% Top-1
-            - **Hardware:** 8× A100 GPUs
-            """)
+            output = gr.Label(num_top_classes=5, label="Predictions")
     
     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)
+    **Model:** ResNet50 (25.5M params) | **Accuracy:** 77.09%
     
-    Built with PyTorch • Trained on AWS p4d.24xlarge • Top 10% from-scratch result
+    [GitHub](https://github.com/Shwethaamrutha/TSAI-S8)
     """)
 
 if __name__ == "__main__":