✨ Upgrade backend to support ConvNeXt automatic detection and ResNet18 fallback

toothmap_backend/models.py

@@ -4,6 +4,8 @@ from torchvision.models.detection import fasterrcnn_resnet50_fpn_v2
 from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
 from torchvision.models import resnet18
 import segmentation_models_pytorch as smp
+import torchvision
+from pathlib import Path
 import os
 
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -33,6 +35,43 @@ def _strip_dp(state: dict) -> dict:
         return {k.replace("module.", "", 1): v for k, v in state.items()}
     return state
 
+def load_convnext_model(ckpt_path: Path, device: torch.device):
+    """Loads ConvNeXt-Tiny customized for 32 FDI classes (>90% Acc Setup)."""
+    m = torchvision.models.convnext_tiny(weights=None)
+    m.classifier[2] = nn.Linear(m.classifier[2].in_features, 32)
+    
+    if not os.path.exists(ckpt_path):
+        print(f"⚠️  ConvNeXt Checkpoint not found: {ckpt_path}. Model will return random garbage until trained.")
+        m.to(device)
+        m.eval()
+        return m
+        
+    state = torch.load(ckpt_path, map_location=device, weights_only=True)
+    # Strip any "module." prefix if trained with DataParallel
+    new_state = {}
+    for k, v in state.items():
+        name = k[7:] if k.startswith("module.") else k
+        new_state[name] = v
+        
+    m.load_state_dict(new_state)
+    m.to(device)
+    m.eval()
+    return m
+
+def load_resnet_model(ckpt_path: Path, device: torch.device):
+    """Loads ResNet18 customized for 32 FDI classes."""
+    m = torchvision.models.resnet18(weights=None)
+    m.fc = nn.Sequential(nn.Dropout(p=0.3), nn.Linear(m.fc.in_features, 32))
+    state = torch.load(ckpt_path, map_location=device, weights_only=True)
+    new_state = {}
+    for k, v in state.items():
+        name = k[7:] if k.startswith("module.") else k
+        new_state[name] = v
+    m.load_state_dict(new_state)
+    m.to(device)
+    m.eval()
+    return m
+
 def _load(path: str, model):
     if not os.path.exists(path):
         print(f"⚠️  Checkpoint not found: {path}")
@@ -49,6 +88,17 @@ def _load(path: str, model):
 def load_models(ckpt_dir: str):
     print(f"📦 Loading models from: {ckpt_dir}")
     MODELS["frcnn"] = _load(os.path.join(ckpt_dir, "frcnn_best.pt"),        build_frcnn()).to(DEVICE).eval()
-    MODELS["cls"]   = _load(os.path.join(ckpt_dir, "resnet18_cls_best.pt"), build_classifier()).to(DEVICE).eval()
+    
+    convnext_path = Path(os.path.join(ckpt_dir, "convnext_best.pt"))
+    resnet_path = Path(os.path.join(ckpt_dir, "resnet18_cls_best.pt"))
+    if convnext_path.exists():
+        MODELS["cls"] = load_convnext_model(convnext_path, DEVICE)
+        print("✅ ConvNeXt loaded.")
+    elif resnet_path.exists():
+        MODELS["cls"] = load_resnet_model(resnet_path, DEVICE)
+        print("✅ ResNet18 loaded.")
+    else:
+        print("⚠️ No classification checkpoint found!")
+        
     MODELS["unet"]  = _load(os.path.join(ckpt_dir, "unet_resnet34_best.pt"), build_unet()).to(DEVICE).eval()
     print("🚀 All PyTorch models ready.")