rpotter6298
271 lines
9.7 KiB
Python
271 lines
9.7 KiB
Python
"""
|
|
portable_versions/refuge_mask_adapter.py
|
|
=========================================
|
|
Optic-disc cropper for REFUGE (and REFUGE2) fundus images, designed as a
|
|
drop-in ``preprocessor`` for ``ImageLoader``.
|
|
|
|
Given an image and its corresponding segmentation mask, it:
|
|
1. Extracts the optic disc region from the mask
|
|
2. Computes a padded bounding box around it
|
|
3. Crops and resizes the original image
|
|
|
|
Dependencies: Pillow, numpy (nothing else)
|
|
|
|
Quickstart
|
|
----------
|
|
from portable_versions.image_loader import ImageLoader, RandomHorizontalFlip, RandomRotation, ColorJitter
|
|
from portable_versions.refuge_mask_adapter import RefugeMaskCropper
|
|
|
|
cropper = RefugeMaskCropper(
|
|
mask_dir="REFUGE/Annotations/Training400/Disc_Cup_Masks",
|
|
scale=1.5, # context around disc (1.0 = tight, 2.0 = lots of context)
|
|
target_size=(200, 200), # output size — should match ImageLoader target_size
|
|
mask_suffix=".bmp", # REFUGE1 uses .bmp; REFUGE2 uses .png
|
|
)
|
|
|
|
loader = ImageLoader(
|
|
target_size=(200, 200),
|
|
normalize=True,
|
|
preprocessor=cropper,
|
|
)
|
|
loader.augmentation = [
|
|
RandomHorizontalFlip(),
|
|
RandomRotation(15),
|
|
ColorJitter(0.2, 0.2, 0.1, 0.05),
|
|
]
|
|
|
|
imgs = loader.get_img(image_paths, augment=True) # (N, 200, 200, 3)
|
|
|
|
REFUGE mask formats
|
|
-------------------
|
|
REFUGE1 Grayscale BMP: background=128, disc=255, cup=0
|
|
REFUGE2 RGB PNG: background detected from image borders, disc/cup by colour
|
|
|
|
Both are handled automatically.
|
|
|
|
Directory structure assumption
|
|
------------------------------
|
|
The cropper looks for the mask with the same stem as the image file, inside
|
|
``mask_dir``. If your layout differs, pass a custom ``mask_path_fn``:
|
|
|
|
cropper = RefugeMaskCropper(
|
|
mask_path_fn=lambda img_path: img_path.with_suffix(".bmp"),
|
|
scale=1.5,
|
|
target_size=(200, 200),
|
|
)
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from collections import Counter
|
|
from pathlib import Path
|
|
from typing import Optional, Tuple
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
|
|
_MASK_DIR_NAMES = {"Disc_Cup_Masks", "Disc_Masks", "Disc_Mask"}
|
|
_MASK_SUFFIXES = {".bmp", ".png"}
|
|
|
|
|
|
class RefugeMaskCropper:
|
|
"""
|
|
Crop a fundus image to the optic disc region using its segmentation mask.
|
|
|
|
Pass the REFUGE root directory and the cropper will automatically index
|
|
all masks underneath it — no need to specify which subdirectory or
|
|
file extension.
|
|
|
|
cropper = RefugeMaskCropper("REFUGE/", scale=1.5)
|
|
loader = ImageLoader(target_size=(200, 200), preprocessor=cropper)
|
|
imgs = loader.get_img(test_set) # test_set = any list of image paths
|
|
|
|
Parameters
|
|
----------
|
|
refuge_root : str or Path
|
|
Top-level REFUGE directory. All mask files under directories named
|
|
``Disc_Cup_Masks``, ``Disc_Masks``, or ``Disc_Mask`` are indexed
|
|
automatically (supports both .bmp and .png).
|
|
scale : float
|
|
Padding multiplier applied to the disc radius.
|
|
1.0 = tight crop, 1.5 = moderate context, 2.5 = lots of context.
|
|
target_size : (height, width)
|
|
Output size after cropping. Should match ``ImageLoader.target_size``.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
refuge_root: str | Path,
|
|
*,
|
|
scale: float = 1.5,
|
|
target_size: Tuple[int, int] = (200, 200),
|
|
) -> None:
|
|
self.refuge_root = Path(refuge_root)
|
|
self.scale = scale
|
|
self.target_size = target_size
|
|
self._index: dict[str, list[Path]] = {}
|
|
self._build_index()
|
|
|
|
def _build_index(self) -> None:
|
|
"""Walk refuge_root and index all mask files by stem (stem → [paths])."""
|
|
for mask_dir in self.refuge_root.rglob("*"):
|
|
if mask_dir.is_dir() and mask_dir.name in _MASK_DIR_NAMES:
|
|
for f in mask_dir.rglob("*"):
|
|
if f.is_file() and f.suffix.lower() in _MASK_SUFFIXES:
|
|
self._index.setdefault(f.stem, []).append(f)
|
|
if not self._index:
|
|
raise FileNotFoundError(
|
|
f"No mask files found under {self.refuge_root!r}. "
|
|
f"Expected directories named: {_MASK_DIR_NAMES}"
|
|
)
|
|
n_masks = sum(len(v) for v in self._index.values())
|
|
print(f"[RefugeMaskCropper] indexed {n_masks} masks ({len(self._index)} unique stems)", flush=True)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Callable interface — drop-in preprocessor for ImageLoader
|
|
# ------------------------------------------------------------------
|
|
|
|
def __call__(
|
|
self,
|
|
img: Image.Image,
|
|
img_path: Optional[str | Path] = None,
|
|
) -> Image.Image:
|
|
stem = Path(img_path).stem if img_path else None
|
|
mask_path = self._lookup(stem, img_path)
|
|
disc_mask = _load_disc_mask(mask_path, img.size)
|
|
box = _mask_to_crop_box(disc_mask, scale=self.scale, img_size=img.size)
|
|
cropped = img.crop(box)
|
|
return cropped.resize(
|
|
(self.target_size[1], self.target_size[0]), Image.Resampling.BILINEAR
|
|
)
|
|
|
|
def _lookup(self, stem: Optional[str], img_path: Optional[str | Path] = None) -> Path:
|
|
if stem is None:
|
|
raise ValueError("img_path is required to match the mask.")
|
|
candidates = self._index.get(stem)
|
|
if not candidates:
|
|
raise KeyError(
|
|
f"No mask found for image stem {stem!r}. "
|
|
f"Available stems (sample): {list(self._index)[:5]}"
|
|
)
|
|
if len(candidates) == 1:
|
|
return candidates[0]
|
|
# Pick the mask whose directory components best overlap with img_path
|
|
# (ignores the filename itself to handle extension differences)
|
|
img_parts = set(Path(img_path).parent.parts) if img_path else set()
|
|
return max(candidates, key=lambda m: len(set(m.parent.parts) & img_parts))
|
|
|
|
def __repr__(self) -> str:
|
|
return (
|
|
f"RefugeMaskCropper(refuge_root={str(self.refuge_root)!r}, "
|
|
f"scale={self.scale}, target_size={self.target_size}, "
|
|
f"masks_indexed={len(self._index)})"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Mask parsing
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _load_disc_mask(mask_path: Path, img_size: Tuple[int, int]) -> np.ndarray:
|
|
"""
|
|
Return a binary disc mask (uint8, 1=disc) from a REFUGE mask file.
|
|
|
|
Handles:
|
|
- Grayscale BMP (REFUGE1): background≈128, disc=255, cup=0
|
|
- RGB PNG (REFUGE2): background detected from image borders
|
|
"""
|
|
mask_img = Image.open(mask_path)
|
|
|
|
if mask_img.mode == "L" or mask_img.mode == "P":
|
|
arr = np.asarray(mask_img.convert("L"), dtype=np.uint8)
|
|
bg = _border_mode(arr)
|
|
disc_mask = (arr != bg).astype(np.uint8)
|
|
else:
|
|
arr = np.asarray(mask_img.convert("RGB"), dtype=np.uint8)
|
|
bg = _border_mode_rgb(arr)
|
|
# disc = any non-background pixel
|
|
bg_mask = np.all(arr == bg, axis=2)
|
|
disc_mask = (~bg_mask).astype(np.uint8)
|
|
|
|
# Ensure mask matches image spatial size
|
|
mh, mw = disc_mask.shape
|
|
iw, ih = img_size
|
|
if (mw, mh) != (iw, ih):
|
|
disc_img = Image.fromarray(disc_mask * 255).resize((iw, ih), Image.NEAREST)
|
|
disc_mask = (np.asarray(disc_img) > 0).astype(np.uint8)
|
|
|
|
return disc_mask
|
|
|
|
|
|
def _border_mode(arr: np.ndarray, border: int = 5) -> int:
|
|
"""Most common pixel value along the image border (grayscale)."""
|
|
h, w = arr.shape
|
|
border_pixels = np.concatenate([
|
|
arr[:border, :].ravel(),
|
|
arr[-border:, :].ravel(),
|
|
arr[:, :border].ravel(),
|
|
arr[:, -border:].ravel(),
|
|
])
|
|
return int(Counter(border_pixels.tolist()).most_common(1)[0][0])
|
|
|
|
|
|
def _border_mode_rgb(arr: np.ndarray, border: int = 5) -> np.ndarray:
|
|
"""Most common RGB colour along the image border."""
|
|
h, w, _ = arr.shape
|
|
border_pixels = np.concatenate([
|
|
arr[:border, :].reshape(-1, 3),
|
|
arr[-border:, :].reshape(-1, 3),
|
|
arr[:, :border].reshape(-1, 3),
|
|
arr[:, -border:].reshape(-1, 3),
|
|
], axis=0)
|
|
tuples = [tuple(row) for row in border_pixels.tolist()]
|
|
most_common = Counter(tuples).most_common(1)[0][0]
|
|
return np.array(most_common, dtype=np.uint8)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Bounding box from mask
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _mask_to_crop_box(
|
|
disc_mask: np.ndarray,
|
|
scale: float,
|
|
img_size: Tuple[int, int],
|
|
) -> Tuple[int, int, int, int]:
|
|
"""
|
|
Compute a square crop box centred on the disc with padding = scale * radius.
|
|
|
|
Returns (left, upper, right, lower) — ready for PIL Image.crop().
|
|
Falls back to the full image if no disc pixels are found.
|
|
"""
|
|
coords = np.argwhere(disc_mask > 0) # (N, 2) in (row, col) order
|
|
if coords.size == 0:
|
|
w, h = img_size
|
|
return (0, 0, w, h)
|
|
|
|
ys, xs = coords[:, 0], coords[:, 1]
|
|
centre_x = float(xs.mean())
|
|
centre_y = float(ys.mean())
|
|
radius = max(float(xs.max() - xs.min()), float(ys.max() - ys.min())) / 2.0
|
|
crop_radius = radius * scale
|
|
|
|
iw, ih = img_size
|
|
left = int(max(0, centre_x - crop_radius))
|
|
upper = int(max(0, centre_y - crop_radius))
|
|
right = int(min(iw, centre_x + crop_radius))
|
|
lower = int(min(ih, centre_y + crop_radius))
|
|
|
|
# Make square by expanding the shorter side
|
|
cw, ch = right - left, lower - upper
|
|
if cw < ch:
|
|
diff = ch - cw
|
|
left = max(0, left - diff // 2)
|
|
right = min(iw, right + diff // 2)
|
|
elif ch < cw:
|
|
diff = cw - ch
|
|
upper = max(0, upper - diff // 2)
|
|
lower = min(ih, lower + diff // 2)
|
|
|
|
return (left, upper, right, lower)
|