ryan/hypertower
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add distributed server implementation and protocol definitions

Browse Source 
- Introduced `protocol.py` for shared data models used in server/client communication, including request and response schemas for registration, job submission, and status updates.
- Implemented `server.py` to manage a SQLite job queue and client registry, handling job polling, status updates, and job completion.
- Created a cheat sheet for server usage, detailing commands for starting the server, submitting jobs, and monitoring clients.
- Added several experiment configuration files for various training setups, including geometry vector injections and baseline ensembles.
This commit is contained in:
rpotter6298
2026-04-28 08:24:25 +02:00
parent 4dea45df78
commit 512ebd13b2
42 changed files with 4468 additions and 612 deletions
Download Patch File Download Diff File Expand all files Collapse all files
v4/classes/profiles/fundus_images.py
+574
View File
@@ -0,0 +1,574 @@
"""fundus_images — disc/cup geometry for fundus image profiles.
Contains all fundus-specific geometry logic: mask parsing, feature computation,
and source-specific loaders. Profiles whose ImageDataView supports geometry
should implement build_geometry_loader(source, **kwargs) and/or
build_seg_map_loader(source, **kwargs) and delegate here.
Geometry vectors (5 scalar CDR features per eye)
-----------------------------------------------
build_geometry_loader("gt", contour_dir=...) → GTGeometryLoader
Seg maps (3-class disc/cup label map per eye, fed to a CNN tower)
-----------------------------------------------------------------
build_seg_map_loader("gt", contour_dir=..., ...) → GTSegMapLoader
build_seg_map_loader("unet", weights_path=..., ...) → UNetSegMapLoader
"""
from __future__ import annotations
import copy
from collections import Counter
from pathlib import Path
from typing import Iterable, Tuple
import numpy as np
import torch
from PIL import Image, ImageDraw
from PIL.Image import Resampling
from torch.utils.data import DataLoader, Dataset
EPS = 1e-6
_FEATURE_DIM = 5
_FEATURE_NAMES = ["area_cdr", "rim_ratio", "vertical_cdr", "horizontal_cdr", "centre_shift"]
_MASK_SIZE = (512, 512) # canonical rasterisation size; CDR ratios are scale-invariant
# ---------------------------------------------------------------------------
# Mask utilities
# ---------------------------------------------------------------------------
def disc_cup_from_mask_image(mask_img: Image.Image) -> Tuple[np.ndarray, np.ndarray]:
"""Return binary (disc, cup) masks from a REFUGE-style colour annotation image."""
arr = np.asarray(mask_img)
if arr.ndim == 3:
h, w, c = arr.shape
border = np.concatenate(
[arr[0, :, :], arr[-1, :, :], arr[:, 0, :], arr[:, -1, :]], axis=0
)
bg_color = Counter(map(tuple, border)).most_common(1)[0][0]
colors = Counter(map(tuple, arr.reshape(-1, c)))
colors.pop(bg_color, None)
disc = (~np.all(arr == bg_color, axis=-1)).astype(np.uint8)
if colors:
cup_color = min(colors.keys(), key=lambda col: sum(col))
cup = np.all(arr == cup_color, axis=-1).astype(np.uint8)
else:
cup = np.zeros((h, w), dtype=np.uint8)
else:
border = np.concatenate([arr[0, :], arr[-1, :], arr[:, 0], arr[:, -1]])
bg_value = Counter(border.tolist()).most_common(1)[0][0]
disc = (arr != bg_value).astype(np.uint8)
fg = arr[arr != bg_value]
cup = (arr == int(np.min(fg))).astype(np.uint8) if fg.size > 0 else np.zeros_like(arr)
cup = (cup > 0) & (disc > 0)
return disc.astype(np.uint8), cup.astype(np.uint8)
def _contour_to_mask(coords: np.ndarray, size: Tuple[int, int]) -> np.ndarray:
"""Rasterize a polygon contour (Nx2 xy array) into a binary mask of (width, height)."""
from PIL import ImageDraw
if coords is None or coords.size == 0:
return np.zeros((size[1], size[0]), dtype=np.uint8)
img = Image.new("L", size, 0)
draw = ImageDraw.Draw(img)
draw.polygon([tuple(map(float, pt)) for pt in coords], outline=1, fill=1)
return np.array(img, dtype=np.uint8)
# ---------------------------------------------------------------------------
# Feature computation
# ---------------------------------------------------------------------------
def compute_geometry_features(disc_mask: np.ndarray, cup_mask: np.ndarray) -> np.ndarray:
"""Compute 5 cup/disc structural descriptors from binary masks.
Returns float32 [area_cdr, rim_ratio, vertical_cdr, horizontal_cdr, centre_shift].
"""
disc = (disc_mask > 0).astype(np.float32)
cup = (cup_mask > 0).astype(np.float32)
disc_area = disc.sum()
cup_area = cup.sum()
area_cdr = cup_area / (disc_area + EPS)
rim_ratio = (disc_area - cup_area) / (disc_area + EPS)
disc_h = float(np.any(disc > 0, axis=1).sum())
cup_h = float(np.any(cup > 0, axis=1).sum())
disc_w = float(np.any(disc > 0, axis=0).sum())
cup_w = float(np.any(cup > 0, axis=0).sum())
vertical_cdr = cup_h / (disc_h + EPS)
horizontal_cdr = cup_w / (disc_w + EPS)
def _centre(m: np.ndarray) -> Tuple[float, float]:
coords = np.argwhere(m > 0)
if coords.size == 0:
return 0.5, 0.5
ys, xs = coords[:, 0], coords[:, 1]
return float(xs.mean()) / m.shape[1], float(ys.mean()) / m.shape[0]
dcx, dcy = _centre(disc)
ccx, ccy = _centre(cup)
centre_shift = float(np.hypot(ccx - dcx, ccy - dcy))
return np.array(
[area_cdr, rim_ratio, vertical_cdr, horizontal_cdr, centre_shift],
dtype=np.float32,
)
# ---------------------------------------------------------------------------
# Loaders
# ---------------------------------------------------------------------------
class GTGeometryLoader:
"""Pre-computes per-eye geometry vectors from PAPILA GT contour annotations.
File naming: RET{pid:03d}{eye}_{disc|cup}_exp{n}.txt
Averages exp1 and exp2 when both are present; zero vector for missing entries.
Usage:
loader = GTGeometryLoader(contour_dir)
loader.precompute(df, patient_col="Patient ID")
vecs = loader.all_vectors() # {(pid, eye): ndarray}
"""
_EXPERTS = (1, 2)
feature_dim = _FEATURE_DIM
feature_names = _FEATURE_NAMES
def __init__(self, contour_dir: str | Path) -> None:
self._dir = Path(contour_dir)
self._cache: dict[tuple, np.ndarray] = {}
def precompute(self, df, patient_col: str = "Patient ID") -> None:
n_ok = 0
for _, row in df.iterrows():
pid = int(row[patient_col])
eye = str(row.get("eyeID", "OD"))
key = (pid, eye)
if key in self._cache:
continue
vec = self._compute(pid, eye)
self._cache[key] = vec if vec is not None else np.zeros(self.feature_dim, dtype=np.float32)
if vec is not None:
n_ok += 1
print(f"[GTGeometryLoader] {n_ok}/{len(self._cache)} geometry vectors computed", flush=True)
def all_vectors(self) -> dict:
return dict(self._cache)
def _compute(self, pid: int, eye: str) -> "np.ndarray | None":
stem = f"RET{pid:03d}{eye}"
vecs: list[np.ndarray] = []
for exp in self._EXPERTS:
disc_path = self._dir / f"{stem}_disc_exp{exp}.txt"
cup_path = self._dir / f"{stem}_cup_exp{exp}.txt"
if not disc_path.exists():
continue
try:
disc_c = np.loadtxt(disc_path)
if disc_c.ndim == 1:
disc_c = disc_c.reshape(-1, 2)
disc_mask = _contour_to_mask(disc_c, _MASK_SIZE)
if cup_path.exists():
cup_c = np.loadtxt(cup_path)
if cup_c.ndim == 1:
cup_c = cup_c.reshape(-1, 2)
cup_mask = _contour_to_mask(cup_c, _MASK_SIZE)
else:
cup_mask = np.zeros((_MASK_SIZE[1], _MASK_SIZE[0]), dtype=np.uint8)
cup_mask = ((cup_mask > 0) & (disc_mask > 0)).astype(np.uint8)
vecs.append(compute_geometry_features(disc_mask, cup_mask))
except Exception:
continue
if not vecs:
return None
return np.stack(vecs).mean(axis=0).astype(np.float32)
# ---------------------------------------------------------------------------
# Seg-map utilities (shared by GT and UNet seg-map loaders)
# ---------------------------------------------------------------------------
def _combine_disc_cup(disc: np.ndarray, cup: np.ndarray) -> np.ndarray:
"""Merge binary disc + cup masks into a uint8 label map: 0=bg, 1=rim, 2=cup."""
disc = (disc > 0).astype(np.uint8)
cup = ((cup > 0) & (disc > 0)).astype(np.uint8)
return (disc + cup).astype(np.uint8)
def _crop_to_disc_bbox(seg_map: np.ndarray) -> np.ndarray:
"""Crop a label map tightly to the disc bounding box (anywhere seg_map > 0)."""
rows = np.any(seg_map > 0, axis=1)
cols = np.any(seg_map > 0, axis=0)
if not rows.any():
return seg_map
r0, r1 = int(np.argmax(rows)), int(len(rows) - 1 - np.argmax(rows[::-1]))
c0, c1 = int(np.argmax(cols)), int(len(cols) - 1 - np.argmax(cols[::-1]))
return seg_map[r0:r1 + 1, c0:c1 + 1]
def _seg_map_to_array(
seg_map: np.ndarray, channels: int, target_size: int
) -> np.ndarray:
"""Resize a {0,1,2} seg map and convert to a (C, H, W) float32 array.
channels=1 → (1, H, W) values in {0, 0.5, 1.0}
channels=3 → (3, H, W) one-hot [bg, rim, cup]
"""
pil = Image.fromarray(seg_map.astype(np.uint8), mode="L").resize(
(target_size, target_size), Resampling.NEAREST
)
arr = np.array(pil, dtype=np.uint8)
if channels == 1:
return (arr.astype(np.float32) / 2.0)[None, :, :]
if channels == 3:
return np.stack([
(arr == 0).astype(np.float32),
(arr == 1).astype(np.float32),
(arr == 2).astype(np.float32),
], axis=0)
raise ValueError(f"channels must be 1 or 3, got {channels}")
def _load_papila_contour(path: Path) -> np.ndarray:
"""Load (x, y) contour pairs from a PAPILA whitespace/comma-delimited text file."""
for delim in (",", None):
try:
arr = np.loadtxt(str(path), delimiter=delim, comments="#", dtype=np.float32)
if arr.size > 0 and arr.ndim >= 1:
if arr.ndim == 1:
arr = arr.reshape(-1, 2)
if arr.shape[1] >= 2:
return arr[:, :2]
except Exception:
continue
return np.zeros((0, 2), dtype=np.float32)
def _papila_disc_cup_masks(
pid: int, eye: str, contour_dir: Path, image_size: Tuple[int, int],
mask_size: int, experts: Tuple[int, ...] = (1, 2),
) -> Tuple[np.ndarray, np.ndarray] | None:
"""Average masks across PAPILA experts. Returns (disc, cup) at mask_size, or None."""
stem = f"RET{pid:03d}{eye}"
discs, cups = [], []
for exp in experts:
disc_path = contour_dir / f"{stem}_disc_exp{exp}.txt"
cup_path = contour_dir / f"{stem}_cup_exp{exp}.txt"
if not disc_path.exists():
continue
disc_c = _load_papila_contour(disc_path)
if len(disc_c) < 3:
continue
disc_m = _rasterise_polygon(disc_c, image_size, mask_size)
if cup_path.exists():
cup_c = _load_papila_contour(cup_path)
cup_m = (_rasterise_polygon(cup_c, image_size, mask_size)
if len(cup_c) >= 3 else np.zeros_like(disc_m))
else:
cup_m = np.zeros_like(disc_m)
discs.append(disc_m)
cups.append(cup_m)
if not discs:
return None
disc = (np.mean(discs, axis=0) > 0.5).astype(np.uint8)
cup = (np.mean(cups, axis=0) > 0.5).astype(np.uint8)
return disc, cup
def _rasterise_polygon(
coords: np.ndarray, image_size: Tuple[int, int], target_size: int
) -> np.ndarray:
"""Rasterise an (N, 2) polygon contour into a (target_size, target_size) binary mask.
image_size is (width, height) of the coord space (the original fundus image).
"""
if coords is None or len(coords) < 3:
return np.zeros((target_size, target_size), dtype=np.uint8)
img = Image.new("L", image_size, 0)
pts = [tuple(map(float, p)) for p in coords]
ImageDraw.Draw(img).polygon(pts, outline=1, fill=1)
img = img.resize((target_size, target_size), Resampling.NEAREST)
return (np.array(img, dtype=np.uint8) > 0).astype(np.uint8)
# ---------------------------------------------------------------------------
# Seg-map loaders
# ---------------------------------------------------------------------------
class GTSegMapLoader:
"""Pre-computes per-eye disc/cup seg maps from PAPILA GT contour annotations.
Output: dict {(pid, eye): np.ndarray (C, H, W) float32} cached for the fold.
"""
def __init__(
self,
contour_dir: str | Path,
*,
channels: int = 3,
mask_size: int = 512,
target_size: int = 224,
crop_to_disc: bool = True,
) -> None:
self._dir = Path(contour_dir)
self._channels = channels
self._mask_size = mask_size
self._target_size = target_size
self._crop = crop_to_disc
self._cache: dict[tuple, np.ndarray] = {}
@property
def cache_dim(self) -> tuple[int, int, int]:
return (self._channels, self._target_size, self._target_size)
def precompute(self, samples: Iterable[Tuple[int, str, Path]]) -> None:
"""samples: iterable of (pid, eye, image_path) tuples."""
n_ok = 0
blank = np.zeros((self._mask_size, self._mask_size), dtype=np.uint8)
for pid, eye, image_path in samples:
key = (pid, eye)
if key in self._cache:
continue
with Image.open(image_path) as _im:
im_size = _im.size # (W, H)
res = _papila_disc_cup_masks(
pid, eye, self._dir, im_size, self._mask_size,
)
if res is None:
seg = blank
else:
disc, cup = res
seg = _combine_disc_cup(disc, cup)
n_ok += 1
if self._crop:
seg = _crop_to_disc_bbox(seg)
self._cache[key] = _seg_map_to_array(seg, self._channels, self._target_size)
print(
f"[GTSegMapLoader] {n_ok}/{len(self._cache)} GT seg maps computed "
f"(channels={self._channels}, target={self._target_size})",
flush=True,
)
def all_seg_maps(self) -> dict:
return self._cache
class _UNetFTDataset(Dataset):
"""Pre-cached (image_tensor, mask_tensor) pairs for fine-tuning a UNet.
Decode + resize + normalize + GT mask rasterisation are all deterministic,
so we do them once at construction and store float32 tensors on CPU. This
drops per-batch cost to a tensor lookup + GPU transfer.
"""
def __init__(
self,
records: list, # list of (pid, eye, image_path)
contour_dir: Path,
segmenter, # UNetSegmenter — used for image preprocessing
):
import time
S = segmenter.target_size
self._imgs: list[torch.Tensor] = []
self._masks: list[torch.Tensor] = []
print(
f"[UNetSegMapLoader] pre-caching {len(records)} (image, mask) pairs "
f"at {S}×{S}...",
flush=True,
)
t0 = time.time()
report = max(1, len(records) // 4)
for i, (pid, eye, image_path) in enumerate(records, 1):
with Image.open(image_path) as raw:
im_size = raw.size
img_t = segmenter.preprocess(raw).detach().cpu()
res = _papila_disc_cup_masks(pid, eye, contour_dir, im_size, S)
if res is None:
disc = np.zeros((S, S), dtype=np.uint8)
cup = np.zeros_like(disc)
else:
disc, cup = res
mask_t = torch.from_numpy(np.stack([disc, cup], axis=0).astype(np.float32))
self._imgs.append(img_t)
self._masks.append(mask_t)
if i % report == 0 or i == len(records):
print(
f" [UNet ft cache] {i}/{len(records)} ({time.time() - t0:.1f}s)",
flush=True,
)
def __len__(self) -> int:
return len(self._imgs)
def __getitem__(self, idx: int):
return self._imgs[idx], self._masks[idx]
class UNetSegMapLoader:
"""Pre-computes per-eye seg maps via a REFUGE-pretrained UNet.
Optionally fine-tunes the UNet per fold on the training split's GT contours.
Output: dict {(pid, eye): np.ndarray (C, H, W) float32} cached for the fold.
"""
def __init__(
self,
weights_path: str | Path,
*,
contour_dir: str | Path,
channels: int = 3,
target_size: int = 224,
unet_size: int = 512,
normalize: str = "per_image",
threshold: float = 0.5,
crop_to_disc: bool = True,
finetune_epochs: int = 0,
finetune_lr: float = 1e-5,
finetune_batch_size: int = 4,
device: str | None = None,
) -> None:
from v4.classes.accessory.unet import UNetSegmenter
self._weights_path = Path(weights_path)
self._contour_dir = Path(contour_dir)
self._channels = channels
self._target_size = target_size
self._threshold = threshold
self._crop = crop_to_disc
self._ft_epochs = finetune_epochs
self._ft_lr = finetune_lr
self._ft_batch_size = finetune_batch_size
self._segmenter = UNetSegmenter(
target_size=unet_size, normalize=normalize, device=device,
).load_weights(self._weights_path)
self._base_state = copy.deepcopy(self._segmenter.model.state_dict())
self._cache: dict[tuple, np.ndarray] = {}
@property
def cache_dim(self) -> tuple[int, int, int]:
return (self._channels, self._target_size, self._target_size)
def reset_cache(self) -> None:
"""Clear cached seg maps (call between folds)."""
self._cache.clear()
def reset_weights(self) -> None:
"""Restore base REFUGE weights (undo any prior fine-tuning)."""
self._segmenter.model.load_state_dict(copy.deepcopy(self._base_state))
def finetune(self, train_samples: list) -> None:
"""Fine-tune the UNet on the training fold's GT contours.
train_samples: list of (pid, eye, image_path) tuples — train split only.
"""
if self._ft_epochs <= 0:
return
ds = _UNetFTDataset(train_samples, self._contour_dir, self._segmenter)
loader = DataLoader(
ds, batch_size=self._ft_batch_size, shuffle=True, num_workers=0,
)
print(
f"[UNetSegMapLoader] fine-tuning UNet for {self._ft_epochs} epochs "
f"on {len(train_samples)} samples (lr={self._ft_lr}, "
f"batch_size={self._ft_batch_size})",
flush=True,
)
self._segmenter.finetune(
loader, epochs=self._ft_epochs, lr=self._ft_lr,
log_prefix="[UNet ft]",
)
def precompute(self, samples: Iterable[Tuple[int, str, Path]]) -> None:
"""Run UNet inference on every sample and cache the resulting seg map."""
import time
samples = list(samples)
todo = [s for s in samples if (s[0], s[1]) not in self._cache]
if not todo:
return
print(
f"[UNetSegMapLoader] running UNet inference on {len(todo)} images...",
flush=True,
)
t0 = time.time()
report = max(1, len(todo) // 4)
for i, (pid, eye, image_path) in enumerate(todo, 1):
with Image.open(image_path) as raw:
disc, cup = self._segmenter.predict(raw, threshold=self._threshold)
seg = _combine_disc_cup(disc, cup)
if self._crop:
seg = _crop_to_disc_bbox(seg)
self._cache[(pid, eye)] = _seg_map_to_array(
seg, self._channels, self._target_size,
)
if i % report == 0 or i == len(todo):
print(
f" [UNet inf] {i}/{len(todo)} ({time.time() - t0:.1f}s)",
flush=True,
)
print(
f"[UNetSegMapLoader] {len(todo)} seg maps cached via UNet "
f"(channels={self._channels}, target={self._target_size})",
flush=True,
)
def all_seg_maps(self) -> dict:
return self._cache
# ---------------------------------------------------------------------------
# Factories
# ---------------------------------------------------------------------------
def build_geometry_loader(source: str, **kwargs):
"""Return the appropriate geometry-vector loader for the given source string.
Parameters
----------
source : "gt" | "unet"
contour_dir : (gt) path to contour annotation directory
"""
if source == "gt":
contour_dir = kwargs.get("contour_dir")
if contour_dir is None:
raise ValueError("build_geometry_loader source='gt' requires contour_dir")
return GTGeometryLoader(contour_dir)
raise NotImplementedError(f"build_geometry_loader: source={source!r} not implemented")
def build_seg_map_loader(source: str, **kwargs):
"""Return the appropriate seg-map loader for the given source string.
Parameters
----------
source : "gt" | "unet"
GT kwargs:
contour_dir, channels=3, mask_size=512, target_size=224, crop_to_disc=True
UNet kwargs:
weights_path, contour_dir, channels=3, target_size=224, unet_size=512,
normalize="per_image", threshold=0.5, crop_to_disc=True,
finetune_epochs=0, finetune_lr=1e-5, finetune_batch_size=4, device=None
"""
if source == "gt":
if "contour_dir" not in kwargs:
raise ValueError("build_seg_map_loader source='gt' requires contour_dir")
return GTSegMapLoader(**kwargs)
if source == "unet":
if "weights_path" not in kwargs:
raise ValueError("build_seg_map_loader source='unet' requires weights_path")
if "contour_dir" not in kwargs:
raise ValueError(
"build_seg_map_loader source='unet' requires contour_dir "
"(needed for per-fold fine-tuning, even if finetune_epochs=0)"
)
return UNetSegMapLoader(**kwargs)
raise NotImplementedError(f"build_seg_map_loader: source={source!r} not implemented")