rpotter6298
820 lines
31 KiB
Python
Executable File
820 lines
31 KiB
Python
Executable File
"""REFUGE glaucoma classification with rotation-based TTT."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Callable, Dict, Iterable, List, Optional, Sequence, Tuple
|
|
import random
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
import torch
|
|
from torch import nn
|
|
from torch.utils.data import DataLoader, Dataset
|
|
from torchvision import models, transforms
|
|
from torchvision.transforms import functional as TF
|
|
import torch.nn.functional as F
|
|
from sklearn.metrics import roc_auc_score
|
|
from skimage.transform import warp_polar
|
|
from tqdm import tqdm
|
|
|
|
from classes.geometry_features import (
|
|
FEATURE_DIM,
|
|
EPS,
|
|
compute_geometry_features,
|
|
disc_cup_from_mask_image,
|
|
)
|
|
from classes.refuge_preprocessing import RefugePreprocessing, RefugeSample
|
|
from classes.refuge_segmentation import RefugeSegmentation
|
|
from classes.unet_segmenter import UNetSegmenter
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Dataset utilities
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _default_image_transform(size: int = 256) -> transforms.Compose:
|
|
return transforms.Compose(
|
|
[
|
|
transforms.Resize((size, size)),
|
|
transforms.ToTensor(),
|
|
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
|
]
|
|
)
|
|
|
|
|
|
def _augment_image_transform(size: int = 256) -> transforms.Compose:
|
|
return transforms.Compose(
|
|
[
|
|
transforms.Resize((size, size)),
|
|
transforms.RandomHorizontalFlip(),
|
|
transforms.RandomRotation(10),
|
|
transforms.ColorJitter(0.1, 0.1, 0.1, 0.05),
|
|
transforms.ToTensor(),
|
|
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
|
]
|
|
)
|
|
|
|
|
|
def _crop_from_geometry(image: Image.Image, geometry: Dict[str, float], size: int = 256) -> Image.Image:
|
|
cx, cy = geometry["centre_x"], geometry["centre_y"]
|
|
r = geometry["crop_radius"]
|
|
left = max(0.0, cx - r)
|
|
upper = max(0.0, cy - r)
|
|
right = min(image.width, cx + r)
|
|
lower = min(image.height, cy + r)
|
|
crop = image.crop((left, upper, right, lower))
|
|
return crop.resize((size, size), Image.BILINEAR)
|
|
|
|
|
|
def _geometry_from_mask(mask: np.ndarray, scale: float) -> Dict[str, float]:
|
|
mask = np.asarray(mask) > 0
|
|
coords = np.argwhere(mask)
|
|
if coords.size == 0:
|
|
raise RuntimeError("Empty mask; cannot derive geometry")
|
|
ys, xs = coords[:, 0], coords[:, 1]
|
|
centre_x = float(xs.mean())
|
|
centre_y = float(ys.mean())
|
|
width = float(xs.max() - xs.min())
|
|
height = float(ys.max() - ys.min())
|
|
diameter = max(width, height)
|
|
radius = diameter / 2.0
|
|
crop_radius = radius * scale
|
|
return {
|
|
"centre_x": centre_x,
|
|
"centre_y": centre_y,
|
|
"radius": radius,
|
|
"crop_radius": crop_radius,
|
|
"crop_size": crop_radius * 2.0,
|
|
}
|
|
|
|
|
|
def _compute_feature_vector(disc_mask: np.ndarray, cup_mask: np.ndarray) -> np.ndarray:
|
|
return compute_geometry_features(disc_mask, cup_mask)
|
|
|
|
|
|
def _compute_polar_image(crop: Image.Image, size: int) -> Image.Image:
|
|
arr = np.asarray(crop).astype(np.float32) / 255.0
|
|
radius = min(arr.shape[0], arr.shape[1]) / 2.0
|
|
polar = warp_polar(
|
|
arr,
|
|
radius=radius,
|
|
scaling="linear",
|
|
channel_axis=-1,
|
|
)
|
|
polar = np.clip(polar, 0.0, 1.0)
|
|
polar_img = Image.fromarray((polar * 255).astype(np.uint8))
|
|
return polar_img.resize((size, size), Image.BILINEAR)
|
|
|
|
|
|
def _crop_mask_from_geometry(mask: np.ndarray, geometry: Dict[str, float], size: int) -> np.ndarray:
|
|
mask_img = Image.fromarray((mask > 0).astype(np.uint8) * 255)
|
|
cx, cy = geometry["centre_x"], geometry["centre_y"]
|
|
r = geometry["crop_radius"]
|
|
left = max(0.0, cx - r)
|
|
upper = max(0.0, cy - r)
|
|
right = min(mask_img.width, cx + r)
|
|
lower = min(mask_img.height, cy + r)
|
|
crop = mask_img.crop((left, upper, right, lower)).resize((size, size), Image.NEAREST)
|
|
return (np.asarray(crop) > 0).astype(np.uint8)
|
|
|
|
|
|
@dataclass
|
|
class RefugeClassificationRecord:
|
|
sample: RefugeSample
|
|
geometry: Dict[str, float]
|
|
disc_mask: Optional[np.ndarray] = None
|
|
cup_mask: Optional[np.ndarray] = None
|
|
|
|
|
|
class RefugeClassificationDataset(Dataset):
|
|
def __init__(
|
|
self,
|
|
records: Sequence[RefugeClassificationRecord],
|
|
transform: transforms.Compose,
|
|
polar_transform: transforms.Compose,
|
|
size: int = 256,
|
|
) -> None:
|
|
self.records = list(records)
|
|
self.transform = transform
|
|
self.polar_transform = polar_transform
|
|
self.size = size
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.records)
|
|
|
|
def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
|
|
rec = self.records[idx]
|
|
image = Image.open(rec.sample.image_path).convert("RGB")
|
|
crop = _crop_from_geometry(image, rec.geometry, size=self.size)
|
|
polar_image = _compute_polar_image(crop, size=self.size)
|
|
tensor = self.transform(crop)
|
|
polar_tensor = self.polar_transform(polar_image)
|
|
|
|
features = np.zeros((FEATURE_DIM,), dtype=np.float32)
|
|
if rec.disc_mask is not None and rec.cup_mask is not None:
|
|
disc_crop = _crop_mask_from_geometry(rec.disc_mask, rec.geometry, self.size)
|
|
cup_crop = _crop_mask_from_geometry(rec.cup_mask, rec.geometry, self.size)
|
|
features = _compute_feature_vector(disc_crop, cup_crop)
|
|
|
|
feature_tensor = torch.from_numpy(features).float()
|
|
label = rec.sample.label
|
|
if label is None:
|
|
raise ValueError(f"Sample {rec.sample.sample_id} is missing glaucoma label")
|
|
return {
|
|
"image": tensor,
|
|
"polar": polar_tensor,
|
|
"features": feature_tensor,
|
|
"label": torch.tensor(label, dtype=torch.long),
|
|
"sample_id": rec.sample.sample_id,
|
|
}
|
|
|
|
|
|
class RefugeTTTDataset(Dataset):
|
|
"""Dataset providing unlabeled crops for test-time training."""
|
|
|
|
def __init__(self, records: Sequence[RefugeClassificationRecord], transform: transforms.Compose, size: int = 256) -> None:
|
|
self.records = list(records)
|
|
self.transform = transform
|
|
self.size = size
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.records)
|
|
|
|
def __getitem__(self, idx: int) -> torch.Tensor:
|
|
rec = self.records[idx]
|
|
image = Image.open(rec.sample.image_path).convert("RGB")
|
|
crop = _crop_from_geometry(image, rec.geometry, size=self.size)
|
|
return self.transform(crop)
|
|
|
|
|
|
class UNetGeometryProvider:
|
|
"""Callable wrapper that derives disc geometry using a trained UNetSegmenter."""
|
|
|
|
def __init__(
|
|
self,
|
|
segmenter: UNetSegmenter,
|
|
threshold: float = 0.5,
|
|
tta: bool = False,
|
|
) -> None:
|
|
self.segmenter = segmenter
|
|
self.threshold = threshold
|
|
self.tta = tta
|
|
self.segmenter.model.eval()
|
|
|
|
def __call__(self, sample: RefugeSample, scale: float) -> Tuple[Dict[str, float], np.ndarray, np.ndarray]:
|
|
image = Image.open(sample.image_path).convert("RGB")
|
|
resized = self.segmenter.preprocess_image(image)
|
|
tensor = transforms.ToTensor()(resized)
|
|
tensor = self.segmenter._normalize_tensor(tensor)
|
|
tensor = tensor.unsqueeze(0).to(self.segmenter.device)
|
|
with torch.no_grad():
|
|
logits = self.segmenter.model(tensor)
|
|
if self.tta:
|
|
t_h = torch.flip(tensor, dims=[3])
|
|
log_h = self.segmenter.model(t_h)
|
|
log_h = torch.flip(log_h, dims=[3])
|
|
t_v = torch.flip(tensor, dims=[2])
|
|
log_v = self.segmenter.model(t_v)
|
|
log_v = torch.flip(log_v, dims=[2])
|
|
logits = (logits + log_h + log_v) / 3.0
|
|
probs = torch.sigmoid(logits)[0].cpu().numpy()
|
|
|
|
disc_pred = (probs[0] > self.threshold).astype(np.uint8) * 255
|
|
cup_pred = (probs[1] > self.threshold).astype(np.uint8) * 255
|
|
disc_img = Image.fromarray(disc_pred, mode="L").resize(image.size, Image.NEAREST)
|
|
cup_img = Image.fromarray(cup_pred, mode="L").resize(image.size, Image.NEAREST)
|
|
disc_mask = (np.array(disc_img, dtype=np.uint8) > 0).astype(np.uint8)
|
|
cup_mask = (np.array(cup_img, dtype=np.uint8) > 0).astype(np.uint8)
|
|
cup_mask = (cup_mask > 0) & (disc_mask > 0)
|
|
cup_mask = cup_mask.astype(np.uint8)
|
|
geom = _geometry_from_mask(disc_mask, scale)
|
|
return geom, disc_mask, cup_mask
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Classification module
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class ArcMarginProduct(nn.Module):
|
|
"""Additive angular margin (ArcFace) head."""
|
|
|
|
def __init__(
|
|
self,
|
|
in_features: int,
|
|
out_features: int,
|
|
s: float = 30.0,
|
|
m: float = 0.5,
|
|
easy_margin: bool = False,
|
|
) -> None:
|
|
super().__init__()
|
|
self.in_features = in_features
|
|
self.out_features = out_features
|
|
self.s = float(s)
|
|
self.m = float(m)
|
|
self.easy_margin = easy_margin
|
|
self.weight = nn.Parameter(torch.empty(out_features, in_features))
|
|
nn.init.xavier_uniform_(self.weight)
|
|
|
|
self.cos_m = math.cos(m)
|
|
self.sin_m = math.sin(m)
|
|
self.th = math.cos(math.pi - m)
|
|
self.mm = math.sin(math.pi - m) * m
|
|
|
|
def forward(self, input: torch.Tensor, label: Optional[torch.Tensor] = None) -> torch.Tensor:
|
|
cosine = F.linear(F.normalize(input), F.normalize(self.weight))
|
|
if label is None:
|
|
return cosine * self.s
|
|
|
|
sine = torch.sqrt(torch.clamp(1.0 - cosine.pow(2), min=0.0))
|
|
phi = cosine * self.cos_m - sine * self.sin_m
|
|
if self.easy_margin:
|
|
phi = torch.where(cosine > 0, phi, cosine)
|
|
else:
|
|
phi = torch.where(cosine > self.th, phi, cosine - self.mm)
|
|
|
|
one_hot = torch.zeros_like(cosine)
|
|
one_hot.scatter_(1, label.view(-1, 1), 1.0)
|
|
logits = (one_hot * phi) + ((1.0 - one_hot) * cosine)
|
|
logits *= self.s
|
|
return logits
|
|
|
|
|
|
class RefugeClassification:
|
|
"""Train and evaluate REFUGE glaucoma classifiers with TTT support."""
|
|
|
|
def __init__(
|
|
self,
|
|
preprocessing: RefugePreprocessing,
|
|
segmentation: RefugeSegmentation,
|
|
backbone: Optional[nn.Module] = None,
|
|
geometry_fn: Optional[
|
|
Callable[
|
|
[RefugeSample, float],
|
|
Tuple[Dict[str, float], Optional[np.ndarray], Optional[np.ndarray]],
|
|
]
|
|
] = None,
|
|
cache_dir: Optional[Path] = None,
|
|
use_all_labeled: bool = False,
|
|
auto_val_ratio: float = 0.1,
|
|
use_margin: bool = False,
|
|
margin_s: float = 30.0,
|
|
margin_m: float = 0.5,
|
|
) -> None:
|
|
self.preprocessing = preprocessing
|
|
self.segmentation = segmentation
|
|
if backbone is not None:
|
|
self.backbone = backbone
|
|
in_features = getattr(self.backbone, "_feature_dim", None)
|
|
if in_features is None:
|
|
if hasattr(self.backbone, "fc") and hasattr(self.backbone.fc, "in_features"):
|
|
in_features = self.backbone.fc.in_features # type: ignore[attr-defined]
|
|
self.backbone.fc = nn.Identity() # type: ignore[attr-defined]
|
|
else:
|
|
raise ValueError(
|
|
"Provided backbone must have '_feature_dim' or expose fc.in_features"
|
|
)
|
|
else:
|
|
self.backbone = self._default_backbone()
|
|
in_features = getattr(self.backbone, "_feature_dim", None)
|
|
if in_features is None:
|
|
in_features = self.backbone.fc.in_features # type: ignore[attr-defined]
|
|
self.backbone.fc = nn.Identity() # type: ignore[attr-defined]
|
|
self.feature_dim = in_features
|
|
self.use_polar = True
|
|
self.extra_feature_dim = FEATURE_DIM
|
|
combined_dim = self.feature_dim * (1 + int(self.use_polar)) + self.extra_feature_dim
|
|
self.margin_s = float(margin_s)
|
|
self.margin_m = float(margin_m)
|
|
self.use_margin = bool(use_margin)
|
|
if self.use_margin:
|
|
self.classifier_head = ArcMarginProduct(
|
|
combined_dim, 2, s=self.margin_s, m=self.margin_m
|
|
)
|
|
else:
|
|
self.classifier_head = nn.Linear(combined_dim, 2)
|
|
self.rotation_head = nn.Linear(self.feature_dim, 4)
|
|
|
|
self.train_dataset: Optional[RefugeClassificationDataset] = None
|
|
self.val_dataset: Optional[RefugeClassificationDataset] = None
|
|
self.train_loader: Optional[DataLoader] = None
|
|
self.val_loader: Optional[DataLoader] = None
|
|
self.ttt_transform = _default_image_transform()
|
|
self.train_transform = _augment_image_transform()
|
|
self.eval_transform = _default_image_transform()
|
|
self.polar_transform = _default_image_transform()
|
|
self.crop_scale = 2.5
|
|
self.crop_size = 256
|
|
self.geometry_cache: Dict[
|
|
str, Tuple[Dict[str, float], Optional[np.ndarray], Optional[np.ndarray]]
|
|
] = {}
|
|
self.train_records: List[RefugeClassificationRecord] = []
|
|
self.val_records: List[RefugeClassificationRecord] = []
|
|
self._geometry_fn = geometry_fn
|
|
self.cache_dir = cache_dir
|
|
if self.cache_dir is not None:
|
|
self.cache_dir.mkdir(parents=True, exist_ok=True)
|
|
self.use_all_labeled = use_all_labeled
|
|
self.auto_val_ratio = auto_val_ratio
|
|
|
|
# ------------------------------------------------------------------
|
|
@staticmethod
|
|
def _default_backbone() -> nn.Module:
|
|
weights = models.ResNet50_Weights.IMAGENET1K_V2
|
|
model = models.resnet50(weights=weights)
|
|
in_features = model.fc.in_features
|
|
model.fc = nn.Identity()
|
|
setattr(model, "_feature_dim", in_features)
|
|
return model
|
|
|
|
# ------------------------------------------------------------------
|
|
def build_datasets(
|
|
self,
|
|
crop_scale: float = 2.5,
|
|
crop_size: int = 256,
|
|
batch_size: int = 16,
|
|
num_workers: int = 4,
|
|
) -> None:
|
|
self.crop_scale = crop_scale
|
|
self.crop_size = crop_size
|
|
self.train_transform = _augment_image_transform(crop_size)
|
|
self.eval_transform = _default_image_transform(crop_size)
|
|
self.ttt_transform = _default_image_transform(crop_size)
|
|
self.polar_transform = _default_image_transform(crop_size)
|
|
|
|
manifest = list(self.preprocessing.build_manifest())
|
|
train_records: List[RefugeClassificationRecord] = []
|
|
val_records: List[RefugeClassificationRecord] = []
|
|
|
|
allowed_splits = {"train", "val"}
|
|
candidates = [
|
|
sample
|
|
for sample in manifest
|
|
if sample.label is not None and sample.split in allowed_splits
|
|
]
|
|
|
|
print(
|
|
f"[classifier] Building datasets from {len(candidates)} labelled samples (train/val)"
|
|
)
|
|
|
|
for sample in tqdm(
|
|
candidates,
|
|
desc="Preparing records",
|
|
unit="sample",
|
|
leave=False,
|
|
):
|
|
try:
|
|
geom, disc_mask, cup_mask = self._resolve_geometry(sample, crop_scale)
|
|
except RuntimeError:
|
|
continue
|
|
record = RefugeClassificationRecord(
|
|
sample=sample,
|
|
geometry=geom,
|
|
disc_mask=disc_mask,
|
|
cup_mask=cup_mask,
|
|
)
|
|
if sample.split == "train" or (
|
|
self.use_all_labeled and sample.split == "val"
|
|
):
|
|
train_records.append(record)
|
|
else:
|
|
val_records.append(record)
|
|
|
|
if (not val_records or self.use_all_labeled) and train_records and self.auto_val_ratio > 0.0:
|
|
rng = random.Random(42)
|
|
label_groups: Dict[int, List[RefugeClassificationRecord]] = {}
|
|
for rec in train_records:
|
|
label = int(rec.sample.label or 0)
|
|
label_groups.setdefault(label, []).append(rec)
|
|
|
|
new_train: List[RefugeClassificationRecord] = []
|
|
new_val: List[RefugeClassificationRecord] = []
|
|
for recs in label_groups.values():
|
|
rng.shuffle(recs)
|
|
if len(recs) <= 1:
|
|
new_train.extend(recs)
|
|
continue
|
|
val_count = max(1, int(round(len(recs) * self.auto_val_ratio)))
|
|
if val_count >= len(recs):
|
|
val_count = len(recs) - 1
|
|
new_val.extend(recs[:val_count])
|
|
new_train.extend(recs[val_count:])
|
|
|
|
if not new_val:
|
|
# Fallback: ensure at least one validation sample if possible
|
|
if len(new_train) > 1:
|
|
new_val.append(new_train.pop())
|
|
|
|
if new_val:
|
|
val_records = new_val
|
|
train_records = new_train
|
|
|
|
self.train_records = train_records
|
|
self.val_records = val_records
|
|
|
|
print(
|
|
f"[classifier] Records ready → train: {len(train_records)}, val: {len(val_records)}"
|
|
)
|
|
|
|
self.train_dataset = RefugeClassificationDataset(
|
|
train_records,
|
|
transform=self.train_transform,
|
|
polar_transform=self.polar_transform,
|
|
size=crop_size,
|
|
)
|
|
self.val_dataset = RefugeClassificationDataset(
|
|
val_records,
|
|
transform=self.eval_transform,
|
|
polar_transform=self.polar_transform,
|
|
size=crop_size,
|
|
)
|
|
|
|
self.train_loader = DataLoader(
|
|
self.train_dataset,
|
|
batch_size=batch_size,
|
|
shuffle=True,
|
|
num_workers=num_workers,
|
|
pin_memory=True,
|
|
)
|
|
self.val_loader = DataLoader(
|
|
self.val_dataset,
|
|
batch_size=batch_size,
|
|
shuffle=False,
|
|
num_workers=num_workers,
|
|
pin_memory=True,
|
|
)
|
|
|
|
print(
|
|
"[classifier] DataLoaders prepared — training batches will start shortly"
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
def _resolve_geometry(
|
|
self, sample: RefugeSample, scale: float
|
|
) -> Tuple[Dict[str, float], Optional[np.ndarray], Optional[np.ndarray]]:
|
|
key = self._cache_key(sample.sample_id, scale)
|
|
cached = self.geometry_cache.get(key)
|
|
if cached is not None:
|
|
return cached
|
|
|
|
cache_path = self._cache_path(sample.sample_id, scale)
|
|
if cache_path is not None and cache_path.exists():
|
|
data = np.load(cache_path, allow_pickle=False)
|
|
geom = {
|
|
"centre_x": float(data["centre_x"]),
|
|
"centre_y": float(data["centre_y"]),
|
|
"radius": float(data["radius"]),
|
|
"crop_radius": float(data["crop_radius"]),
|
|
"crop_size": float(data["crop_size"]),
|
|
}
|
|
disc_mask = None
|
|
cup_mask = None
|
|
if int(data["has_disc"]):
|
|
disc_mask = data["disc_mask"].astype(np.uint8)
|
|
if int(data["has_cup"]):
|
|
cup_mask = data["cup_mask"].astype(np.uint8)
|
|
self.geometry_cache[key] = (geom, disc_mask, cup_mask)
|
|
return geom, disc_mask, cup_mask
|
|
|
|
disc_mask: Optional[np.ndarray] = None
|
|
cup_mask: Optional[np.ndarray] = None
|
|
|
|
if sample.mask_path and sample.mask_path.exists():
|
|
mask_img = Image.open(sample.mask_path).convert("RGB")
|
|
disc_mask, cup_mask = disc_cup_from_mask_image(mask_img)
|
|
geom = _geometry_from_mask(disc_mask, scale)
|
|
elif self._geometry_fn is not None:
|
|
geom, disc_mask, cup_mask = self._geometry_fn(sample, scale)
|
|
else:
|
|
geom = self.segmentation.infer_disc_geometry(sample, scale=scale)
|
|
try:
|
|
pred_mask = self.segmentation.predict_mask(sample).numpy()
|
|
disc_mask = pred_mask.astype(np.uint8)
|
|
except Exception:
|
|
disc_mask = None
|
|
cup_mask = None
|
|
|
|
if cache_path is not None:
|
|
try:
|
|
np.savez_compressed(
|
|
cache_path,
|
|
centre_x=geom["centre_x"],
|
|
centre_y=geom["centre_y"],
|
|
radius=geom["radius"],
|
|
crop_radius=geom["crop_radius"],
|
|
crop_size=geom.get("crop_size", geom["crop_radius"] * 2.0),
|
|
disc_mask=disc_mask if disc_mask is not None else np.array([], dtype=np.uint8),
|
|
cup_mask=cup_mask if cup_mask is not None else np.array([], dtype=np.uint8),
|
|
has_disc=int(disc_mask is not None),
|
|
has_cup=int(cup_mask is not None),
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
self.geometry_cache[key] = (geom, disc_mask, cup_mask)
|
|
return geom, disc_mask, cup_mask
|
|
|
|
def set_geometry_fn(
|
|
self,
|
|
geometry_fn: Optional[
|
|
Callable[
|
|
[RefugeSample, float],
|
|
Tuple[Dict[str, float], Optional[np.ndarray], Optional[np.ndarray]],
|
|
]
|
|
],
|
|
) -> None:
|
|
self._geometry_fn = geometry_fn
|
|
self.geometry_cache.clear()
|
|
|
|
def build_records_for_samples(
|
|
self,
|
|
samples: Sequence[RefugeSample],
|
|
crop_scale: Optional[float] = None,
|
|
progress_prefix: Optional[str] = None,
|
|
) -> List[RefugeClassificationRecord]:
|
|
scale = crop_scale if crop_scale is not None else self.crop_scale
|
|
records: List[RefugeClassificationRecord] = []
|
|
iterator: Iterable[RefugeSample]
|
|
if progress_prefix is not None:
|
|
iterator = tqdm(samples, desc=progress_prefix, unit="sample", leave=False)
|
|
else:
|
|
iterator = samples
|
|
for sample in iterator:
|
|
if sample.label is None:
|
|
continue
|
|
try:
|
|
geom, disc_mask, cup_mask = self._resolve_geometry(sample, scale)
|
|
except RuntimeError:
|
|
continue
|
|
records.append(
|
|
RefugeClassificationRecord(
|
|
sample=sample,
|
|
geometry=geom,
|
|
disc_mask=disc_mask,
|
|
cup_mask=cup_mask,
|
|
)
|
|
)
|
|
return records
|
|
|
|
def _cache_key(self, sample_id: str, scale: float) -> str:
|
|
scale_tag = int(round(scale * 100))
|
|
return f"{sample_id}_s{scale_tag}"
|
|
|
|
def _cache_path(self, sample_id: str, scale: float) -> Optional[Path]:
|
|
if self.cache_dir is None:
|
|
return None
|
|
return self.cache_dir / f"{self._cache_key(sample_id, scale)}.npz"
|
|
|
|
# ------------------------------------------------------------------
|
|
def train(
|
|
self,
|
|
epochs: int = 30,
|
|
lr: float = 1e-4,
|
|
weight_decay: float = 1e-4,
|
|
device: Optional[str] = None,
|
|
rotation_weight: float = 0.5,
|
|
checkpoint_dir: Optional[Path] = None,
|
|
) -> Dict[str, float]:
|
|
if self.train_loader is None or self.val_loader is None:
|
|
self.build_datasets()
|
|
|
|
device = device or ("cuda" if torch.cuda.is_available() else "cpu")
|
|
self.backbone.to(device)
|
|
self.classifier_head.to(device)
|
|
self.rotation_head.to(device)
|
|
|
|
params = list(self.backbone.parameters()) + list(self.classifier_head.parameters()) + list(self.rotation_head.parameters())
|
|
optimizer = torch.optim.Adam(params, lr=lr, weight_decay=weight_decay)
|
|
clf_loss = nn.CrossEntropyLoss()
|
|
rot_loss = nn.CrossEntropyLoss()
|
|
|
|
best_auc = 0.0
|
|
history: Dict[str, float] = {}
|
|
|
|
epoch_iter = tqdm(range(1, epochs + 1), desc="Epochs", unit="epoch")
|
|
|
|
print(
|
|
f"[classifier] Starting training for {epochs} epochs with batch size {self.train_loader.batch_size}"
|
|
)
|
|
|
|
for epoch in epoch_iter:
|
|
self.backbone.train()
|
|
self.classifier_head.train()
|
|
self.rotation_head.train()
|
|
running_loss = 0.0
|
|
|
|
batch_iter = tqdm(
|
|
self.train_loader, # type: ignore[arg-type]
|
|
desc=f"Train {epoch}/{epochs}",
|
|
leave=False,
|
|
unit="batch",
|
|
)
|
|
|
|
for batch in batch_iter:
|
|
images = batch["image"].to(device)
|
|
polars = batch["polar"].to(device)
|
|
extra_feats = batch["features"].to(device)
|
|
labels = batch["label"].to(device)
|
|
optimizer.zero_grad()
|
|
|
|
feats_img = self.backbone(images)
|
|
feats = feats_img
|
|
if self.use_polar:
|
|
feats_polar = self.backbone(polars)
|
|
feats = torch.cat([feats, feats_polar], dim=1)
|
|
if self.extra_feature_dim > 0:
|
|
feats = torch.cat([feats, extra_feats], dim=1)
|
|
if self.use_margin:
|
|
logits = self.classifier_head(feats, labels)
|
|
else:
|
|
logits = self.classifier_head(feats)
|
|
loss_cls = clf_loss(logits, labels)
|
|
|
|
rot_imgs, rot_labels = self._build_rotation_batch(images)
|
|
feats_rot = self.backbone(rot_imgs)
|
|
logits_rot = self.rotation_head(feats_rot)
|
|
loss_rot = rot_loss(logits_rot, rot_labels)
|
|
|
|
loss = loss_cls + rotation_weight * loss_rot
|
|
loss.backward()
|
|
optimizer.step()
|
|
running_loss += loss.item() * images.size(0)
|
|
|
|
train_loss = running_loss / len(self.train_loader.dataset) # type: ignore[arg-type]
|
|
metrics = self.evaluate(device=device)
|
|
history[f"epoch_{epoch}_loss"] = train_loss
|
|
history[f"epoch_{epoch}_auc"] = metrics.get("auc", float("nan"))
|
|
|
|
auc_val = metrics.get("auc", 0.0)
|
|
epoch_iter.set_postfix(loss=f"{train_loss:.4f}", auc=f"{auc_val:.4f}")
|
|
|
|
if auc_val > best_auc:
|
|
best_auc = metrics["auc"]
|
|
if checkpoint_dir is not None:
|
|
checkpoint_dir.mkdir(parents=True, exist_ok=True)
|
|
torch.save({
|
|
"backbone": self.backbone.state_dict(),
|
|
"classifier": self.classifier_head.state_dict(),
|
|
"rotation": self.rotation_head.state_dict(),
|
|
}, checkpoint_dir / "refuge_classifier_best.pt")
|
|
|
|
return {"best_auc": best_auc, **history}
|
|
|
|
# ------------------------------------------------------------------
|
|
def evaluate(
|
|
self,
|
|
split: str = "val",
|
|
apply_ttt: bool = False,
|
|
device: Optional[str] = None,
|
|
) -> Dict[str, float]:
|
|
if split != "val":
|
|
raise ValueError("Only validation split supported currently")
|
|
if self.val_loader is None:
|
|
self.build_datasets()
|
|
|
|
device = device or ("cuda" if torch.cuda.is_available() else "cpu")
|
|
self.backbone.to(device)
|
|
self.classifier_head.to(device)
|
|
self.rotation_head.to(device)
|
|
|
|
if apply_ttt:
|
|
ttt_ds = RefugeTTTDataset(self.val_records, transform=self.ttt_transform, size=self.crop_size)
|
|
ttt_loader = DataLoader(ttt_ds, batch_size=32, shuffle=False)
|
|
self.apply_ttt(ttt_loader, device=device)
|
|
|
|
self.backbone.eval()
|
|
self.classifier_head.eval()
|
|
preds: List[float] = []
|
|
targets: List[int] = []
|
|
|
|
with torch.no_grad():
|
|
val_iter = tqdm(self.val_loader, desc="Validate", leave=False, unit="batch")
|
|
for batch in val_iter: # type: ignore[arg-type]
|
|
images = batch["image"].to(device)
|
|
labels = batch["label"].to(device)
|
|
polars = batch["polar"].to(device)
|
|
extra_feats = batch["features"].to(device)
|
|
feats_img = self.backbone(images)
|
|
feats = feats_img
|
|
if self.use_polar:
|
|
feats_polar = self.backbone(polars)
|
|
feats = torch.cat([feats, feats_polar], dim=1)
|
|
if self.extra_feature_dim > 0:
|
|
feats = torch.cat([feats, extra_feats], dim=1)
|
|
if self.use_margin:
|
|
logits = self.classifier_head(feats)
|
|
else:
|
|
logits = self.classifier_head(feats)
|
|
probs = torch.softmax(logits, dim=1)[:, 1]
|
|
preds.extend(probs.cpu().numpy().tolist())
|
|
targets.extend(labels.cpu().numpy().tolist())
|
|
|
|
auc = 0.0
|
|
try:
|
|
if len(set(targets)) > 1:
|
|
auc = float(roc_auc_score(targets, preds))
|
|
except ValueError:
|
|
auc = 0.0
|
|
|
|
return {"auc": auc}
|
|
|
|
# ------------------------------------------------------------------
|
|
def apply_ttt(self, loader: DataLoader, device: Optional[str] = None, steps: int = 1, lr: float = 1e-5) -> None:
|
|
device = device or ("cuda" if torch.cuda.is_available() else "cpu")
|
|
self.backbone.to(device)
|
|
self.rotation_head.to(device)
|
|
self.backbone.train()
|
|
self.rotation_head.train()
|
|
|
|
optimizer = torch.optim.Adam(list(self.backbone.parameters()) + list(self.rotation_head.parameters()), lr=lr)
|
|
criterion = nn.CrossEntropyLoss()
|
|
|
|
for _ in range(steps):
|
|
for batch in tqdm(loader, desc="TTT adapt", leave=False, unit="batch"):
|
|
if isinstance(batch, dict):
|
|
images = batch["image"].to(device)
|
|
else:
|
|
images = batch.to(device)
|
|
optimizer.zero_grad()
|
|
rot_imgs, rot_labels = self._build_rotation_batch(images)
|
|
feats = self.backbone(rot_imgs)
|
|
logits = self.rotation_head(feats)
|
|
loss = criterion(logits, rot_labels)
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
# ------------------------------------------------------------------
|
|
def extract_backbone(self) -> nn.Module:
|
|
return self.backbone
|
|
|
|
def save_checkpoint(self, output_dir: Path) -> None:
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
torch.save({
|
|
"backbone": self.backbone.state_dict(),
|
|
"classifier": self.classifier_head.state_dict(),
|
|
"rotation": self.rotation_head.state_dict(),
|
|
}, output_dir / "refuge_classifier.pt")
|
|
|
|
def load_checkpoint(self, checkpoint_path: Path) -> None:
|
|
payload = torch.load(checkpoint_path, map_location="cpu")
|
|
self.backbone.load_state_dict(payload["backbone"])
|
|
self.classifier_head.load_state_dict(payload["classifier"])
|
|
self.rotation_head.load_state_dict(payload["rotation"])
|
|
|
|
# ------------------------------------------------------------------
|
|
def _build_rotation_batch(self, images: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
rotations = [0, 90, 180, 270]
|
|
rotated = []
|
|
labels = []
|
|
for idx, angle in enumerate(rotations):
|
|
rot = TF.rotate(images, angle)
|
|
rotated.append(rot)
|
|
labels.append(torch.full((images.size(0),), idx, dtype=torch.long, device=images.device))
|
|
batch = torch.cat(rotated, dim=0)
|
|
batch_labels = torch.cat(labels, dim=0)
|
|
return batch, batch_labels
|