Refactor geometry feature loaders and enhance distributed client status reporting
- Updated GTGeometryLoader to streamline geometry vector computation and caching. - Introduced UNetGeometryLoader for UNet-derived geometry vectors. - Added sample collection method in PapilaBundle for better data handling. - Refined GeometrySegEncoder and ImageEncoder to utilize new sample collection. - Enhanced distributed client with heartbeat mechanism for improved job tracking. - Added new configuration files for UNet-derived geometry integration.
This commit is contained in:
rpotter6298
@@ -124,71 +124,47 @@ def compute_geometry_features(disc_mask: np.ndarray, cup_mask: np.ndarray) -> np
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class GTGeometryLoader:
|
||||
"""Pre-computes per-eye geometry vectors from PAPILA GT contour annotations.
|
||||
"""Per-eye 5-feature CDR 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}
|
||||
Loads contour coordinates from per-expert text files, rasterises them at
|
||||
the original image's pixel space (so polygons aren't clipped), then
|
||||
computes the CDR feature vector. Expert masks are merged before feature
|
||||
computation to match the seg-map path.
|
||||
"""
|
||||
|
||||
_EXPERTS = (1, 2)
|
||||
|
||||
feature_dim = _FEATURE_DIM
|
||||
feature_names = _FEATURE_NAMES
|
||||
|
||||
def __init__(self, contour_dir: str | Path) -> None:
|
||||
self._dir = Path(contour_dir)
|
||||
def __init__(self, contour_dir: str | Path, *, mask_size: int = _MASK_SIZE[0]) -> None:
|
||||
self._dir = Path(contour_dir)
|
||||
self._mask_size = mask_size
|
||||
self._cache: dict[tuple, np.ndarray] = {}
|
||||
|
||||
def precompute(self, df, patient_col: str = "Patient ID") -> None:
|
||||
def reset_cache(self) -> None:
|
||||
self._cache.clear()
|
||||
|
||||
def precompute(self, samples: Iterable[Tuple[int, str, Path]]) -> None:
|
||||
n_ok = 0
|
||||
for _, row in df.iterrows():
|
||||
pid = int(row[patient_col])
|
||||
eye = str(row.get("eyeID", "OD"))
|
||||
for pid, eye, image_path in samples:
|
||||
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:
|
||||
with Image.open(image_path) as img:
|
||||
image_size = img.size # (W, H)
|
||||
res = _papila_disc_cup_masks(
|
||||
pid, eye, self._dir, image_size, self._mask_size,
|
||||
)
|
||||
if res is None:
|
||||
self._cache[key] = np.zeros(self.feature_dim, dtype=np.float32)
|
||||
else:
|
||||
disc, cup = res
|
||||
self._cache[key] = compute_geometry_features(disc, cup)
|
||||
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)
|
||||
@@ -411,12 +387,103 @@ class _UNetFTDataset(Dataset):
|
||||
return self._imgs[idx], self._masks[idx]
|
||||
|
||||
|
||||
class _PapilaUNetMaskPipeline:
|
||||
"""Per-fold UNet mask producer.
|
||||
|
||||
Owns a UNetSegmenter, handles the per-fold lifecycle:
|
||||
- reset_weights() restores REFUGE base state (call at start of each fold)
|
||||
- finetune(train_samples) fine-tunes on the train split's GT contours
|
||||
- predict(samples) returns {(pid, eye): (disc_mask, cup_mask)} raw masks
|
||||
|
||||
Downstream loaders interpret the raw masks differently — seg-map loader
|
||||
crops/resizes/encodes for a CNN, geometry-vector loader computes 5 CDR
|
||||
features. Sharing this pipeline avoids duplicating UNet load + finetune
|
||||
when both seg-map and feature-vector outputs are needed in one config.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
weights_path: str | Path,
|
||||
*,
|
||||
contour_dir: str | Path,
|
||||
unet_size: int = 512,
|
||||
normalize: str = "per_image",
|
||||
threshold: float = 0.5,
|
||||
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._contour_dir = Path(contour_dir)
|
||||
self._threshold = threshold
|
||||
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(Path(weights_path))
|
||||
self._base_state = copy.deepcopy(self._segmenter.model.state_dict())
|
||||
|
||||
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."""
|
||||
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"[UNetMaskPipeline] 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 predict(
|
||||
self, samples: Iterable[Tuple[int, str, Path]],
|
||||
) -> dict[tuple, Tuple[np.ndarray, np.ndarray]]:
|
||||
"""Run inference. Returns {(pid, eye): (disc_mask, cup_mask)} raw uint8."""
|
||||
import time
|
||||
samples = list(samples)
|
||||
if not samples:
|
||||
return {}
|
||||
print(
|
||||
f"[UNetMaskPipeline] running UNet inference on {len(samples)} images...",
|
||||
flush=True,
|
||||
)
|
||||
result: dict[tuple, Tuple[np.ndarray, np.ndarray]] = {}
|
||||
t0 = time.time()
|
||||
report = max(1, len(samples) // 4)
|
||||
for i, (pid, eye, image_path) in enumerate(samples, 1):
|
||||
with Image.open(image_path) as raw:
|
||||
disc, cup = self._segmenter.predict(raw, threshold=self._threshold)
|
||||
result[(pid, eye)] = (disc, cup)
|
||||
if i % report == 0 or i == len(samples):
|
||||
print(
|
||||
f" [UNet inf] {i}/{len(samples)} ({time.time() - t0:.1f}s)",
|
||||
flush=True,
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
class UNetSegMapLoader:
|
||||
"""Pre-computes per-eye seg maps via a REFUGE-pretrained UNet.
|
||||
"""Per-eye CNN-ready seg maps via a REFUGE-pretrained UNet.
|
||||
|
||||
Optionally fine-tunes the UNet per fold on the training split's GT contours.
|
||||
Wraps a `_PapilaUNetMaskPipeline` and post-processes raw masks into
|
||||
(C, H, W) float32 arrays sized for a downstream CNN.
|
||||
|
||||
Output: dict {(pid, eye): np.ndarray (C, H, W) float32} cached for the fold.
|
||||
Output: dict {(pid, eye): np.ndarray (C, H, W) float32}
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -435,22 +502,20 @@ class UNetSegMapLoader:
|
||||
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._pipeline = _PapilaUNetMaskPipeline(
|
||||
weights_path,
|
||||
contour_dir=contour_dir,
|
||||
unet_size=unet_size,
|
||||
normalize=normalize,
|
||||
threshold=threshold,
|
||||
finetune_epochs=finetune_epochs,
|
||||
finetune_lr=finetune_lr,
|
||||
finetune_batch_size=finetune_batch_size,
|
||||
device=device,
|
||||
)
|
||||
self._channels = channels
|
||||
self._target_size = target_size
|
||||
self._crop = crop_to_disc
|
||||
self._cache: dict[tuple, np.ndarray] = {}
|
||||
|
||||
@property
|
||||
@@ -458,89 +523,130 @@ class UNetSegMapLoader:
|
||||
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))
|
||||
self._pipeline.reset_weights()
|
||||
|
||||
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]",
|
||||
)
|
||||
self._pipeline.finetune(train_samples)
|
||||
|
||||
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)
|
||||
todo = [s for s in samples if (s[0], s[1]) not in self._cache]
|
||||
masks = self._pipeline.predict(todo)
|
||||
for (pid, eye), (disc, cup) in masks.items():
|
||||
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,
|
||||
)
|
||||
if masks:
|
||||
print(
|
||||
f"[UNetSegMapLoader] {len(masks)} seg maps cached "
|
||||
f"(channels={self._channels}, target={self._target_size})",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
def all_seg_maps(self) -> dict:
|
||||
return self._cache
|
||||
|
||||
|
||||
class UNetGeometryLoader:
|
||||
"""Per-eye 5-feature CDR vectors derived from UNet-predicted masks.
|
||||
|
||||
Same UNet lifecycle as `UNetSegMapLoader` but the output is a 5-vector
|
||||
(compute_geometry_features over the predicted disc/cup masks) rather
|
||||
than a CNN-ready seg map. Designed to slot into ImageEncoder's
|
||||
geometry_source mechanism for vector-style geometry injection.
|
||||
"""
|
||||
|
||||
feature_dim = _FEATURE_DIM
|
||||
feature_names = _FEATURE_NAMES
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
weights_path: str | Path,
|
||||
*,
|
||||
contour_dir: str | Path,
|
||||
unet_size: int = 512,
|
||||
normalize: str = "per_image",
|
||||
threshold: float = 0.5,
|
||||
finetune_epochs: int = 0,
|
||||
finetune_lr: float = 1e-5,
|
||||
finetune_batch_size: int = 4,
|
||||
device: str | None = None,
|
||||
) -> None:
|
||||
self._pipeline = _PapilaUNetMaskPipeline(
|
||||
weights_path,
|
||||
contour_dir=contour_dir,
|
||||
unet_size=unet_size,
|
||||
normalize=normalize,
|
||||
threshold=threshold,
|
||||
finetune_epochs=finetune_epochs,
|
||||
finetune_lr=finetune_lr,
|
||||
finetune_batch_size=finetune_batch_size,
|
||||
device=device,
|
||||
)
|
||||
self._cache: dict[tuple, np.ndarray] = {}
|
||||
|
||||
def reset_cache(self) -> None:
|
||||
self._cache.clear()
|
||||
|
||||
def reset_weights(self) -> None:
|
||||
self._pipeline.reset_weights()
|
||||
|
||||
def finetune(self, train_samples: list) -> None:
|
||||
self._pipeline.finetune(train_samples)
|
||||
|
||||
def precompute(self, samples: Iterable[Tuple[int, str, Path]]) -> None:
|
||||
todo = [s for s in samples if (s[0], s[1]) not in self._cache]
|
||||
masks = self._pipeline.predict(todo)
|
||||
for (pid, eye), (disc, cup) in masks.items():
|
||||
self._cache[(pid, eye)] = compute_geometry_features(disc, cup)
|
||||
if masks:
|
||||
print(
|
||||
f"[UNetGeometryLoader] {len(masks)} geometry vectors cached "
|
||||
f"(dim={self.feature_dim})",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
def all_vectors(self) -> dict:
|
||||
return dict(self._cache)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Factories
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def build_geometry_loader(source: str, **kwargs):
|
||||
"""Return the appropriate geometry-vector loader for the given source string.
|
||||
"""Return the appropriate geometry-vector loader for the given source.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
source : "gt" | "unet"
|
||||
contour_dir : (gt) path to contour annotation directory
|
||||
source : "gt" | "unet"
|
||||
|
||||
GT kwargs:
|
||||
contour_dir
|
||||
UNet kwargs:
|
||||
weights_path, contour_dir, unet_size=512, normalize="per_image",
|
||||
threshold=0.5, finetune_epochs=0, finetune_lr=1e-5,
|
||||
finetune_batch_size=4, device=None
|
||||
"""
|
||||
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)
|
||||
if source == "unet":
|
||||
if "weights_path" not in kwargs:
|
||||
raise ValueError("build_geometry_loader source='unet' requires weights_path")
|
||||
if "contour_dir" not in kwargs:
|
||||
raise ValueError(
|
||||
"build_geometry_loader source='unet' requires contour_dir "
|
||||
"(needed for per-fold fine-tuning, even if finetune_epochs=0)"
|
||||
)
|
||||
return UNetGeometryLoader(**kwargs)
|
||||
raise NotImplementedError(f"build_geometry_loader: source={source!r} not implemented")
|
||||
|
||||
|
||||
|
||||
@@ -507,6 +507,25 @@ class PapilaBundle:
|
||||
"""
|
||||
return [self._bundle.patient_col, "eyeID"]
|
||||
|
||||
# ── Sample collection (for tower early_pass) ─────────────────────────────
|
||||
|
||||
def collect_samples(self, df: pd.DataFrame | None) -> list[tuple]:
|
||||
"""Build (pid, eye, image_path) tuples from a split DataFrame.
|
||||
|
||||
Used by tower early_pass implementations that need per-eye image paths
|
||||
(UNet inference, contour rasterisation, etc.). Returns [] for an
|
||||
empty/None df.
|
||||
"""
|
||||
if df is None or len(df) == 0:
|
||||
return []
|
||||
pc = self._bundle.patient_col
|
||||
out: list[tuple] = []
|
||||
for _, row in df.iterrows():
|
||||
pid = int(row[pc])
|
||||
eye = str(row.get("eyeID", "OD"))
|
||||
out.append((pid, eye, self.image.get_image_path(pid, eye)))
|
||||
return out
|
||||
|
||||
# ── Backward-compat delegates ────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
|
||||
@@ -126,16 +126,14 @@ class GeometrySegEncoder(TowerBase):
|
||||
# ── EPC early_pass ───────────────────────────────────────────────────────
|
||||
|
||||
def early_pass(self, context) -> None:
|
||||
data = context.require("data")
|
||||
split = context.require("split")
|
||||
data = context.require("data")
|
||||
split = context.require("split")
|
||||
|
||||
train_samples = self._collect_samples(split.train, data)
|
||||
all_samples = self._collect_samples(split.train, data)
|
||||
all_samples += self._collect_samples(split.val, data)
|
||||
train_samples = data.collect_samples(split.train)
|
||||
all_samples = train_samples + data.collect_samples(split.val)
|
||||
if split.test is not None:
|
||||
all_samples += self._collect_samples(split.test, data)
|
||||
all_samples += data.collect_samples(split.test)
|
||||
|
||||
# Reset per-fold state if loader supports it (UNet only).
|
||||
if hasattr(self._loader, "reset_cache"):
|
||||
self._loader.reset_cache()
|
||||
if hasattr(self._loader, "reset_weights"):
|
||||
@@ -172,18 +170,6 @@ class GeometrySegEncoder(TowerBase):
|
||||
|
||||
# ── Internals ────────────────────────────────────────────────────────────
|
||||
|
||||
def _collect_samples(self, df, data) -> list:
|
||||
"""Build (pid, eye, image_path) tuples from a split DataFrame."""
|
||||
if df is None or len(df) == 0:
|
||||
return []
|
||||
pc = data.patient_col
|
||||
out = []
|
||||
for _, row in df.iterrows():
|
||||
pid = int(row[pc])
|
||||
eye = str(row.get("eyeID", "OD"))
|
||||
out.append((pid, eye, data.image.get_image_path(pid, eye)))
|
||||
return out
|
||||
|
||||
@staticmethod
|
||||
def _augment_array(arr: np.ndarray) -> np.ndarray:
|
||||
"""Random flip + 90° rotation on a (C, H, W) seg-map array."""
|
||||
|
||||
@@ -151,11 +151,12 @@ class ImageEncoder(TowerBase):
|
||||
|
||||
def early_pass(self, context) -> None:
|
||||
"""Per-fold setup: warm tensor cache (if enabled), publish geometry vectors."""
|
||||
data = context.require("data")
|
||||
data = context.require("data")
|
||||
split = context.require("split")
|
||||
|
||||
if self._cache_transformed:
|
||||
self._tensor_cache.clear()
|
||||
n = self._warm_tensor_cache(data, context.require("split"))
|
||||
n = self._warm_tensor_cache(data, split)
|
||||
print(
|
||||
f"[ImageEncoder] warmed transformed-tensor cache for {n} entries "
|
||||
f"({self._name})",
|
||||
@@ -164,7 +165,20 @@ class ImageEncoder(TowerBase):
|
||||
|
||||
if self._geom_loader is None:
|
||||
return
|
||||
self._geom_loader.precompute(data.df, patient_col=data.patient_col)
|
||||
|
||||
train_samples = data.collect_samples(split.train)
|
||||
all_samples = train_samples + data.collect_samples(split.val)
|
||||
if split.test is not None:
|
||||
all_samples += data.collect_samples(split.test)
|
||||
|
||||
if hasattr(self._geom_loader, "reset_cache"):
|
||||
self._geom_loader.reset_cache()
|
||||
if hasattr(self._geom_loader, "reset_weights"):
|
||||
self._geom_loader.reset_weights()
|
||||
if hasattr(self._geom_loader, "finetune"):
|
||||
self._geom_loader.finetune(train_samples)
|
||||
|
||||
self._geom_loader.precompute(all_samples)
|
||||
vecs = self._geom_loader.all_vectors()
|
||||
context.put(self.EPC_GEOMETRY_KEY, vecs)
|
||||
print(
|
||||
|
||||
@@ -0,0 +1,142 @@
|
||||
{
|
||||
"_notes": [
|
||||
"Geometry vector injection via UNet (no GT, no seg-CNN tower).",
|
||||
"img tower runs UNet per-fold to predict disc/cup masks, computes 5 CDR",
|
||||
"features per eye, publishes to EPC; cd tower consumes them.",
|
||||
"Same shape as ensemble_fused.json — just adds the UNet-derived geometry hook."
|
||||
],
|
||||
"run_name": "v4/ensemble_fused_geom_unet",
|
||||
"num_classes": 2,
|
||||
"label_filter": [0, 1],
|
||||
"split_identity_level": 1,
|
||||
"eval_stage": "hb",
|
||||
"save_predictions": true,
|
||||
"seed": 1234,
|
||||
"folds": 5,
|
||||
"fold_seed": 100,
|
||||
"output_root": "v4/results",
|
||||
"out_dir_tags": ["binary"],
|
||||
|
||||
"data": {
|
||||
"module": "v4.classes.profiles.v4papila",
|
||||
"args": {
|
||||
"image_dir": "Papila/FundusImages",
|
||||
"clinical_dir": "Papila/ClinicalData",
|
||||
"label_col": "Diagnosis",
|
||||
"iop_corr_method": "ratio",
|
||||
"iop_drop_raw": true,
|
||||
"exclude_cols": ["Axial_Length"],
|
||||
"in_memory_cache": true
|
||||
}
|
||||
},
|
||||
|
||||
"towers": [
|
||||
{
|
||||
"name": "img",
|
||||
"module": "v4.classes.towers.image_tower",
|
||||
"class": "ImageEncoder",
|
||||
"data_source": "image",
|
||||
"epc_supplies": ["geometry_vectors"],
|
||||
"args": {
|
||||
"backbone": "refugelike",
|
||||
"freeze_ratio": 0.0,
|
||||
"augment": true,
|
||||
"geometry_source": "unet",
|
||||
"weights_path": "models/v2/refuge/segmentation/per_image/best.pt",
|
||||
"contour_dir": "Papila/ExpertsSegmentations/Contours",
|
||||
"unet_size": 512,
|
||||
"normalize": "per_image",
|
||||
"threshold": 0.5,
|
||||
"finetune_epochs": 10,
|
||||
"finetune_lr": 1e-5,
|
||||
"finetune_batch_size": 4
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "cd",
|
||||
"module": "v4.classes.towers.clinical_tower",
|
||||
"class": "ClinicalEncoder",
|
||||
"data_source": "matrix",
|
||||
"epc_requests": ["geometry_vectors"],
|
||||
"args": {
|
||||
"hidden_dim": 128,
|
||||
"geom_dim": 5
|
||||
}
|
||||
}
|
||||
],
|
||||
|
||||
"stages": [
|
||||
{
|
||||
"name": "cd_warm",
|
||||
"type": "warm",
|
||||
"tower": "cd",
|
||||
"head_name": "cd_aux",
|
||||
"level": "eye",
|
||||
"epochs": 40
|
||||
},
|
||||
{
|
||||
"name": "img_aux",
|
||||
"type": "head",
|
||||
"input": "img",
|
||||
"train_with": "nt",
|
||||
"bcd": true
|
||||
},
|
||||
{
|
||||
"name": "cd_aux",
|
||||
"type": "head",
|
||||
"input": "cd",
|
||||
"train_with": "nt",
|
||||
"bcd": true
|
||||
},
|
||||
{
|
||||
"name": "nt",
|
||||
"type": "fusion",
|
||||
"module": "v4.classes.bridges.fusion_bridge",
|
||||
"class": "FusionBridge",
|
||||
"inputs": ["img", "cd"],
|
||||
"level": "eye",
|
||||
"epochs": 36,
|
||||
"train_towers": true,
|
||||
"warmup": {
|
||||
"tower_epochs": 3,
|
||||
"fused_epochs": 3
|
||||
},
|
||||
"args": {
|
||||
"fusion_dim": 256
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "nt_head",
|
||||
"type": "head",
|
||||
"input": "nt",
|
||||
"train_with": "nt"
|
||||
},
|
||||
{
|
||||
"name": "hb",
|
||||
"type": "fusion",
|
||||
"module": "v4.classes.bridges.hyperbridge",
|
||||
"class": "HyperBridge",
|
||||
"inputs": { "a": "nt", "b": "nt" },
|
||||
"level": "patient",
|
||||
"epochs": 10,
|
||||
"args": {
|
||||
"hidden_dim": 256,
|
||||
"mode": "embedding_mlp"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "hb_head",
|
||||
"type": "head",
|
||||
"input": "hb",
|
||||
"train_with": "hb",
|
||||
"args": { "dropout": 0.3 }
|
||||
}
|
||||
],
|
||||
|
||||
"training": {
|
||||
"lr": 1e-4,
|
||||
"batch_size": 8,
|
||||
"bcd_prob": 0.5,
|
||||
"tune_binary_threshold": true
|
||||
}
|
||||
}
|
||||
@@ -126,11 +126,12 @@ def _clients_table(api: _API) -> str:
|
||||
c["hostname"],
|
||||
c["gpu_info"][:30],
|
||||
s["state"],
|
||||
(s.get("job_id") or "-")[:12],
|
||||
s.get("run_name") or "-",
|
||||
prog,
|
||||
_ago(c["last_seen"]),
|
||||
])
|
||||
headers = ["ID", "HOST", "GPU", "STATE", "RUN", "PROGRESS", "SEEN"]
|
||||
headers = ["ID", "HOST", "GPU", "STATE", "JOB_ID", "RUN", "PROGRESS", "SEEN"]
|
||||
widths = [max(len(str(r[i])) for r in ([headers] + rows)) for i in range(len(headers))]
|
||||
sep = " "
|
||||
lines = []
|
||||
|
||||
+31
-11
@@ -232,6 +232,20 @@ def _run_job(job: JobSpec, server: _Server,
|
||||
log_dir.mkdir(parents=True, exist_ok=True)
|
||||
log_file = log_dir / f"job_{job.job_id}.log"
|
||||
ctx: dict = {}
|
||||
last_push = [time.time()] # mutable holder so _tail and _heartbeat share it
|
||||
|
||||
def _push():
|
||||
server.push_status(StatusPush(
|
||||
state="running",
|
||||
job_id=job.job_id,
|
||||
run_name=job.run_name,
|
||||
fold=ctx.get("fold"),
|
||||
stage=ctx.get("stage"),
|
||||
epoch=ctx.get("epoch"),
|
||||
total_epochs=ctx.get("total_epochs"),
|
||||
last_val_auc=ctx.get("last_val_auc"),
|
||||
))
|
||||
last_push[0] = time.time()
|
||||
|
||||
def _tail(path: Path):
|
||||
with open(path, "r") as f:
|
||||
@@ -242,16 +256,7 @@ def _run_job(job: JobSpec, server: _Server,
|
||||
info = _parse_line(raw)
|
||||
ctx.update(info)
|
||||
if "epoch" in info:
|
||||
server.push_status(StatusPush(
|
||||
state="running",
|
||||
job_id=job.job_id,
|
||||
run_name=job.run_name,
|
||||
fold=ctx.get("fold"),
|
||||
stage=ctx.get("stage"),
|
||||
epoch=ctx.get("epoch"),
|
||||
total_epochs=ctx.get("total_epochs"),
|
||||
last_val_auc=ctx.get("last_val_auc"),
|
||||
))
|
||||
_push()
|
||||
elif proc.poll() is not None:
|
||||
for raw in f:
|
||||
print(raw, end="", flush=True)
|
||||
@@ -259,6 +264,18 @@ def _run_job(job: JobSpec, server: _Server,
|
||||
else:
|
||||
time.sleep(0.05)
|
||||
|
||||
def _heartbeat():
|
||||
# Push a status update every ~30s even when no log line is parsed.
|
||||
# Prevents the server's reaper from declaring this client stale during
|
||||
# long deterministic blocks (UNet fine-tune, data load, etc.).
|
||||
while proc.poll() is None:
|
||||
time.sleep(5)
|
||||
if time.time() - last_push[0] >= 30:
|
||||
try:
|
||||
_push()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
with open(log_file, "w") as logf:
|
||||
proc = subprocess.Popen(
|
||||
cmd,
|
||||
@@ -268,10 +285,13 @@ def _run_job(job: JobSpec, server: _Server,
|
||||
start_new_session=True,
|
||||
)
|
||||
|
||||
tailer = threading.Thread(target=_tail, args=(log_file,), daemon=True)
|
||||
tailer = threading.Thread(target=_tail, args=(log_file,), daemon=True)
|
||||
heartbeat = threading.Thread(target=_heartbeat, daemon=True)
|
||||
tailer.start()
|
||||
heartbeat.start()
|
||||
proc.wait()
|
||||
tailer.join(timeout=5)
|
||||
heartbeat.join(timeout=5)
|
||||
log_file.unlink(missing_ok=True)
|
||||
|
||||
success = proc.returncode == 0
|
||||
|
||||
Binary file not shown.
@@ -48,13 +48,20 @@ _DB_PATH: Path = Path("v4/distributed/jobs.db")
|
||||
_REPO_ROOT: Path = Path.cwd()
|
||||
_CLIENT_TTL: int = 120 # seconds before a client is considered gone
|
||||
_MAX_ATTEMPTS: int = 3 # max times a job is retried before being left as failed
|
||||
_SERVER_START_TS: float = 0.0 # set in main(); used as a reaper grace window
|
||||
|
||||
_clients: dict[str, ClientInfo] = {}
|
||||
_clients_lock = threading.Lock()
|
||||
|
||||
|
||||
def _reap_stale_clients():
|
||||
"""Background thread: remove silent clients and re-queue their running jobs."""
|
||||
"""Background thread: remove silent clients and re-queue their running jobs.
|
||||
|
||||
On startup, the in-memory `_clients` dict is empty until clients re-register
|
||||
via the `please_reregister` mechanism. We skip the running-job re-queue pass
|
||||
for the first `_CLIENT_TTL` seconds after startup so still-alive clients have
|
||||
time to come back; otherwise the reaper would orphan their jobs.
|
||||
"""
|
||||
while True:
|
||||
time.sleep(30)
|
||||
cutoff = datetime.now(timezone.utc).timestamp() - _CLIENT_TTL
|
||||
@@ -73,6 +80,10 @@ def _reap_stale_clients():
|
||||
del _clients[cid]
|
||||
known_ids = set(_clients.keys())
|
||||
|
||||
in_grace = (time.time() - _SERVER_START_TS) < _CLIENT_TTL
|
||||
if in_grace:
|
||||
continue
|
||||
|
||||
with _db() as conn:
|
||||
rows = conn.execute(
|
||||
"SELECT job_id, assigned_to FROM jobs WHERE state='running'"
|
||||
@@ -454,12 +465,13 @@ def main():
|
||||
if not args.token:
|
||||
ap.error("--token is required (or set HT_TOKEN)")
|
||||
|
||||
global _TOKEN, _DB_PATH, _REPO_ROOT, _CLIENT_TTL, _MAX_ATTEMPTS
|
||||
global _TOKEN, _DB_PATH, _REPO_ROOT, _CLIENT_TTL, _MAX_ATTEMPTS, _SERVER_START_TS
|
||||
_TOKEN = args.token
|
||||
_DB_PATH = Path(args.db)
|
||||
_REPO_ROOT = Path(args.root).resolve() if args.root else Path.cwd()
|
||||
_CLIENT_TTL = args.client_ttl
|
||||
_MAX_ATTEMPTS = args.max_attempts
|
||||
_SERVER_START_TS = time.time()
|
||||
_init_db()
|
||||
|
||||
reaper = threading.Thread(target=_reap_stale_clients, daemon=True)
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
[
|
||||
{
|
||||
"_note": "GT-derived 5-feature CDR vector injected into cd via EPC. Re-run with bug fix (rasterisation now in original image space, not 512x512 canvas). 3 reps.",
|
||||
"run_name": "experiments/tri_v1/geom_vec_gt",
|
||||
"reps": 3
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,7 @@
|
||||
[
|
||||
{
|
||||
"_note": "UNet-derived 5-feature CDR vector injected into cd via EPC. 3 reps.",
|
||||
"run_name": "experiments/tri_v1/geom_vec_unet",
|
||||
"reps": 3
|
||||
}
|
||||
]
|
||||
Reference in New Issue
Block a user