rpotter6298
af813bbb62
- 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.
224 lines
8.2 KiB
Python
224 lines
8.2 KiB
Python
"""geometry_tower — GeometrySegEncoder for v4.
|
|
|
|
A CNN tower that takes a per-eye disc/cup *segmentation map* as input (rather
|
|
than the raw fundus image) and contributes its pooled embedding to fusion.
|
|
|
|
Seg maps are produced by an underlying loader (GT contour rasterisation or
|
|
UNet inference) during early_pass, then cached per fold.
|
|
|
|
UNet fine-tuning lives in early_pass too — the loader's `finetune(train_samples)`
|
|
call uses only the training split, then precompute() runs inference on all
|
|
fold samples (train + val + test).
|
|
|
|
Config example:
|
|
{
|
|
"name": "geom",
|
|
"module": "v4.classes.towers.geometry_tower",
|
|
"class": "GeometrySegEncoder",
|
|
"data_source": "image",
|
|
"args": {
|
|
"backbone": "resnet18",
|
|
"channels": 3,
|
|
"target_size": 224,
|
|
"augment": true,
|
|
"seg_source": "gt",
|
|
"contour_dir": "Papila/ExpertsSegmentations/Contours"
|
|
}
|
|
}
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torch import nn
|
|
from torchvision import models
|
|
|
|
_REPO_ROOT = Path(__file__).resolve().parents[3]
|
|
if str(_REPO_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(_REPO_ROOT))
|
|
|
|
from v4.classes.accessory.backbones import build_backbone
|
|
from v4.classes.towerbase import TowerBase
|
|
|
|
|
|
class GeometrySegEncoder(TowerBase):
|
|
"""CNN tower over disc/cup segmentation maps.
|
|
|
|
image_data : ImageDataView — provides get_image_path(*ids) and side_map.
|
|
Must implement build_seg_map_loader(source, **kwargs).
|
|
backbone : backbone key (see accessory/backbones.py)
|
|
channels : 1 (label map in [0,1]) or 3 (one-hot bg/rim/cup)
|
|
target_size : CNN input spatial size (cached arrays already at this size)
|
|
augment : random flip + 90° rotation at training time
|
|
freeze_ratio : fraction of early backbone blocks to freeze in [0, 1]
|
|
seg_source : passed to image_data.build_seg_map_loader (e.g. "gt", "unet")
|
|
**seg_kwargs : forwarded to build_seg_map_loader
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
image_data,
|
|
backbone: str = "resnet18",
|
|
channels: int = 3,
|
|
target_size: int = 224,
|
|
augment: bool = True,
|
|
freeze_ratio: float = 0.0,
|
|
seg_source: str = "gt",
|
|
**seg_kwargs: Any,
|
|
):
|
|
super().__init__()
|
|
self.image_data = image_data
|
|
self._channels = channels
|
|
self._target_size = target_size
|
|
self._augment = augment
|
|
|
|
if not hasattr(image_data, "build_seg_map_loader"):
|
|
raise TypeError(
|
|
f"GeometrySegEncoder requires image_data to implement "
|
|
f"build_seg_map_loader(), but {type(image_data).__name__} does not."
|
|
)
|
|
loader_kwargs = {
|
|
"channels": channels,
|
|
"target_size": target_size,
|
|
**seg_kwargs,
|
|
}
|
|
self._loader = image_data.build_seg_map_loader(seg_source, **loader_kwargs)
|
|
self._seg_cache: dict = {}
|
|
self._seg_source = seg_source
|
|
|
|
self.backbone, self._base_dim, self._blocks = build_backbone(backbone, freeze_ratio)
|
|
if channels != 3:
|
|
self._adapt_first_conv(channels)
|
|
|
|
print(
|
|
f"[GeometrySegEncoder] backbone={backbone} channels={channels} "
|
|
f"target_size={target_size} seg_source={seg_source}",
|
|
flush=True,
|
|
)
|
|
|
|
# ── TowerBase interface ──────────────────────────────────────────────────
|
|
|
|
@property
|
|
def out_dim(self) -> int:
|
|
return self._base_dim
|
|
|
|
@property
|
|
def _side_map(self) -> dict[str, str]:
|
|
return self.image_data.side_map
|
|
|
|
def _get(self, *ids) -> torch.Tensor:
|
|
key = tuple(ids)
|
|
arr = self._seg_cache.get(key)
|
|
if arr is None:
|
|
arr = np.zeros(
|
|
(self._channels, self._target_size, self._target_size),
|
|
dtype=np.float32,
|
|
)
|
|
if self.training and self._augment:
|
|
arr = self._augment_array(arr)
|
|
return torch.from_numpy(np.ascontiguousarray(arr))
|
|
|
|
# ── EPC early_pass ───────────────────────────────────────────────────────
|
|
|
|
def early_pass(self, context) -> None:
|
|
data = context.require("data")
|
|
split = context.require("split")
|
|
|
|
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._loader, "reset_cache"):
|
|
self._loader.reset_cache()
|
|
if hasattr(self._loader, "reset_weights"):
|
|
self._loader.reset_weights()
|
|
if hasattr(self._loader, "finetune"):
|
|
self._loader.finetune(train_samples)
|
|
|
|
self._loader.precompute(all_samples)
|
|
self._seg_cache = self._loader.all_seg_maps()
|
|
print(
|
|
f"[GeometrySegEncoder] cached {len(self._seg_cache)} seg maps for fold",
|
|
flush=True,
|
|
)
|
|
|
|
# ── nn.Module forward ────────────────────────────────────────────────────
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
y = self.backbone(x)
|
|
if y.dim() > 2:
|
|
y = y.flatten(1)
|
|
return y
|
|
|
|
# ── Utilities ────────────────────────────────────────────────────────────
|
|
|
|
def set_freeze_ratio(self, ratio: float) -> None:
|
|
r = max(0.0, min(1.0, float(ratio)))
|
|
n_freeze = int(math.floor(len(self._blocks) * r))
|
|
for b in self._blocks:
|
|
for p in b.parameters():
|
|
p.requires_grad = True
|
|
for b in self._blocks[:n_freeze]:
|
|
for p in b.parameters():
|
|
p.requires_grad = False
|
|
|
|
# ── Internals ────────────────────────────────────────────────────────────
|
|
|
|
@staticmethod
|
|
def _augment_array(arr: np.ndarray) -> np.ndarray:
|
|
"""Random flip + 90° rotation on a (C, H, W) seg-map array."""
|
|
if np.random.rand() < 0.5:
|
|
arr = arr[:, :, ::-1]
|
|
if np.random.rand() < 0.5:
|
|
arr = arr[:, ::-1, :]
|
|
k = int(np.random.randint(0, 4))
|
|
if k:
|
|
arr = np.rot90(arr, k=k, axes=(1, 2))
|
|
return arr
|
|
|
|
def _adapt_first_conv(self, in_channels: int) -> None:
|
|
"""Replace the first Conv2d to accept a non-3-channel input.
|
|
|
|
Pretrained weights are averaged across the original input channels and
|
|
broadcast across the new ones.
|
|
"""
|
|
first = self._find_first_conv(self.backbone)
|
|
new = nn.Conv2d(
|
|
in_channels,
|
|
first.out_channels,
|
|
kernel_size=first.kernel_size,
|
|
stride=first.stride,
|
|
padding=first.padding,
|
|
bias=first.bias is not None,
|
|
)
|
|
with torch.no_grad():
|
|
new.weight.copy_(
|
|
first.weight.mean(dim=1, keepdim=True).expand_as(new.weight)
|
|
)
|
|
if first.bias is not None:
|
|
new.bias.copy_(first.bias)
|
|
self._replace_first_conv(self.backbone, new)
|
|
|
|
@staticmethod
|
|
def _find_first_conv(module: nn.Module) -> nn.Conv2d:
|
|
for m in module.modules():
|
|
if isinstance(m, nn.Conv2d):
|
|
return m
|
|
raise RuntimeError("No Conv2d found in backbone")
|
|
|
|
@classmethod
|
|
def _replace_first_conv(cls, module: nn.Module, new_conv: nn.Conv2d) -> bool:
|
|
for name, child in module.named_children():
|
|
if isinstance(child, nn.Conv2d):
|
|
setattr(module, name, new_conv)
|
|
return True
|
|
if cls._replace_first_conv(child, new_conv):
|
|
return True
|
|
return False
|