Add ResNet baseline and ConvNeXt v2 backend Dockerfile 1

app.py

@@ -25,6 +25,9 @@ app.add_middleware(
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# ------------------
+# Load baseline model (ResNet)
+# ------------------
 # ------------------
 # Load baseline model (ResNet)
 # ------------------
@@ -33,15 +36,23 @@ resnet_ckpt = torch.load(
     map_location=device
 )
 
-resnet_classes = resnet_ckpt.get("classes") or []
+# Handle different checkpoint formats
+if isinstance(resnet_ckpt, dict) and "model" in resnet_ckpt:
+    resnet_state = resnet_ckpt["model"]
+    resnet_classes = resnet_ckpt.get("classes", [])
+else:
+    resnet_state = resnet_ckpt
+    resnet_classes = []
+
 resnet_num_classes = len(resnet_classes) if resnet_classes else 9
 resnet_mold_idx = (
     resnet_classes.index("mold")
-    if resnet_classes else 4
+    if resnet_classes and "mold" in resnet_classes
+    else 4
 )
 
 resnet_model = MultiTaskResNet50(resnet_num_classes).to(device)
-resnet_model.load_state_dict(resnet_ckpt["model"])
+resnet_model.load_state_dict(resnet_state)
 resnet_model.eval()