fix: app.py to accept model.py

app.py

@@ -6,7 +6,6 @@ import pandas as pd
 import importlib.util
 import spaces
 
-
 # 1. Load leaderboard data
 dataset = load_dataset("gydou/5190_Spring_Final_Hidden_Data_Released")["test"]
 images = dataset["image"]
@@ -20,17 +19,25 @@ def get_student_transform(preprocess_file):
     spec.loader.exec_module(preprocess_mod)
     return preprocess_mod.get_transform()
 
-# 3. Model testing/evaluation function
+# 3. Function to dynamically import student's model class
+def get_student_model_class(model_class_file):
+    spec = importlib.util.spec_from_file_location("student_model", model_class_file.name)
+    model_mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(model_mod)
+    # Expect the class to be named "StudentModel"
+    return model_mod.StudentModel
+
+# 4. Model testing/evaluation function
 @spaces.GPU
-def test_model(model_file, preprocess_file):
+def test_model(state_dict_file, model_class_file, preprocess_file):
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    # Load full student model (assumes torch.save(model, ...))
-    model = torch.load(model_file.name, map_location=device)
+    # Dynamically import model class and preprocessing
+    StudentModel = get_student_model_class(model_class_file)
+    model = StudentModel().to(device)
+    model.load_state_dict(torch.load(state_dict_file.name, map_location=device))
     model.eval()
-    model.to(device)
 
-    # Dynamically get the transform
     transform = get_student_transform(preprocess_file)
 
     preds = []
@@ -48,8 +55,8 @@ def test_model(model_file, preprocess_file):
 
 leaderboard = pd.DataFrame(columns=["Name", "Score"])
 
-def submit(model_file, preprocess_file, name):
-    score = test_model(model_file, preprocess_file)
+def submit(state_dict_file, model_class_file, preprocess_file, name):
+    score = test_model(state_dict_file, model_class_file, preprocess_file)
     global leaderboard
     leaderboard.loc[len(leaderboard)] = [name, score]
     leaderboard_sorted = leaderboard.sort_values("Score")
@@ -60,20 +67,21 @@ with gr.Blocks() as demo:
     gr.Markdown(
         """
         # IMG2GPS Leaderboard
-        Upload your trained PyTorch model (`student_model.pt`)
-        and the preprocessing module (`preprocess.py`) with a `get_transform()` function.
+        Upload your model weights (`model_state_dict.pt`), the model class definition (`model.py`)
+        (with class named StudentModel), and preprocessing (`preprocess.py` with get_transform).
         """
     )
     with gr.Row():
         name = gr.Text(label="Name/Alias")
-        model_file = gr.File(label="PyTorch Model (.pt)")
-        preprocess_file = gr.File(label="Preprocessing (.py)")
+        state_dict_file = gr.File(label="PyTorch Model State Dict (.pt)")
+        model_class_file = gr.File(label="Model Class (`model.py`)")
+        preprocess_file = gr.File(label="Preprocessing (`preprocess.py`)")
     outputs = [
         gr.Number(label="Mean Coordinate Error"),
         gr.Dataframe(headers=["Name", "Score"], label="Leaderboard")
     ]
     submit_btn = gr.Button("Submit")
-    submit_btn.click(fn=submit, inputs=[model_file, preprocess_file, name], outputs=outputs)
+    submit_btn.click(fn=submit, inputs=[state_dict_file, model_class_file, preprocess_file, name], outputs=outputs)
 
 if __name__ == "__main__":
     demo.launch()