Use ConvNeXt multitask model and new checkpoint

Dockerfile

@@ -15,7 +15,7 @@ RUN pip install --no-cache-dir --upgrade pip && \
     pip install --no-cache-dir -r requirements.txt
 
 COPY *.py ./
-COPY resnet50_multitask_mold.pth ./
+COPY best_convnext_multitask.pth ./
 
 EXPOSE 7860
 

app.py

@@ -4,7 +4,7 @@ from PIL import Image
 import torch, io
 from torchvision import transforms
 
-from model import MultiTaskResNet50
+from model import MultiTaskResNet50, MultiTaskConvNeXt
 from decision import final_decision
 from advanced_decision import (
     mc_uncertainty,
@@ -14,7 +14,7 @@ from advanced_decision import (
 from gradcam import GradCAM
 from dino import load_dino, build_embeddings, similarity
 
-app = FastAPI(title="Mold Detection API v2")
+app = FastAPI(title="Mold Detection API v2 (ConvNeXt)")
 
 app.add_middleware(
     CORSMiddleware,
@@ -24,15 +24,20 @@ app.add_middleware(
 )
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-mold_idx = 4
 
 # ------------------
-# Load main model
+# Load main model (ConvNeXt)
 # ------------------
-model = MultiTaskResNet50().to(device)
-model.load_state_dict(
-    torch.load("resnet50_multitask_mold.pth", map_location=device)
-)
+# Expecting checkpoint with keys:
+#   - "model": state_dict
+#   - "classes": list of class names (length N, mold at some index)
+ckpt = torch.load("best_convnext_multitask.pth", map_location=device)
+classes = ckpt.get("classes") or []
+num_classes = len(classes) if classes else 9
+mold_idx = classes.index("mold") if classes else 4
+
+model = MultiTaskConvNeXt(num_classes).to(device)
+model.load_state_dict(ckpt["model"])
 model.eval()
 
 # ------------------
@@ -48,9 +53,12 @@ transform = transforms.Compose([
 ])
 
 # ------------------
-# Grad-CAM
-# ------------------
-gradcam = GradCAM(model, model.backbone.layer4[-1].conv3)
+# 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:
+    target_layer = model.backbone.features[-1].block[-1].dwconv
+gradcam = GradCAM(model, target_layer)
 
 # ------------------
 # DINO (lazy loaded)

best_convnext_multitask.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fb0e73f75e0b9fbc2a548dec97791598d99b67791eec44d2aa35053f1e27e342
+size 350441583

model.py

@@ -2,6 +2,7 @@ import torch
 import torch.nn as nn
 from torchvision import models
 
+
 class MultiTaskResNet50(nn.Module):
     def __init__(self, num_classes=9):
         super().__init__()
@@ -11,10 +12,50 @@ class MultiTaskResNet50(nn.Module):
         self.class_head = nn.Linear(feat_dim, num_classes)
         self.bio_head = nn.Linear(feat_dim, 2)
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor):
         feats = self.backbone(x)
         return {
             "class": self.class_head(feats),
-            "bio": self.bio_head(feats)
+            "bio": self.bio_head(feats),
+        }
+
+
+class MultiTaskConvNeXt(nn.Module):
+    """
+    ConvNeXt-Base backbone with two heads:
+    - N-class structural/mold classifier
+    - 2-class biological vs non-biological head
+
+    Mirrors the training setup from the ConvNeXt Kaggle notebook.
+    """
+
+    def __init__(self, num_classes: int):
+        super().__init__()
+
+        # We load task-specific weights, so no ImageNet weights here.
+        self.backbone = models.convnext_base(weights=None)
+
+        # ConvNeXt classifier is [LayerNorm2d, Flatten, Linear]
+        feat_dim = self.backbone.classifier[2].in_features
+        self.backbone.classifier = nn.Identity()
+
+        self.pool = nn.AdaptiveAvgPool2d((1, 1))
+        self.class_head = nn.Linear(feat_dim, num_classes)
+        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.
+        self.last_conv = self.backbone.features[-1].block[-1].dwconv
+
+    def forward(self, x: torch.Tensor):
+        feats = self.backbone.features(x)
+        feats = self.pool(feats)
+        feats = torch.flatten(feats, 1)
+        feats = self.dropout(feats)
+
+        return {
+            "class": self.class_head(feats),
+            "bio": self.bio_head(feats),
         }
 
+