Fix ConvNeXt checkpoint loading and Grad-CAM layer selection

app.py

@@ -1,10 +1,10 @@
-from fastapi import FastAPI, UploadFile
+from fastapi import FastAPI, UploadFile, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from PIL import Image
 import torch, io
 from torchvision import transforms
 
-from model import MultiTaskResNet50, MultiTaskConvNeXt
+from model import MultiTaskResNet50, MultiTaskConvNeXt, find_last_conv2d
 from decision import final_decision
 from advanced_decision import (
     mc_uncertainty,
@@ -12,7 +12,7 @@ from advanced_decision import (
     final_decision_v2
 )
 from gradcam import GradCAM
-from dino import load_dino, build_embeddings, similarity
+from typing import Optional
 
 app = FastAPI(title="Mold Detection API v2 (ConvNeXt)")
 
@@ -53,26 +53,41 @@ transform = transforms.Compose([
 ])
 
 # ------------------
-# Grad-CAM (use exposed last_conv from ConvNeXt wrapper)
-# If missing, fall back to a reasonable conv layer
-target_layer = getattr(model, "last_conv", None)
-if target_layer is None:
-    # ConvNeXt features[-1] is a ConvNeXt block with a depthwise conv `dwconv`
-    target_layer = model.backbone.features[-1].dwconv
-gradcam = GradCAM(model, target_layer)
+# Grad-CAM target layer (computed, not stored in model state_dict)
+# ------------------
+target_layer = find_last_conv2d(model.backbone)
+gradcam = GradCAM(model, target_layer) if target_layer is not None else None
 
 # ------------------
 # DINO (lazy loaded)
 # ------------------
-dino = None
+dino: Optional[object] = None
 mold_embs = None
 
 
 def ensure_dino():
     global dino, mold_embs
     if dino is None:
-        dino = load_dino(device)
-        mold_embs = build_embeddings(dino, transform, device)
+        try:
+            from dino import load_dino, build_embeddings
+        except ModuleNotFoundError as e:
+            # Local/dev env might not have optional deps like `datasets`.
+            raise HTTPException(
+                status_code=503,
+                detail=(
+                    "DINO dependencies are not installed. "
+                    "Install extras with: pip install datasets scikit-learn"
+                ),
+            ) from e
+
+        try:
+            dino = load_dino(device)
+            mold_embs = build_embeddings(dino, transform, device)
+        except Exception as e:
+            raise HTTPException(
+                status_code=503,
+                detail=f"Failed to initialize DINO reference embeddings: {e}",
+            ) from e
 
 
 # ------------------
@@ -89,6 +104,8 @@ async def predict_v1(file: UploadFile):
 @app.post("/predict/v2")
 async def predict_v2(file: UploadFile):
     ensure_dino()
+    # Import similarity lazily (only needed for v2)
+    from dino import similarity
 
     img = Image.open(io.BytesIO(await file.read())).convert("RGB")
     img_t = transform(img).to(device)

model.py

@@ -3,6 +3,19 @@ import torch.nn as nn
 from torchvision import models
 
 
+def find_last_conv2d(module: nn.Module) -> nn.Conv2d | None:
+    """
+    Returns the last nn.Conv2d found in a module traversal.
+    Important: we do NOT attach this as a child module on the model instance,
+    otherwise it becomes part of state_dict and breaks checkpoint loading.
+    """
+    last = None
+    for m in module.modules():
+        if isinstance(m, nn.Conv2d):
+            last = m
+    return last
+
+
 class MultiTaskResNet50(nn.Module):
     def __init__(self, num_classes=9):
         super().__init__()
@@ -44,11 +57,6 @@ class MultiTaskConvNeXt(nn.Module):
         self.bio_head = nn.Linear(feat_dim, 2)
         self.dropout = nn.Dropout(p=0.1)
 
-        # Expose a sensible last conv layer ref for Grad-CAM usage.
-        # In torchvision ConvNeXt, each element of `features` is a ConvNeXt block
-        # and has a depthwise conv named `dwconv`.
-        self.last_conv = self.backbone.features[-1].dwconv
-
     def forward(self, x: torch.Tensor):
         feats = self.backbone.features(x)
         feats = self.pool(feats)