Update app.py

app.py

@@ -5,18 +5,22 @@ from PIL import Image
 import os
 import time
 import joblib
+import pickle
+import torch
+import torch.nn as nn
+
+# ===============================
+# Step 1: Load Test Dataset
+# ===============================
 
-# Load labels
 label_df = pd.read_csv("labels.csv")  # columns: filename,label
 test_dir = "test_images"
 
-# Preprocessing: convert image to feature vector
 def image_to_features(image: Image.Image) -> np.ndarray:
-    image = image.resize((64, 64))  # resize to consistent shape
-    array = np.asarray(image.convert("L"))  # grayscale
+    image = image.resize((64, 64))
+    array = np.asarray(image.convert("L"))
     return array.flatten()
 
-# Build test dataset
 X_test = []
 y_test = []
 
@@ -31,25 +35,83 @@ for _, row in label_df.iterrows():
 X_test = np.array(X_test)
 y_test = np.array(y_test)
 
+# ===============================
+# Step 2: Define Leaderboard File
+# ===============================
+
 LEADERBOARD_PATH = "leaderboard.csv"
 if not os.path.exists(LEADERBOARD_PATH):
     pd.DataFrame(columns=["Name", "Accuracy", "Avg Time (ms)"]).to_csv(LEADERBOARD_PATH, index=False)
 
+# ===============================
+# Step 3: Define Model Wrappers
+# ===============================
+
+class SimpleTorchModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc = nn.Sequential(
+            nn.Linear(64 * 64, 128),
+            nn.ReLU(),
+            nn.Linear(128, 2)
+        )
+
+    def forward(self, x):
+        return self.fc(x)
+
+class SKLearnWrapper:
+    def __init__(self, model):
+        self.model = model
+
+    def predict(self, X: np.ndarray) -> np.ndarray:
+        return self.model.predict(X)
+
+class TorchPredictWrapper:
+    def __init__(self, model_path: str):
+        self.model = SimpleTorchModel()
+        self.model.load_state_dict(torch.load(model_path, map_location="cpu"))
+        self.model.eval()
+
+    def predict(self, X: np.ndarray) -> np.ndarray:
+        with torch.no_grad():
+            tensor = torch.tensor(X, dtype=torch.float32)
+            logits = self.model(tensor)
+            return torch.argmax(logits, dim=1).numpy()
+
+# ===============================
+# Step 4: Load Model From File
+# ===============================
+
+def load_model(file_path: str):
+    try:
+        model = joblib.load(file_path)
+        return SKLearnWrapper(model)
+    except:
+        try:
+            with open(file_path, "rb") as f:
+                model = pickle.load(f)
+            return SKLearnWrapper(model)
+        except:
+            try:
+                return TorchPredictWrapper(file_path)
+            except Exception as e:
+                raise RuntimeError(f"Unrecognized model format or failed to load: {e}")
+
+# ===============================
+# Step 5: Evaluation Logic
+# ===============================
+
 def evaluate_model(file, name):
     try:
-        # Load the student's model
-        model = joblib.load(file.name)
+        model = load_model(file.name)
 
-        # Measure time for prediction
         start = time.time()
         y_pred = model.predict(X_test)
-        elapsed = (time.time() - start) * 1000  # ms
+        elapsed = (time.time() - start) * 1000  # in ms
 
-        # Calculate accuracy
         accuracy = 100 * (y_pred == y_test).mean()
         avg_time = elapsed / len(X_test)
 
-        # Update leaderboard
         leaderboard = pd.read_csv(LEADERBOARD_PATH)
         leaderboard = pd.concat([leaderboard, pd.DataFrame([{
             "Name": name,
@@ -63,15 +125,19 @@ def evaluate_model(file, name):
     except Exception as e:
         return f"❌ Error during evaluation: {e}"
 
+# ===============================
+# Step 6: Gradio Interface
+# ===============================
+
 interface = gr.Interface(
     fn=evaluate_model,
     inputs=[
-        gr.File(label="Upload your `.joblib` model"),
+        gr.File(label="Upload your `.joblib`, `.pkl`, or `.pth` model"),
         gr.Text(label="Your (group) name(s)")
     ],
     outputs=gr.Dataframe(label="Leaderboard"),
-    title="Olive Fly Detection Leaderboard",
-    description="Upload a `.joblib` file (scikit-learn model with .predict()). We will test it on flattened images and show your accuracy and time."
+    title="Insect Classifier Leaderboard",
+    description="Upload a `.joblib`, `.pkl`, or `.pth` model that implements `.predict(X)` on 64×64 grayscale image vectors. We will evaluate your model and update the leaderboard."
 )
 
 interface.launch()