reworked iop_corr, added explainability tools and plotting tools, cleanup codebase
This commit is contained in:
rpotter6298
Executable
+894
@@ -0,0 +1,894 @@
|
||||
"""U-Net based optic disc/cup segmenter for REFUGE + Papila."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import os
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Iterable, List, Optional, Set, Tuple
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from PIL import Image, ImageDraw, ImageOps
|
||||
from PIL.Image import Resampling
|
||||
from skimage import measure
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from torchvision import transforms
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
@dataclass
|
||||
class ManifestEntry:
|
||||
sample_id: str
|
||||
dataset: str
|
||||
image_path: Path
|
||||
annotation_disc: Path
|
||||
annotation_cup: Path
|
||||
annotation_type_disc: str
|
||||
annotation_type_cup: str
|
||||
split: str # train / holdout / etc.
|
||||
|
||||
|
||||
class UNet(nn.Module):
|
||||
def __init__(
|
||||
self, in_channels: int = 3, base_channels: int = 32, out_channels: int = 2
|
||||
):
|
||||
super().__init__()
|
||||
self.enc1 = self._block(in_channels, base_channels)
|
||||
self.enc2 = self._block(base_channels, base_channels * 2)
|
||||
self.enc3 = self._block(base_channels * 2, base_channels * 4)
|
||||
self.enc4 = self._block(base_channels * 4, base_channels * 8)
|
||||
|
||||
self.pool = nn.MaxPool2d(2)
|
||||
self.bottleneck = self._block(base_channels * 8, base_channels * 16)
|
||||
|
||||
self.up4 = nn.ConvTranspose2d(
|
||||
base_channels * 16, base_channels * 8, 2, stride=2
|
||||
)
|
||||
self.dec4 = self._block(base_channels * 16, base_channels * 8)
|
||||
self.up3 = nn.ConvTranspose2d(base_channels * 8, base_channels * 4, 2, stride=2)
|
||||
self.dec3 = self._block(base_channels * 8, base_channels * 4)
|
||||
self.up2 = nn.ConvTranspose2d(base_channels * 4, base_channels * 2, 2, stride=2)
|
||||
self.dec2 = self._block(base_channels * 4, base_channels * 2)
|
||||
self.up1 = nn.ConvTranspose2d(base_channels * 2, base_channels, 2, stride=2)
|
||||
self.dec1 = self._block(base_channels * 2, base_channels)
|
||||
|
||||
self.out_conv = nn.Conv2d(base_channels, out_channels, kernel_size=1)
|
||||
|
||||
@staticmethod
|
||||
def _block(in_ch: int, out_ch: int) -> nn.Module:
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1, bias=False),
|
||||
nn.BatchNorm2d(out_ch),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1, bias=False),
|
||||
nn.BatchNorm2d(out_ch),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
e1 = self.enc1(x)
|
||||
e2 = self.enc2(self.pool(e1))
|
||||
e3 = self.enc3(self.pool(e2))
|
||||
e4 = self.enc4(self.pool(e3))
|
||||
b = self.bottleneck(self.pool(e4))
|
||||
|
||||
d4 = self.up4(b)
|
||||
d4 = torch.cat([d4, e4], dim=1)
|
||||
d4 = self.dec4(d4)
|
||||
d3 = self.up3(d4)
|
||||
d3 = torch.cat([d3, e3], dim=1)
|
||||
d3 = self.dec3(d3)
|
||||
d2 = self.up2(d3)
|
||||
d2 = torch.cat([d2, e2], dim=1)
|
||||
d2 = self.dec2(d2)
|
||||
d1 = self.up1(d2)
|
||||
d1 = torch.cat([d1, e1], dim=1)
|
||||
d1 = self.dec1(d1)
|
||||
return self.out_conv(d1)
|
||||
|
||||
|
||||
class SegmentationDataset(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
entries: List[ManifestEntry],
|
||||
segmenter: "UNetSegmenter",
|
||||
augment: bool,
|
||||
) -> None:
|
||||
self.entries = entries
|
||||
self.segmenter = segmenter
|
||||
self.augment = augment
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.entries)
|
||||
|
||||
def __getitem__(self, idx: int):
|
||||
entry = self.entries[idx]
|
||||
image = self.segmenter.load_preprocessed_image(entry)
|
||||
disc_mask, cup_mask = self.segmenter.load_masks(entry)
|
||||
|
||||
if self.augment:
|
||||
image = self.segmenter.jitter_image(image)
|
||||
image, disc_mask, cup_mask = self.segmenter.augment_geometric(
|
||||
image, disc_mask, cup_mask
|
||||
)
|
||||
image_tensor = transforms.ToTensor()(image)
|
||||
image_tensor = self.segmenter._normalize_tensor(image_tensor)
|
||||
|
||||
mask = np.stack([disc_mask, cup_mask], axis=0).astype(np.float32)
|
||||
mask_tensor = torch.from_numpy(mask)
|
||||
return image_tensor, mask_tensor
|
||||
|
||||
|
||||
class UNetSegmenter:
|
||||
def __init__(
|
||||
self,
|
||||
manifest_path: Path,
|
||||
device: Optional[str] = None,
|
||||
cup_weight: float = 1.0,
|
||||
disc_weight: float = 1.0,
|
||||
target_size: int = 512,
|
||||
val_ratio: float = 0.1,
|
||||
train_datasets: Optional[Iterable[str]] = None,
|
||||
val_datasets: Optional[Iterable[str]] = None,
|
||||
holdout_datasets: Optional[Iterable[str]] = None,
|
||||
normalize: str = "none",
|
||||
use_stronger_aug: bool = False,
|
||||
mask_cache_dir: Optional[Path] = None,
|
||||
image_cache_dir: Optional[Path] = None,
|
||||
in_memory_cache: bool = False,
|
||||
loader_workers: int = 0,
|
||||
) -> None:
|
||||
self.manifest_path = manifest_path
|
||||
self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
|
||||
self.cup_weight = cup_weight
|
||||
self.disc_weight = disc_weight
|
||||
self.target_size = target_size
|
||||
self.val_ratio = val_ratio
|
||||
self.normalize = (normalize or "none").lower()
|
||||
self.use_stronger_aug = bool(use_stronger_aug)
|
||||
self.mask_cache_dir = Path(mask_cache_dir).resolve() if mask_cache_dir else None
|
||||
if self.mask_cache_dir:
|
||||
self.mask_cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.image_cache_dir = Path(image_cache_dir).resolve() if image_cache_dir else None
|
||||
if self.image_cache_dir:
|
||||
self.image_cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.in_memory_cache = bool(in_memory_cache)
|
||||
self._mem_image_cache: dict[str, np.ndarray] = {}
|
||||
self._mem_mask_cache: dict[str, Tuple[np.ndarray, np.ndarray]] = {}
|
||||
self.loader_workers = max(0, int(loader_workers))
|
||||
|
||||
self.train_dataset_filter = self._normalize_filter(train_datasets)
|
||||
self.val_dataset_filter = self._normalize_filter(val_datasets)
|
||||
self.holdout_dataset_filter = self._normalize_filter(holdout_datasets)
|
||||
|
||||
self.model = UNet().to(self.device)
|
||||
self._manifest: List[ManifestEntry] = []
|
||||
self.train_entries: List[ManifestEntry] = []
|
||||
self.val_entries: List[ManifestEntry] = []
|
||||
self.holdout_entries: List[ManifestEntry] = []
|
||||
self.read_manifest()
|
||||
|
||||
def prebuild_in_memory_cache(
|
||||
self,
|
||||
*,
|
||||
cache_workers: int = 0,
|
||||
include_train: bool = True,
|
||||
include_val: bool = True,
|
||||
include_holdout: bool = False,
|
||||
) -> None:
|
||||
if not self.in_memory_cache:
|
||||
return
|
||||
selected: List[ManifestEntry] = []
|
||||
if include_train:
|
||||
selected.extend(self.train_entries)
|
||||
if include_val:
|
||||
selected.extend(self.val_entries)
|
||||
if include_holdout:
|
||||
selected.extend(self.holdout_entries)
|
||||
if not selected:
|
||||
return
|
||||
|
||||
# Deduplicate by cache key.
|
||||
dedup = {}
|
||||
for entry in selected:
|
||||
dedup[self._entry_cache_key(entry)] = entry
|
||||
entries = list(dedup.values())
|
||||
workers = max(0, int(cache_workers))
|
||||
print(
|
||||
f"[UNetSegmenter] prebuilding in-memory cache for {len(entries)} samples "
|
||||
f"(cache_workers={workers})",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
def _warm_one(entry: ManifestEntry) -> None:
|
||||
self.load_preprocessed_image(entry)
|
||||
self.load_masks(entry)
|
||||
|
||||
if workers <= 1:
|
||||
for entry in tqdm(entries, desc="Warm cache", unit="sample"):
|
||||
_warm_one(entry)
|
||||
else:
|
||||
with ThreadPoolExecutor(max_workers=workers) as ex:
|
||||
futures = [ex.submit(_warm_one, entry) for entry in entries]
|
||||
for fut in tqdm(as_completed(futures), total=len(futures), desc="Warm cache", unit="sample"):
|
||||
fut.result()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def read_manifest(self) -> None:
|
||||
df = pd.read_csv(self.manifest_path)
|
||||
entries: List[ManifestEntry] = []
|
||||
for _, row in df.iterrows():
|
||||
entry = ManifestEntry(
|
||||
sample_id=row["sample_id"],
|
||||
dataset=row["dataset"],
|
||||
image_path=Path(row["image_path"]),
|
||||
annotation_disc=Path(row["annotation_disc"]),
|
||||
annotation_cup=Path(row["annotation_cup"]),
|
||||
annotation_type_disc=row["annotation_type_disc"],
|
||||
annotation_type_cup=row["annotation_type_cup"],
|
||||
split=row["split"],
|
||||
)
|
||||
entries.append(entry)
|
||||
self._manifest = entries
|
||||
self.holdout_entries = [e for e in entries if e.split == "holdout"]
|
||||
if self.holdout_dataset_filter is not None:
|
||||
self.holdout_entries = [
|
||||
e for e in self.holdout_entries if e.dataset in self.holdout_dataset_filter
|
||||
]
|
||||
|
||||
trainable = [e for e in entries if e.split != "holdout"]
|
||||
if self.train_dataset_filter is not None:
|
||||
trainable = [
|
||||
e for e in trainable if e.dataset in self.train_dataset_filter
|
||||
]
|
||||
|
||||
if not trainable:
|
||||
self.val_entries = []
|
||||
self.train_entries = []
|
||||
return
|
||||
|
||||
val_pool = trainable
|
||||
if self.val_dataset_filter is not None:
|
||||
filtered = [e for e in trainable if e.dataset in self.val_dataset_filter]
|
||||
if filtered:
|
||||
val_pool = filtered
|
||||
|
||||
if len(trainable) == 1:
|
||||
val_count = 0
|
||||
else:
|
||||
val_count = max(1, int(len(trainable) * self.val_ratio))
|
||||
val_count = min(val_count, len(val_pool), len(trainable) - 1)
|
||||
|
||||
selected_val: List[ManifestEntry] = []
|
||||
if val_count > 0:
|
||||
selected_val = list(val_pool[:val_count])
|
||||
self.val_entries = selected_val
|
||||
selected_ids = {id(item) for item in selected_val}
|
||||
self.train_entries = [e for e in trainable if id(e) not in selected_ids]
|
||||
|
||||
if not self.train_entries and trainable:
|
||||
# Fallback when filtering removed all train entries (e.g. val_count forced entire set)
|
||||
self.train_entries = trainable
|
||||
self.val_entries = []
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def preprocess_image(self, image: Image.Image) -> Image.Image:
|
||||
return image.resize((self.target_size, self.target_size), Resampling.BILINEAR)
|
||||
|
||||
def jitter_image(self, image: Image.Image) -> Image.Image:
|
||||
# Photometric jitter only; geometric ops are applied jointly (image+mask)
|
||||
return transforms.ColorJitter(0.1, 0.1, 0.1, 0.05)(image)
|
||||
|
||||
def augment_geometric(
|
||||
self,
|
||||
image: Image.Image,
|
||||
disc_mask: np.ndarray,
|
||||
cup_mask: np.ndarray,
|
||||
) -> tuple[Image.Image, np.ndarray, np.ndarray]:
|
||||
if not self.use_stronger_aug:
|
||||
return image, disc_mask, cup_mask
|
||||
|
||||
img = image
|
||||
disc_pil = Image.fromarray((disc_mask > 0).astype(np.uint8) * 255)
|
||||
cup_pil = Image.fromarray((cup_mask > 0).astype(np.uint8) * 255)
|
||||
|
||||
# Random horizontal flip
|
||||
if np.random.rand() < 0.5:
|
||||
img = ImageOps.mirror(img)
|
||||
disc_pil = ImageOps.mirror(disc_pil)
|
||||
cup_pil = ImageOps.mirror(cup_pil)
|
||||
# Random vertical flip
|
||||
if np.random.rand() < 0.5:
|
||||
img = ImageOps.flip(img)
|
||||
disc_pil = ImageOps.flip(disc_pil)
|
||||
cup_pil = ImageOps.flip(cup_pil)
|
||||
# Random rotation (multiples of 90° to keep masks aligned)
|
||||
rotations = np.random.choice([0, 90, 180, 270])
|
||||
if rotations:
|
||||
img = img.rotate(rotations, expand=False)
|
||||
disc_pil = disc_pil.rotate(rotations, expand=False)
|
||||
cup_pil = cup_pil.rotate(rotations, expand=False)
|
||||
|
||||
disc_mask = (np.array(disc_pil) > 0).astype(np.float32)
|
||||
cup_mask = (np.array(cup_pil) > 0).astype(np.float32)
|
||||
return img, disc_mask, cup_mask
|
||||
|
||||
@staticmethod
|
||||
def _slugify(text: str) -> str:
|
||||
return "".join(ch if ch.isalnum() or ch in ("-", "_") else "_" for ch in text)
|
||||
|
||||
def _entry_cache_key(self, entry: ManifestEntry) -> str:
|
||||
return self._slugify(f"{entry.dataset}_{entry.sample_id}_sz{self.target_size}")
|
||||
|
||||
def _mask_cache_path(self, entry: ManifestEntry) -> Optional[Path]:
|
||||
if self.mask_cache_dir is None:
|
||||
return None
|
||||
slug = self._slugify(f"{entry.dataset}_{entry.sample_id}")
|
||||
fname = f"{slug}_sz{self.target_size}.npz"
|
||||
return self.mask_cache_dir / fname
|
||||
|
||||
def _image_cache_path(self, entry: ManifestEntry) -> Optional[Path]:
|
||||
if self.image_cache_dir is None:
|
||||
return None
|
||||
slug = self._slugify(f"{entry.dataset}_{entry.sample_id}")
|
||||
fname = f"{slug}_img_sz{self.target_size}.npz"
|
||||
return self.image_cache_dir / fname
|
||||
|
||||
def _load_image_cache(self, cache_path: Path) -> Optional[Image.Image]:
|
||||
try:
|
||||
data = np.load(str(cache_path), allow_pickle=False)
|
||||
arr = data["image"].astype(np.uint8, copy=False)
|
||||
if arr.ndim != 3 or arr.shape[2] != 3:
|
||||
return None
|
||||
return Image.fromarray(arr, mode="RGB")
|
||||
except Exception:
|
||||
with suppress(OSError, FileNotFoundError):
|
||||
cache_path.unlink()
|
||||
return None
|
||||
|
||||
def _save_image_cache(self, cache_path: Optional[Path], image: Image.Image) -> None:
|
||||
if cache_path is None:
|
||||
return
|
||||
cache_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp_path = cache_path.with_suffix(cache_path.suffix + ".tmp.npz")
|
||||
try:
|
||||
arr = np.asarray(image, dtype=np.uint8)
|
||||
np.savez_compressed(tmp_path, image=arr)
|
||||
os.replace(tmp_path, cache_path)
|
||||
except Exception:
|
||||
with suppress(OSError, FileNotFoundError):
|
||||
tmp_path.unlink()
|
||||
|
||||
def load_preprocessed_image(self, entry: ManifestEntry) -> Image.Image:
|
||||
key = self._entry_cache_key(entry)
|
||||
if self.in_memory_cache:
|
||||
cached = self._mem_image_cache.get(key)
|
||||
if cached is not None:
|
||||
return Image.fromarray(cached, mode="RGB")
|
||||
cache_path = self._image_cache_path(entry)
|
||||
if cache_path and cache_path.exists():
|
||||
cached = self._load_image_cache(cache_path)
|
||||
if cached is not None:
|
||||
if self.in_memory_cache:
|
||||
self._mem_image_cache[key] = np.asarray(cached, dtype=np.uint8)
|
||||
return cached
|
||||
image = Image.open(entry.image_path).convert("RGB")
|
||||
image = self.preprocess_image(image)
|
||||
if self.in_memory_cache:
|
||||
self._mem_image_cache[key] = np.asarray(image, dtype=np.uint8)
|
||||
self._save_image_cache(cache_path, image)
|
||||
return image
|
||||
|
||||
def _load_mask_cache(self, cache_path: Path) -> Optional[Tuple[np.ndarray, np.ndarray]]:
|
||||
try:
|
||||
data = np.load(str(cache_path), allow_pickle=False)
|
||||
disc = data["disc"].astype(np.float32)
|
||||
cup = data["cup"].astype(np.float32)
|
||||
return disc, cup
|
||||
except Exception:
|
||||
with suppress(OSError, FileNotFoundError):
|
||||
cache_path.unlink()
|
||||
return None
|
||||
|
||||
def _save_mask_cache(
|
||||
self,
|
||||
cache_path: Optional[Path],
|
||||
disc_mask: np.ndarray,
|
||||
cup_mask: np.ndarray,
|
||||
) -> None:
|
||||
if cache_path is None:
|
||||
return
|
||||
cache_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp_path = cache_path.with_suffix(cache_path.suffix + ".tmp.npz")
|
||||
try:
|
||||
np.savez_compressed(
|
||||
tmp_path,
|
||||
disc=disc_mask.astype(np.uint8),
|
||||
cup=cup_mask.astype(np.uint8),
|
||||
)
|
||||
os.replace(tmp_path, cache_path)
|
||||
except Exception:
|
||||
with suppress(OSError, FileNotFoundError):
|
||||
tmp_path.unlink()
|
||||
|
||||
def _normalize_tensor(self, tensor: torch.Tensor) -> torch.Tensor:
|
||||
if self.normalize == "per_image":
|
||||
mean = tensor.mean(dim=(1, 2), keepdim=True)
|
||||
std = tensor.std(dim=(1, 2), keepdim=True).clamp(min=1e-6)
|
||||
return (tensor - mean) / std
|
||||
if self.normalize == "imagenet":
|
||||
mean = torch.tensor([0.485, 0.456, 0.406]).view(-1, 1, 1)
|
||||
std = torch.tensor([0.229, 0.224, 0.225]).view(-1, 1, 1)
|
||||
return (tensor - mean) / std
|
||||
return tensor
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def extract_masks_from_image(
|
||||
self,
|
||||
mask_path: Path,
|
||||
disc_color: Optional[tuple[int, int, int]] = None,
|
||||
cup_color: Optional[tuple[int, int, int]] = None,
|
||||
) -> Tuple[np.ndarray, Optional[np.ndarray], Tuple[int, int]]:
|
||||
raw = Image.open(mask_path)
|
||||
arr = np.array(raw)
|
||||
if arr.ndim == 2:
|
||||
h, w = arr.shape
|
||||
flat = arr.reshape(-1).astype(np.int64, copy=False)
|
||||
edges = np.concatenate([arr[0, :], arr[-1, :], arr[:, 0], arr[:, -1]], axis=0).astype(np.int64, copy=False)
|
||||
edge_counts = np.bincount(edges, minlength=256)
|
||||
bg_val = int(np.argmax(edge_counts))
|
||||
counts = np.bincount(flat, minlength=256)
|
||||
counts[bg_val] = 0
|
||||
vals = np.where(counts > 0)[0]
|
||||
if vals.size < 1:
|
||||
raise ValueError(f"Mask {mask_path} does not contain discernible labels")
|
||||
# Disc = ALL non-background pixels (full optic disc: rim + cup combined).
|
||||
# Previously this was rim-only, which caused the cup structural prior
|
||||
# (cup & disc) to produce empty cup masks since cup and rim don't overlap.
|
||||
disc_mask = (arr != bg_val).astype(np.uint8)
|
||||
# Cup = the darkest non-background value (0 in REFUGE = inner cup region).
|
||||
# Using min-value rather than frequency avoids swapping when cup area > rim area.
|
||||
cup_val = int(np.min(vals)) if vals.size > 1 else None
|
||||
cup_mask = (arr == cup_val).astype(np.uint8) if cup_val is not None else np.zeros_like(disc_mask, dtype=np.uint8)
|
||||
return disc_mask, cup_mask if cup_mask.any() else None, (w, h)
|
||||
|
||||
image = raw.convert("RGB")
|
||||
arr = np.array(image)
|
||||
h, w, c = arr.shape
|
||||
|
||||
if disc_color is None or cup_color is None:
|
||||
# Fast color discovery via NumPy (avoid Python-level per-pixel tuple counting).
|
||||
edges = np.concatenate(
|
||||
[arr[0, :, :], arr[-1, :, :], arr[:, 0, :], arr[:, -1, :]], axis=0
|
||||
)
|
||||
edge_colors, edge_counts = np.unique(edges.reshape(-1, c), axis=0, return_counts=True)
|
||||
bg_color_np = edge_colors[int(np.argmax(edge_counts))]
|
||||
|
||||
colors_np, counts_np = np.unique(arr.reshape(-1, c), axis=0, return_counts=True)
|
||||
keep = np.any(colors_np != bg_color_np.reshape(1, -1), axis=1)
|
||||
colors_np = colors_np[keep]
|
||||
counts_np = counts_np[keep]
|
||||
if colors_np.shape[0] < 1:
|
||||
raise ValueError(f"Mask {mask_path} does not contain discernible labels")
|
||||
order = np.argsort(-counts_np)
|
||||
colors_np = colors_np[order]
|
||||
disc_color = tuple(int(v) for v in colors_np[0].tolist())
|
||||
cup_color = (
|
||||
tuple(int(v) for v in colors_np[1].tolist())
|
||||
if colors_np.shape[0] > 1
|
||||
else None
|
||||
)
|
||||
|
||||
disc_mask = np.zeros((h, w), dtype=np.uint8)
|
||||
cup_mask = np.zeros((h, w), dtype=np.uint8)
|
||||
|
||||
if disc_color is not None:
|
||||
disc_mask[np.all(arr == disc_color, axis=-1)] = 1
|
||||
if cup_color is not None:
|
||||
cup_mask[np.all(arr == cup_color, axis=-1)] = 1
|
||||
|
||||
return disc_mask, cup_mask if cup_mask.any() else None, (w, h)
|
||||
|
||||
def load_contour_from_file(self, contour_path: Path) -> np.ndarray:
|
||||
# Fast path: contour files are typically CSV or whitespace-delimited x,y pairs.
|
||||
try:
|
||||
arr = np.loadtxt(str(contour_path), delimiter=",", comments="#", dtype=np.float32)
|
||||
except Exception:
|
||||
try:
|
||||
arr = np.loadtxt(str(contour_path), comments="#", dtype=np.float32)
|
||||
except Exception:
|
||||
return np.zeros((0, 2), dtype=np.float32)
|
||||
if arr.size == 0:
|
||||
return np.zeros((0, 2), dtype=np.float32)
|
||||
if arr.ndim == 1:
|
||||
if arr.shape[0] < 2:
|
||||
return np.zeros((0, 2), dtype=np.float32)
|
||||
arr = arr.reshape(1, -1)
|
||||
if arr.shape[1] < 2:
|
||||
return np.zeros((0, 2), dtype=np.float32)
|
||||
return arr[:, :2].astype(np.float32, copy=False)
|
||||
|
||||
def coords_to_mask(
|
||||
self,
|
||||
coords: Optional[np.ndarray],
|
||||
size: Tuple[int, int],
|
||||
) -> np.ndarray:
|
||||
if coords is None or len(coords) == 0:
|
||||
return np.zeros((self.target_size, self.target_size), dtype=np.float32)
|
||||
|
||||
width, height = map(int, size)
|
||||
target_shape = (height, width)
|
||||
arr = np.asarray(coords)
|
||||
if arr.size == 0:
|
||||
return np.zeros((self.target_size, self.target_size), dtype=np.float32)
|
||||
|
||||
if arr.ndim == 2 and arr.shape[-1] != 2:
|
||||
mask = (arr > 0).astype(np.uint8)
|
||||
return self._resize_mask(mask)
|
||||
|
||||
if arr.ndim > 2:
|
||||
arr = arr.reshape(-1, arr.shape[-1])
|
||||
arr = arr.astype(float, copy=False)
|
||||
if arr.shape[-1] != 2:
|
||||
raise ValueError(f"Expected coordinate pairs, got shape {arr.shape}")
|
||||
|
||||
points = [tuple(map(float, pt)) for pt in arr]
|
||||
if len(points) < 3:
|
||||
return np.zeros(target_shape, dtype=np.float32)
|
||||
|
||||
img = Image.new("L", size, 0)
|
||||
draw = ImageDraw.Draw(img)
|
||||
draw.polygon(points, outline=1, fill=1)
|
||||
mask = np.array(img, dtype=np.uint8)
|
||||
return self._resize_mask(mask)
|
||||
|
||||
def _resize_mask(self, mask: np.ndarray) -> np.ndarray:
|
||||
img = Image.fromarray((mask > 0).astype(np.uint8) * 255)
|
||||
img = img.resize((self.target_size, self.target_size), Resampling.NEAREST)
|
||||
return (np.array(img, dtype=np.uint8) > 0).astype(np.float32)
|
||||
|
||||
def load_masks(self, entry: ManifestEntry) -> Tuple[np.ndarray, np.ndarray]:
|
||||
key = self._entry_cache_key(entry)
|
||||
if self.in_memory_cache:
|
||||
cached = self._mem_mask_cache.get(key)
|
||||
if cached is not None:
|
||||
disc_u8, cup_u8 = cached
|
||||
return disc_u8.astype(np.float32), cup_u8.astype(np.float32)
|
||||
cache_path = self._mask_cache_path(entry)
|
||||
if cache_path and cache_path.exists():
|
||||
cached = self._load_mask_cache(cache_path)
|
||||
if cached is not None:
|
||||
if self.in_memory_cache:
|
||||
disc, cup = cached
|
||||
self._mem_mask_cache[key] = (
|
||||
disc.astype(np.uint8),
|
||||
cup.astype(np.uint8),
|
||||
)
|
||||
return cached
|
||||
|
||||
image = Image.open(entry.image_path)
|
||||
size = image.size
|
||||
|
||||
disc_coords = cup_coords = None
|
||||
if entry.annotation_type_disc == "mask":
|
||||
disc_coords, cup_coords_from_disc, size = self.extract_masks_from_image(
|
||||
entry.annotation_disc
|
||||
)
|
||||
if cup_coords_from_disc is not None:
|
||||
cup_coords = cup_coords_from_disc
|
||||
else:
|
||||
disc_coords = self.load_contour_from_file(entry.annotation_disc)
|
||||
|
||||
if entry.annotation_type_cup == "mask":
|
||||
_, cup_coords_from_cup, size_cup = self.extract_masks_from_image(
|
||||
entry.annotation_cup
|
||||
)
|
||||
if cup_coords_from_cup is not None:
|
||||
cup_coords = cup_coords_from_cup
|
||||
if disc_coords is None:
|
||||
disc_coords, _, size = self.extract_masks_from_image(
|
||||
entry.annotation_cup
|
||||
)
|
||||
else:
|
||||
size = size_cup
|
||||
else:
|
||||
cup_coords = self.load_contour_from_file(entry.annotation_cup)
|
||||
|
||||
disc_mask = self.coords_to_mask(disc_coords, size).astype(np.float32)
|
||||
cup_mask = self.coords_to_mask(cup_coords, size).astype(np.float32)
|
||||
if self.in_memory_cache:
|
||||
self._mem_mask_cache[key] = (
|
||||
disc_mask.astype(np.uint8),
|
||||
cup_mask.astype(np.uint8),
|
||||
)
|
||||
self._save_mask_cache(cache_path, disc_mask, cup_mask)
|
||||
return disc_mask, cup_mask
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def build_loaders(self, batch_size: int = 4, num_workers: int = 0) -> Tuple[DataLoader, DataLoader]:
|
||||
train_ds = SegmentationDataset(self.train_entries, self, augment=True)
|
||||
val_ds = SegmentationDataset(self.val_entries, self, augment=False)
|
||||
train_loader = DataLoader(
|
||||
train_ds, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True
|
||||
)
|
||||
val_loader = DataLoader(
|
||||
val_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True
|
||||
)
|
||||
return train_loader, val_loader
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def dice_score(self, preds: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
||||
preds = (preds > 0.5).float()
|
||||
intersection = (preds * targets).sum(dim=(2, 3))
|
||||
union = preds.sum(dim=(2, 3)) + targets.sum(dim=(2, 3))
|
||||
dice = (2 * intersection + 1e-6) / (union + 1e-6)
|
||||
return dice.mean(dim=0)
|
||||
|
||||
def train(
|
||||
self,
|
||||
epochs: int = 40,
|
||||
batch_size: int = 4,
|
||||
lr: float = 1e-3,
|
||||
weight_decay: float = 1e-5,
|
||||
checkpoint_dir: Path = Path("models/unet_segmenter"),
|
||||
) -> None:
|
||||
print(
|
||||
f"[UNetSegmenter] training on device={self.device} "
|
||||
f"(epochs={epochs}, batch_size={batch_size}, workers={self.loader_workers})"
|
||||
)
|
||||
train_loader, val_loader = self.build_loaders(batch_size=batch_size, num_workers=self.loader_workers)
|
||||
optimizer = torch.optim.Adam(
|
||||
self.model.parameters(), lr=lr, weight_decay=weight_decay
|
||||
)
|
||||
criterion = nn.BCEWithLogitsLoss()
|
||||
best_dice = -math.inf
|
||||
checkpoint_dir.mkdir(parents=True, exist_ok=True)
|
||||
best_path = checkpoint_dir / "best.pt"
|
||||
|
||||
epoch_bar = tqdm(range(1, epochs + 1), desc="Epochs", unit="epoch")
|
||||
|
||||
for epoch in epoch_bar:
|
||||
self.model.train()
|
||||
batch_bar = tqdm(
|
||||
train_loader,
|
||||
desc=f"Train {epoch}/{epochs}",
|
||||
leave=False,
|
||||
unit="batch",
|
||||
total=len(train_loader),
|
||||
)
|
||||
train_loss_total = 0.0
|
||||
train_samples = 0
|
||||
for images, masks in batch_bar:
|
||||
images = images.to(self.device)
|
||||
masks = masks.to(self.device)
|
||||
optimizer.zero_grad()
|
||||
logits = self.model(images)
|
||||
loss_disc = criterion(logits[:, 0:1], masks[:, 0:1])
|
||||
loss_cup = criterion(logits[:, 1:2], masks[:, 1:2])
|
||||
loss = self.disc_weight * loss_disc + self.cup_weight * loss_cup
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
batch_size = images.size(0)
|
||||
train_loss_total += loss.item() * batch_size
|
||||
train_samples += batch_size
|
||||
|
||||
train_loss = (
|
||||
train_loss_total / train_samples if train_samples else float("nan")
|
||||
)
|
||||
|
||||
self.model.eval()
|
||||
dices = []
|
||||
val_bar = tqdm(
|
||||
val_loader,
|
||||
desc="Validate",
|
||||
leave=False,
|
||||
unit="batch",
|
||||
total=len(val_loader),
|
||||
)
|
||||
with torch.no_grad():
|
||||
for images, masks in val_bar:
|
||||
images = images.to(self.device)
|
||||
masks = masks.to(self.device)
|
||||
logits = self.model(images)
|
||||
probs = torch.sigmoid(logits)
|
||||
dice = self.dice_score(probs, masks)
|
||||
dices.append(dice.cpu())
|
||||
if dices:
|
||||
mean_dice = torch.stack(dices).mean(dim=0)
|
||||
disc_dice = mean_dice[0].item()
|
||||
cup_dice = mean_dice[1].item()
|
||||
weight_sum = self.disc_weight + self.cup_weight
|
||||
score = (
|
||||
(self.disc_weight * disc_dice + self.cup_weight * cup_dice)
|
||||
/ weight_sum
|
||||
if weight_sum
|
||||
else 0.0
|
||||
)
|
||||
epoch_bar.set_postfix(
|
||||
loss=f"{train_loss:.4f}",
|
||||
dice_disc=f"{disc_dice:.3f}",
|
||||
dice_cup=f"{cup_dice:.3f}",
|
||||
dice_w=f"{score:.3f}",
|
||||
)
|
||||
else:
|
||||
disc_dice = cup_dice = 0.0
|
||||
score = 0.0
|
||||
epoch_bar.set_postfix(loss=f"{train_loss:.4f}")
|
||||
|
||||
if score > best_dice:
|
||||
best_dice = score
|
||||
torch.save({"model": self.model.state_dict()}, best_path)
|
||||
|
||||
if best_path.exists():
|
||||
state = torch.load(best_path, map_location=self.device)
|
||||
self.model.load_state_dict(state["model"])
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def evaluate_holdout(
|
||||
self, output_dir: Path = Path("analysis_data/segmenter_eval")
|
||||
) -> pd.DataFrame:
|
||||
return self.evaluate_dataset(split_filter={"holdout"}, output_dir=output_dir)
|
||||
|
||||
@staticmethod
|
||||
def overlay_masks(
|
||||
image: Image.Image, disc: np.ndarray, cup: np.ndarray
|
||||
) -> Image.Image:
|
||||
overlay = image.copy()
|
||||
disc_img = Image.fromarray((disc * 255).astype(np.uint8))
|
||||
cup_img = Image.fromarray((cup * 255).astype(np.uint8))
|
||||
disc_color = Image.new("RGBA", image.size, (255, 0, 0, 0))
|
||||
cup_color = Image.new("RGBA", image.size, (0, 255, 0, 0))
|
||||
disc_color.paste((255, 0, 0, 100), mask=disc_img)
|
||||
cup_color.paste((0, 255, 0, 100), mask=cup_img)
|
||||
overlay = overlay.convert("RGBA")
|
||||
overlay = Image.alpha_composite(overlay, disc_color)
|
||||
overlay = Image.alpha_composite(overlay, cup_color)
|
||||
return overlay.convert("RGB")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
@staticmethod
|
||||
def _normalize_filter(values: Optional[Iterable[str]]) -> Optional[Set[str]]:
|
||||
if values is None:
|
||||
return None
|
||||
if isinstance(values, str):
|
||||
return {values}
|
||||
return {str(item) for item in values}
|
||||
|
||||
@staticmethod
|
||||
def _dice_from_masks(pred: np.ndarray, target: np.ndarray) -> float:
|
||||
pred = (pred > 0).astype(np.float32)
|
||||
target = (target > 0).astype(np.float32)
|
||||
intersection = float((pred * target).sum())
|
||||
denom = float(pred.sum() + target.sum())
|
||||
return (2.0 * intersection + 1e-6) / (denom + 1e-6)
|
||||
|
||||
def get_entries(
|
||||
self,
|
||||
dataset_filter: Optional[Iterable[str]] = None,
|
||||
split_filter: Optional[Iterable[str]] = None,
|
||||
) -> List[ManifestEntry]:
|
||||
dataset_set = self._normalize_filter(dataset_filter)
|
||||
split_set = self._normalize_filter(split_filter)
|
||||
entries = self._manifest
|
||||
if dataset_set is not None:
|
||||
entries = [e for e in entries if e.dataset in dataset_set]
|
||||
if split_set is not None:
|
||||
entries = [e for e in entries if e.split in split_set]
|
||||
return list(entries)
|
||||
|
||||
def evaluate_dataset(
|
||||
self,
|
||||
dataset_filter: Optional[Iterable[str]] = None,
|
||||
split_filter: Optional[Iterable[str]] = None,
|
||||
output_dir: Path = Path("analysis_data/segmenter_eval"),
|
||||
save_overlays: bool = True,
|
||||
metrics_path: Optional[Path] = None,
|
||||
threshold: float = 0.5,
|
||||
tta: bool = False,
|
||||
) -> pd.DataFrame:
|
||||
entries = self.get_entries(
|
||||
dataset_filter=dataset_filter, split_filter=split_filter
|
||||
)
|
||||
if not entries:
|
||||
return pd.DataFrame(
|
||||
columns=[
|
||||
"sample_id",
|
||||
"dataset",
|
||||
"split",
|
||||
"dice_disc",
|
||||
"dice_cup",
|
||||
]
|
||||
)
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
if metrics_path is None:
|
||||
suffix_parts = []
|
||||
if dataset_filter is not None:
|
||||
suffix_parts.append("-".join(sorted(self._normalize_filter(dataset_filter))))
|
||||
if split_filter is not None:
|
||||
suffix_parts.append("-".join(sorted(self._normalize_filter(split_filter))))
|
||||
suffix = "_".join(part for part in suffix_parts if part)
|
||||
csv_name = f"metrics{'_' + suffix if suffix else ''}.csv"
|
||||
metrics_path = output_dir / csv_name
|
||||
|
||||
records = []
|
||||
self.model.eval()
|
||||
progress = tqdm(
|
||||
entries,
|
||||
desc="Evaluate",
|
||||
unit="sample",
|
||||
leave=False,
|
||||
)
|
||||
for entry in progress:
|
||||
orig_image = Image.open(entry.image_path).convert("RGB")
|
||||
image = self.preprocess_image(orig_image)
|
||||
tensor = transforms.ToTensor()(image)
|
||||
tensor = self._normalize_tensor(tensor)
|
||||
tensor = tensor.unsqueeze(0).to(self.device)
|
||||
with torch.no_grad():
|
||||
logits = self.model(tensor)
|
||||
if tta:
|
||||
t_h = torch.flip(tensor, dims=[3])
|
||||
log_h = self.model(t_h)
|
||||
log_h = torch.flip(log_h, dims=[3])
|
||||
t_v = torch.flip(tensor, dims=[2])
|
||||
log_v = self.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] > threshold).astype(np.uint8)
|
||||
cup_pred = (probs[1] > threshold).astype(np.uint8)
|
||||
# Structural prior: cup within disc
|
||||
cup_pred = (cup_pred > 0) & (disc_pred > 0)
|
||||
cup_pred = cup_pred.astype(np.uint8)
|
||||
|
||||
disc_gt, cup_gt = self.load_masks(entry)
|
||||
disc_gt = disc_gt.astype(np.uint8)
|
||||
cup_gt = cup_gt.astype(np.uint8)
|
||||
|
||||
dice_disc = self._dice_from_masks(disc_pred, disc_gt)
|
||||
dice_cup = self._dice_from_masks(cup_pred, cup_gt)
|
||||
|
||||
records.append(
|
||||
{
|
||||
"sample_id": entry.sample_id,
|
||||
"dataset": entry.dataset,
|
||||
"split": entry.split,
|
||||
"dice_disc": dice_disc,
|
||||
"dice_cup": dice_cup,
|
||||
}
|
||||
)
|
||||
|
||||
progress.set_postfix(
|
||||
dice_disc=f"{dice_disc:.3f}", dice_cup=f"{dice_cup:.3f}"
|
||||
)
|
||||
|
||||
if save_overlays:
|
||||
overlay_gt = self.overlay_masks(image, disc_gt, cup_gt)
|
||||
overlay_pred = self.overlay_masks(image, disc_pred, cup_pred)
|
||||
combined = Image.new("RGB", (image.width * 2, image.height))
|
||||
combined.paste(overlay_gt, (0, 0))
|
||||
combined.paste(overlay_pred, (image.width, 0))
|
||||
combined.save(output_dir / f"{entry.sample_id}_eval.png")
|
||||
|
||||
metrics_df = pd.DataFrame(records)
|
||||
summary = metrics_df[["dice_disc", "dice_cup"]].mean()
|
||||
summary_row = {
|
||||
"sample_id": "__mean__",
|
||||
"dataset": "summary",
|
||||
"split": "summary",
|
||||
"dice_disc": summary["dice_disc"],
|
||||
"dice_cup": summary["dice_cup"],
|
||||
}
|
||||
metrics_with_summary = pd.concat(
|
||||
[metrics_df, pd.DataFrame([summary_row])], ignore_index=True
|
||||
)
|
||||
metrics_with_summary.to_csv(metrics_path, index=False)
|
||||
return metrics_with_summary
|
||||
Reference in New Issue
Block a user