rpotter6298
473 lines
19 KiB
Python
473 lines
19 KiB
Python
"""Optic-disc image croppers and preprocessor factory for V2."""
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import Dict, Optional, Tuple
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import torch
|
|
from PIL import Image, ImageDraw
|
|
from torchvision import transforms
|
|
|
|
from v3.classes.geometry_features import compute_geometry_features, disc_cup_from_mask_image
|
|
from v3.classes.unet_segmenter import UNetSegmenter
|
|
|
|
|
|
def _geometry_from_mask(mask: np.ndarray, scale: float) -> Dict:
|
|
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,
|
|
}
|
|
|
|
|
|
class UNetImageCropper:
|
|
def __init__(
|
|
self,
|
|
manifest_path: Path,
|
|
weights_path: Path,
|
|
normalize: str = "per_image",
|
|
threshold: float = 0.5,
|
|
tta: bool = False,
|
|
scale: float = 2.5,
|
|
target_size: int = 224,
|
|
cache_dir: Optional[Path] = None,
|
|
) -> None:
|
|
self.segmenter = UNetSegmenter(
|
|
manifest_path=manifest_path,
|
|
normalize=normalize,
|
|
)
|
|
state = torch.load(weights_path, map_location=self.segmenter.device)
|
|
state_dict = state.get("model", state)
|
|
self.segmenter.model.load_state_dict(state_dict)
|
|
self.segmenter.model.to(self.segmenter.device)
|
|
self.segmenter.model.eval()
|
|
|
|
self.threshold = threshold
|
|
self.tta = tta
|
|
self.scale = scale
|
|
self.target_size = target_size
|
|
self.cache_dir = Path(cache_dir) if cache_dir is not None else None
|
|
if self.cache_dir is not None:
|
|
self.cache_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
self.to_tensor = transforms.ToTensor()
|
|
|
|
def _cache_path(self, image_path: Path) -> Optional[Path]:
|
|
if self.cache_dir is None:
|
|
return None
|
|
stem = image_path.stem
|
|
return self.cache_dir / f"{stem}_s{int(self.scale * 100)}.npz"
|
|
|
|
def clear_cache(self) -> None:
|
|
if self.cache_dir is None or not self.cache_dir.exists():
|
|
return
|
|
removed = sum(1 for f in self.cache_dir.glob("*.npz") if f.unlink() or True)
|
|
print(f"[UNetImageCropper] Cleared {removed} cached crop files from {self.cache_dir}")
|
|
|
|
def _infer_masks(self, image: Image.Image) -> Optional[Tuple[np.ndarray, np.ndarray]]:
|
|
resized = self.segmenter.preprocess_image(image)
|
|
tensor = self.segmenter._normalize_tensor(
|
|
self.to_tensor(resized).to(self.segmenter.device)
|
|
).unsqueeze(0)
|
|
|
|
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)
|
|
disc_mask = np.array(disc_img, dtype=np.uint8)
|
|
cup_img = Image.fromarray(cup_pred, mode="L").resize(image.size, Image.NEAREST)
|
|
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)
|
|
disc_mask = (disc_mask > 0).astype(np.uint8)
|
|
return disc_mask, cup_mask
|
|
|
|
def _compute_crop_info(self, image: Image.Image, image_path: Path) -> Optional[dict]:
|
|
image_path = Path(image_path).resolve()
|
|
cache_path = self._cache_path(image_path)
|
|
cached_bounds = None
|
|
if cache_path is not None and cache_path.exists():
|
|
data = np.load(cache_path, allow_pickle=False)
|
|
try:
|
|
cached_bounds = {
|
|
"left": float(data["left"]),
|
|
"upper": float(data["upper"]),
|
|
"right": float(data["right"]),
|
|
"lower": float(data["lower"]),
|
|
}
|
|
if "features" in data.files:
|
|
cached_bounds["features"] = data["features"].astype(np.float32)
|
|
return cached_bounds
|
|
except KeyError:
|
|
cached_bounds = None
|
|
|
|
masks = self._infer_masks(image)
|
|
if masks is None:
|
|
return cached_bounds
|
|
disc_mask, cup_mask = masks
|
|
try:
|
|
geom = _geometry_from_mask(disc_mask, self.scale)
|
|
except Exception:
|
|
return cached_bounds
|
|
cx = geom["centre_x"]
|
|
cy = geom["centre_y"]
|
|
r = geom["crop_radius"]
|
|
left = max(0.0, cx - r)
|
|
upper = max(0.0, cy - r)
|
|
right = min(float(image.width), cx + r)
|
|
lower = min(float(image.height), cy + r)
|
|
features = compute_geometry_features(disc_mask, cup_mask)
|
|
|
|
info = {
|
|
"left": left,
|
|
"upper": upper,
|
|
"right": right,
|
|
"lower": lower,
|
|
"features": features,
|
|
}
|
|
if cache_path is not None:
|
|
np.savez(
|
|
cache_path,
|
|
left=left,
|
|
upper=upper,
|
|
right=right,
|
|
lower=lower,
|
|
width=float(image.width),
|
|
height=float(image.height),
|
|
scale=self.scale,
|
|
target_size=self.target_size,
|
|
features=features,
|
|
)
|
|
return info
|
|
|
|
def __call__(self, image: Image.Image, image_path: Path) -> Image.Image:
|
|
info = self._compute_crop_info(image, image_path)
|
|
if info is None:
|
|
return image
|
|
left = info["left"]
|
|
upper = info["upper"]
|
|
right = info["right"]
|
|
lower = info["lower"]
|
|
if right <= left or lower <= upper:
|
|
return image
|
|
crop = image.crop((left, upper, right, lower))
|
|
return crop.resize((self.target_size, self.target_size), Image.BILINEAR)
|
|
|
|
def geometry_features(self, image: Image.Image, image_path: Path) -> Optional[np.ndarray]:
|
|
info = self._compute_crop_info(image, image_path)
|
|
if info is None:
|
|
return None
|
|
features = info.get("features")
|
|
if features is None:
|
|
return None
|
|
return np.asarray(features, dtype=np.float32)
|
|
|
|
def precompute_geometry(self, image_paths) -> None:
|
|
"""Pre-compute geometry features for all image_paths into an in-memory cache.
|
|
Safe to call in the main process; geometry_for_image() can then be called
|
|
from DataLoader workers without touching CUDA.
|
|
"""
|
|
self._geometry_cache: Dict[str, Optional[np.ndarray]] = {}
|
|
paths = list(image_paths)
|
|
print(f"[UNetImageCropper] pre-computing geometry for {len(paths)} images...", flush=True)
|
|
for img_path in paths:
|
|
key = str(Path(img_path).resolve())
|
|
try:
|
|
img = Image.open(img_path).convert("RGB")
|
|
self._geometry_cache[key] = self.geometry_features(img, img_path)
|
|
except Exception:
|
|
self._geometry_cache[key] = None
|
|
n_ok = sum(1 for v in self._geometry_cache.values() if v is not None)
|
|
print(f"[UNetImageCropper] {n_ok}/{len(paths)} geometry vectors computed", flush=True)
|
|
|
|
def geometry_for_image(self, image_path) -> Optional[np.ndarray]:
|
|
"""Return pre-computed geometry vector for image_path (call precompute_geometry first)."""
|
|
cache = getattr(self, "_geometry_cache", None)
|
|
if cache is None:
|
|
raise RuntimeError("Call precompute_geometry() before geometry_for_image()")
|
|
return cache.get(str(Path(image_path).resolve()))
|
|
|
|
|
|
class ManifestImageCropper:
|
|
def __init__(
|
|
self,
|
|
manifest_path: Path,
|
|
scale: float = 2.5,
|
|
target_size: int = 224,
|
|
cache_dir: Optional[Path] = None,
|
|
) -> None:
|
|
self.scale = scale
|
|
self.target_size = target_size
|
|
self.cache_dir = Path(cache_dir) if cache_dir is not None else None
|
|
if self.cache_dir is not None:
|
|
self.cache_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
df = pd.read_csv(manifest_path)
|
|
self.entries: Dict[str, dict] = {}
|
|
for _, row in df.iterrows():
|
|
img_path = Path(row["image_path"]).resolve()
|
|
self.entries[str(img_path)] = {
|
|
"annotation_disc": row.get("annotation_disc"),
|
|
"annotation_cup": row.get("annotation_cup"),
|
|
"annotation_type_disc": row.get("annotation_type_disc"),
|
|
"annotation_type_cup": row.get("annotation_type_cup"),
|
|
}
|
|
|
|
def _cache_path(self, image_path: Path) -> Optional[Path]:
|
|
if self.cache_dir is None:
|
|
return None
|
|
return self.cache_dir / f"{image_path.stem}_s{int(self.scale * 100)}.npz"
|
|
|
|
def clear_cache(self) -> None:
|
|
if self.cache_dir is None or not self.cache_dir.exists():
|
|
return
|
|
removed = sum(1 for f in self.cache_dir.glob("*.npz") if f.unlink() or True)
|
|
print(f"[ManifestImageCropper] Cleared {removed} cached crop files from {self.cache_dir}")
|
|
|
|
@staticmethod
|
|
def _load_contour(path: Path) -> np.ndarray:
|
|
coords = np.loadtxt(path)
|
|
if coords.ndim == 1:
|
|
coords = coords.reshape(-1, 2)
|
|
return coords
|
|
|
|
@staticmethod
|
|
def _contour_to_mask(coords: np.ndarray, size: tuple[int, int]) -> np.ndarray:
|
|
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)
|
|
points = [tuple(map(float, pt)) for pt in coords]
|
|
draw.polygon(points, outline=1, fill=1)
|
|
return np.array(img, dtype=np.uint8)
|
|
|
|
def _load_masks(self, entry: dict, image: Image.Image) -> Optional[Tuple[np.ndarray, np.ndarray]]:
|
|
disc_path = entry.get("annotation_disc")
|
|
cup_path = entry.get("annotation_cup")
|
|
disc_type = (entry.get("annotation_type_disc") or "").lower()
|
|
cup_type = (entry.get("annotation_type_cup") or "").lower()
|
|
|
|
disc_mask: Optional[np.ndarray] = None
|
|
cup_mask: Optional[np.ndarray] = None
|
|
|
|
if disc_path and not pd.isna(disc_path):
|
|
disc_path = Path(disc_path)
|
|
try:
|
|
if disc_type == "mask":
|
|
mask_img = Image.open(disc_path)
|
|
mask_img = mask_img.resize(image.size, Image.NEAREST)
|
|
disc_mask, cup_from_mask = disc_cup_from_mask_image(mask_img)
|
|
if cup_from_mask.sum() > 0:
|
|
cup_mask = cup_from_mask
|
|
elif disc_type == "contour":
|
|
coords = self._load_contour(disc_path)
|
|
disc_mask = self._contour_to_mask(coords, image.size)
|
|
except Exception:
|
|
disc_mask = None
|
|
|
|
if cup_mask is None and cup_path and not pd.isna(cup_path):
|
|
cup_path = Path(cup_path)
|
|
try:
|
|
if cup_type == "mask":
|
|
mask_img = Image.open(cup_path)
|
|
mask_img = mask_img.resize(image.size, Image.NEAREST)
|
|
_, cup_mask = disc_cup_from_mask_image(mask_img)
|
|
elif cup_type == "contour":
|
|
coords = self._load_contour(cup_path)
|
|
cup_mask = self._contour_to_mask(coords, image.size)
|
|
except Exception:
|
|
cup_mask = None
|
|
|
|
if disc_mask is None:
|
|
return None
|
|
disc_mask = (disc_mask > 0).astype(np.uint8)
|
|
if cup_mask is None:
|
|
cup_mask = np.zeros_like(disc_mask, dtype=np.uint8)
|
|
cup_mask = ((cup_mask > 0) & (disc_mask > 0)).astype(np.uint8)
|
|
return disc_mask, cup_mask
|
|
|
|
def _compute_crop_info(self, image: Image.Image, image_path: Path) -> Optional[dict]:
|
|
image_path = Path(image_path).resolve()
|
|
entry = self.entries.get(str(image_path))
|
|
if entry is None:
|
|
return None
|
|
cache_path = self._cache_path(image_path)
|
|
cached_bounds = None
|
|
if cache_path is not None and cache_path.exists():
|
|
data = np.load(cache_path, allow_pickle=False)
|
|
try:
|
|
cached_bounds = {
|
|
"left": float(data["left"]),
|
|
"upper": float(data["upper"]),
|
|
"right": float(data["right"]),
|
|
"lower": float(data["lower"]),
|
|
}
|
|
if "features" in data.files:
|
|
cached_bounds["features"] = data["features"].astype(np.float32)
|
|
return cached_bounds
|
|
except KeyError:
|
|
cached_bounds = None
|
|
|
|
masks = self._load_masks(entry, image)
|
|
if masks is None:
|
|
return cached_bounds
|
|
disc_mask, cup_mask = masks
|
|
try:
|
|
geom = _geometry_from_mask(disc_mask, self.scale)
|
|
except Exception:
|
|
return cached_bounds
|
|
cx = geom["centre_x"]
|
|
cy = geom["centre_y"]
|
|
r = geom["crop_radius"]
|
|
left = max(0.0, cx - r)
|
|
upper = max(0.0, cy - r)
|
|
right = min(float(image.width), cx + r)
|
|
lower = min(float(image.height), cy + r)
|
|
features = compute_geometry_features(disc_mask, cup_mask)
|
|
|
|
info = {
|
|
"left": left,
|
|
"upper": upper,
|
|
"right": right,
|
|
"lower": lower,
|
|
"features": features,
|
|
}
|
|
if cache_path is not None:
|
|
np.savez(
|
|
cache_path,
|
|
left=left,
|
|
upper=upper,
|
|
right=right,
|
|
lower=lower,
|
|
width=float(image.width),
|
|
height=float(image.height),
|
|
scale=self.scale,
|
|
target_size=self.target_size,
|
|
features=features,
|
|
)
|
|
return info
|
|
|
|
def __call__(self, image: Image.Image, image_path: Path) -> Image.Image:
|
|
info = self._compute_crop_info(image, image_path)
|
|
if info is None:
|
|
return image
|
|
left = info["left"]
|
|
upper = info["upper"]
|
|
right = info["right"]
|
|
lower = info["lower"]
|
|
if right <= left or lower <= upper:
|
|
return image
|
|
crop = image.crop((left, upper, right, lower))
|
|
return crop.resize((self.target_size, self.target_size), Image.BILINEAR)
|
|
|
|
def geometry_features(self, image: Image.Image, image_path: Path) -> Optional[np.ndarray]:
|
|
info = self._compute_crop_info(image, image_path)
|
|
if info is None:
|
|
return None
|
|
features = info.get("features")
|
|
if features is None:
|
|
return None
|
|
return np.asarray(features, dtype=np.float32)
|
|
|
|
def precompute_geometry(self, image_paths) -> None:
|
|
"""Pre-compute geometry features for all image_paths into an in-memory cache.
|
|
Safe to call in the main process; geometry_for_image() can then be called
|
|
without re-opening images or re-loading annotations.
|
|
"""
|
|
self._geometry_cache: Dict[str, Optional[np.ndarray]] = {}
|
|
paths = list(image_paths)
|
|
print(f"[ManifestImageCropper] pre-computing geometry for {len(paths)} images...", flush=True)
|
|
for img_path in paths:
|
|
key = str(Path(img_path).resolve())
|
|
entry = self.entries.get(key)
|
|
if entry is None:
|
|
self._geometry_cache[key] = None
|
|
continue
|
|
try:
|
|
img = Image.open(img_path).convert("RGB")
|
|
self._geometry_cache[key] = self.geometry_features(img, img_path)
|
|
except Exception:
|
|
self._geometry_cache[key] = None
|
|
n_ok = sum(1 for v in self._geometry_cache.values() if v is not None)
|
|
print(f"[ManifestImageCropper] {n_ok}/{len(paths)} geometry vectors computed", flush=True)
|
|
|
|
def geometry_for_image(self, image_path) -> Optional[np.ndarray]:
|
|
"""Return pre-computed geometry vector for image_path (call precompute_geometry first)."""
|
|
cache = getattr(self, "_geometry_cache", None)
|
|
if cache is None:
|
|
raise RuntimeError("Call precompute_geometry() before geometry_for_image()")
|
|
return cache.get(str(Path(image_path).resolve()))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Factory
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def build_image_preprocessor_from_args(args):
|
|
"""Construct the correct image cropper from CLI args, or return None."""
|
|
crop_manifest = getattr(args, "img_crop_manifest", None)
|
|
crop_weights = getattr(args, "img_crop_weights", None)
|
|
use_gt = bool(getattr(args, "img_crop_gt", False))
|
|
if not crop_manifest:
|
|
return None
|
|
crop_cache = Path(getattr(args, "img_crop_cache", Path("cache_data/hypertower_crops")))
|
|
persist_cache = bool(getattr(args, "persist_img_crop_cache", False))
|
|
if use_gt:
|
|
pre = ManifestImageCropper(
|
|
manifest_path=Path(crop_manifest),
|
|
scale=getattr(args, "img_crop_scale", 2.5),
|
|
target_size=getattr(args, "img_crop_size", 224),
|
|
cache_dir=crop_cache,
|
|
)
|
|
if not persist_cache:
|
|
pre.clear_cache()
|
|
print(f"[V2 modes] GT disc cropper enabled -> cache at {crop_cache}", flush=True)
|
|
return pre
|
|
if crop_weights:
|
|
pre = UNetImageCropper(
|
|
manifest_path=Path(crop_manifest),
|
|
weights_path=Path(crop_weights),
|
|
normalize=getattr(args, "img_crop_normalize", "per_image"),
|
|
threshold=getattr(args, "img_crop_threshold", 0.5),
|
|
tta=getattr(args, "img_crop_tta", False),
|
|
scale=getattr(args, "img_crop_scale", 2.5),
|
|
target_size=getattr(args, "img_crop_size", 224),
|
|
cache_dir=crop_cache,
|
|
)
|
|
if not persist_cache:
|
|
pre.clear_cache()
|
|
print(f"[V2 modes] UNet disc cropper enabled -> cache at {crop_cache}", flush=True)
|
|
return pre
|
|
print(
|
|
"[V2 modes] img_crop_manifest provided but no --img-crop-gt or --img-crop-weights; cropping disabled.",
|
|
flush=True,
|
|
)
|
|
return None
|