Update app.py

app.py

@@ -4,12 +4,13 @@ import timm
 import torchvision.transforms as T
 from PIL import Image
 import numpy as np
+import pandas as pd
 import json, os
 
 # ------------------------------
 # CONFIGURATION
 # ------------------------------
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DEVICE = "cpu"
 NUM_CLASSES = 200
 
 MODEL_PATHS = {
@@ -23,13 +24,19 @@ MODEL_NAMES = {
 }
 
 # ------------------------------
-# LOAD CLASS NAMES (JSON or TXT)
+# LOAD CLASS NAMES (supports JSON or TXT)
 # ------------------------------
 CLASS_NAMES = []
+
 if os.path.exists("class_names.json"):
     try:
         with open("class_names.json", "r") as f:
-            CLASS_NAMES = json.load(f)
+            data = json.load(f)
+            if isinstance(data, dict):
+                # some JSONs store indices as keys
+                CLASS_NAMES = [v for _, v in sorted(data.items(), key=lambda x: int(x[0]))]
+            elif isinstance(data, list):
+                CLASS_NAMES = data
         print(f"✅ Loaded {len(CLASS_NAMES)} class names from class_names.json")
     except Exception as e:
         print(f"⚠️ Error reading class_names.json: {e}")
@@ -66,17 +73,17 @@ def get_model(model_name):
 # ------------------------------
 def predict(image, model_choice):
     if image is None:
-        return None, "Please upload an image."
+        return pd.DataFrame(), "Please upload an image."
 
     model = get_model(model_choice)
 
-    # Ensure image is valid RGB
+    # Ensure RGB
     try:
         image = image.convert("RGB")
     except Exception as e:
-        return None, f"❌ Error reading image: {e}"
+        return pd.DataFrame(), f"❌ Error reading image: {e}"
 
-    # Pick reasonable size per model
+    # Resize dynamically
     size = 300 if "b3" in model_choice.lower() else 384
     preprocess = T.Compose([
         T.Resize((size, size)),
@@ -86,7 +93,6 @@ def predict(image, model_choice):
     ])
     img_tensor = preprocess(image).unsqueeze(0).to(DEVICE)
 
-    # Inference
     with torch.no_grad():
         logits = model(img_tensor)
         probs = torch.nn.functional.softmax(logits, dim=1)[0]
@@ -95,19 +101,19 @@ def predict(image, model_choice):
     indices = top5.indices.cpu().numpy().tolist()
     scores = top5.values.cpu().numpy().tolist()
 
-    # Map indices to class names
-    if CLASS_NAMES and len(CLASS_NAMES) >= max(indices) + 1:
+    # Map to readable names
+    if CLASS_NAMES and len(CLASS_NAMES) > max(indices):
         names = [CLASS_NAMES[i] for i in indices]
     else:
         names = [f"Class {i}" for i in indices]
 
-    # Result table
-    result_table = {
+    # Make a DataFrame (this avoids 'undefined' issue)
+    df = pd.DataFrame({
         "Class": names,
         "Confidence (%)": [round(s * 100, 2) for s in scores]
-    }
+    })
 
-    return result_table, f"✅ Prediction complete using {model_choice}"
+    return df, f"✅ Prediction complete using {model_choice}"
 
 # ------------------------------
 # GRADIO INTERFACE
@@ -117,7 +123,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
         """
         # 🐦 Bird Species Classifier  
         Upload a bird image or use your camera to identify its species.  
-        Select a model and view the top-5 predictions with confidence scores.
+        Supports multiple EfficientNet models and displays top-5 predictions.
         """
     )
 
@@ -133,26 +139,22 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
             btn = gr.Button("🔍 Classify")
 
         with gr.Column(scale=1):
-            df_out = gr.Dataframe(headers=["Class", "Confidence (%)"],
-                                  label="Top-5 Predictions")
+            df_out = gr.Dataframe(label="Top-5 Predictions",
+                                  interactive=False)
             status = gr.Textbox(label="Status", interactive=False)
 
-    btn.click(fn=predict,
-              inputs=[img_in, model_sel],
+    btn.click(fn=predict, inputs=[img_in, model_sel],
               outputs=[df_out, status])
 
     gr.Markdown(
         """
         ---
-        ### 💡 Tips
-        - Use **EfficientNet-B3** for faster inference.  
-        - Use **EfficientNetV2-S (in21k)** for higher accuracy.  
-        - Images of any size or orientation are automatically resized.  
+        ### 💡 Notes
+        - Works with **any image size or format** (auto-resized).  
+        - **EfficientNet-B3:** fast inference  
+        - **EfficientNetV2-S:** higher accuracy  
         """
     )
 
-# ------------------------------
-# LAUNCH APP
-# ------------------------------
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)