Spaces:
Sleeping
Sleeping
File size: 4,047 Bytes
e73e5ff d6cbcc6 e73e5ff d6cbcc6 e73e5ff d6cbcc6 e73e5ff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | import torch
from torch import nn
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')
print(f"🖥️ Using device: {DEVICE}")
# Shared model registry - avoids Python import scoping issues with module globals
MODELS = {}
def build_frcnn():
model = fasterrcnn_resnet50_fpn_v2(weights=None)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 33)
return model
def build_classifier():
model = resnet18(weights=None)
model.fc = nn.Sequential(nn.Dropout(0.3), nn.Linear(model.fc.in_features, 32))
return model
def build_unet():
return smp.Unet(encoder_name="resnet34", encoder_weights=None, in_channels=3, classes=1)
def _strip_dp(state: dict) -> dict:
"""Strip 'module.' prefix added by nn.DataParallel on Kaggle multi-GPU."""
if any(k.startswith("module.") for k in state.keys()):
print(" ↳ Stripping DataParallel 'module.' prefix from checkpoint")
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}")
return model
try:
state = torch.load(path, map_location=DEVICE, weights_only=False)
state = _strip_dp(state)
model.load_state_dict(state, strict=False)
print(f"✅ Loaded: {os.path.basename(path)}")
except Exception as e:
print(f"❌ Failed loading {os.path.basename(path)}: {e}")
return 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()
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.")
|