Update app.py

app.py

@@ -127,11 +127,22 @@ transform = transforms.Compose([
 IMAGENET_CLASSES = {}
 try:
     with open('imagenet_classes.json', 'r') as f:
-        IMAGENET_CLASSES = json.load(f)
-except:
+        data = json.load(f)
+        
+        # Handle both dict and list formats
+        if isinstance(data, dict):
+            IMAGENET_CLASSES = data
+        elif isinstance(data, list):
+            # Convert list to dict with string indices
+            IMAGENET_CLASSES = {str(i): data[i] for i in range(len(data))}
+        else:
+            raise ValueError("Unexpected JSON format")
+        
+        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)}
-    print("⚠️ Using default class indices")
+    IMAGENET_CLASSES = {str(i): f"Class_{i}" for i in range(1000)}
+    print(f"⚠️ Using default class indices: {e}")
 
 
 # ============================================================================
@@ -211,7 +222,12 @@ 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", size="lg")
             
-            
+            gr.Markdown("""
+            ### 💡 Tips:
+            - Works best with clear, centered objects
+            - Supports 1000 ImageNet classes
+            - Try different images!
+            """)
         
         with gr.Column():
             output = gr.Label(num_top_classes=5, label="Top-5 Predictions")
@@ -230,7 +246,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     ---
     **Links:** [GitHub Code](https://github.com/Shwethaamrutha/TSAI-S8) | [Training Details](https://github.com/Shwethaamrutha/TSAI-S8/blob/main/README.md)
     
-    Built with PyTorch • 
+    Built with PyTorch • Trained on AWS p4d.24xlarge • Top 10% from-scratch result
     """)
 
 if __name__ == "__main__":