Update app.py

app.py

@@ -197,7 +197,7 @@ def predict(image):
         for i in range(5):
             idx = top5_indices[i].item()
             prob = top5_prob[i].item()
-            class_name = IMAGENET_CLASSES.get(idx, f"Class {idx}")
+            class_name = IMAGENET_CLASSES.get(str(idx), f"Class {idx}")
             results[f"{class_name}"] = float(prob)
         
         return results
@@ -220,8 +220,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
     # 🔥 ImageNet ResNet50 Classifier
     
-    **Trained from scratch to 78%+ Top-1 accuracy on ImageNet!**
-    
     Upload any image and get top-5 predictions with confidence scores.
     """)
     
@@ -230,12 +228,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             image_input = gr.Image(type="pil", label="Upload Image")
             predict_btn = gr.Button("Classify Image", variant="primary")
             
-            gr.Markdown("""
-            ### 📝 Tips:
-            - Works best with **clear, centered objects**
-            - Supports **1000 ImageNet classes** (animals, vehicles, objects, etc.)
-            - Try images from different categories!
-            """)
         
         with gr.Column():
             output = gr.Label(num_top_classes=5, label="Top-5 Predictions")
@@ -246,71 +238,19 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             - **Training:** From scratch (no pretrained weights)
             - **Dataset:** ImageNet (1.2M images, 1000 classes)
             - **Accuracy:** 77.09% Top-1 validation
-            - **Training Time:** ~13 hours on 8× A100 GPUs
             
             ### 🔗 Links:
             - [GitHub Repository](https://github.com/Shwethaamrutha/TSAI-S8)
-            - [Training Logs & Details](https://github.com/Shwethaamrutha/TSAI-S8/blob/main/imagenet-training-final/README.md)
-            - [YouTube Demo](https://youtube.com/YOUR_VIDEO_ID)
             """)
-    
-    # Example images
-    gr.Markdown("### 🖼️ Try These Examples:")
-    gr.Examples(
-        examples=[
-            ["examples/dog.jpg"],
-            ["examples/cat.jpg"],
-            ["examples/car.jpg"],
-            ["examples/bird.jpg"],
-        ],
-        inputs=image_input,
-        outputs=output,
-        fn=predict,
-        cache_examples=False,
-    )
+
     
     # Connect button
     predict_btn.click(fn=predict, inputs=image_input, outputs=output)
     
-    gr.Markdown("""
-    ---
-    ### 📊 Training Details:
-    
-    **Phase 1: Initial Training (90 epochs)**
-    - Optimizer: SGD + Nesterov momentum
-    - LR Schedule: OneCycleLR (0.02 → 0.2 → 0.00001)
-    - Regularization: Label smoothing, weight decay, dropout
-    - Result: 76.75%
-    
-    **Phase 2: Fine-tuning (Multiple LR restarts)**
-    - LR=0.001: 76.88% (oscillated)
-    - LR=0.0005: **77.09%** ✅ (best achieved!)
-    - LR=0.0003: 77.02% (similar ceiling)
-    
-    **Result:** 77.09% represents the natural ceiling for standard
-    from-scratch training. Achieving 78%+ requires advanced augmentation
-    techniques (MixUp, CutMix) beyond standard methods.
-    
-    **Key Techniques:**
-    - Mixed precision training (torch.amp)
-    - Distributed training (8 GPUs, DDP)
-    - Robust image loading (handles corrupted files)
-    - Advanced augmentation (crop, flip, color jitter, erasing)
-    
-    ### 💰 Cost Analysis:
-    - Hardware: AWS p4d.24xlarge (8× A100 40GB)
-    - Duration: ~13 hours
-    - Cost: ~$110 (spot pricing)
     
-    ### 📊 Performance Context:
-    - **Industry Baseline:** 70-75% (we beat by 2-7%)
-    - **Good Training:** 75-77% (top tier!)
-    - **Our Result:** 77.09% (top 10% of from-scratch)
-    - **Research-Level:** 78%+ (requires MixUp/CutMix)
     
-    ---
     
-    **Made with ❤️ by [Your Name](https://github.com/Shwethaamrutha)**
+    **Made with ❤️ by Shwetha(https://github.com/Shwethaamrutha)**
     """)
 
 # Launch