Fix Resnet

app.py

@@ -25,6 +25,26 @@ app.add_middleware(
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# ------------------
+# Load baseline model (ResNet)
+# ------------------
+resnet_ckpt = torch.load(
+    "best_resnet_multitask.pth",
+    map_location=device
+)
+
+resnet_classes = resnet_ckpt.get("classes") or []
+resnet_num_classes = len(resnet_classes) if resnet_classes else 9
+resnet_mold_idx = (
+    resnet_classes.index("mold")
+    if resnet_classes else 4
+)
+
+resnet_model = MultiTaskResNet50(resnet_num_classes).to(device)
+resnet_model.load_state_dict(resnet_ckpt["model"])
+resnet_model.eval()
+
+
 # ------------------
 # Load main model (ConvNeXt)
 # ------------------
@@ -150,3 +170,26 @@ async def explain_gradcam(file: UploadFile):
     img_t = transform(img).to(device)
     cam = gradcam.generate(img_t, mold_idx)
     return {"gradcam": cam.tolist()}
+
+
+@app.post("/predict/resnet")
+async def predict_resnet(file: UploadFile):
+    img = Image.open(io.BytesIO(await file.read())).convert("RGB")
+    img_t = transform(img).to(device)
+
+    with torch.no_grad():
+        out = resnet_model(img_t.unsqueeze(0))
+        cp = torch.softmax(out["class"], 1)[0]
+        bp = torch.softmax(out["bio"], 1)[0]
+
+    mold_p = cp[resnet_mold_idx].item()
+    bio_p = bp[1].item()
+
+    decision = final_decision(mold_p, bio_p)
+
+    return {
+        "decision": decision,
+        "mold_probability": round(mold_p, 3),
+        "biological_probability": round(bio_p, 3),
+    }
+