rpotter6298
143 lines
5.1 KiB
Python
Executable File
143 lines
5.1 KiB
Python
Executable File
"""Generate ground-truth mask overlays for REFUGE and Papila samples."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
import sys
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.append(str(ROOT))
|
|
|
|
import numpy as np
|
|
from collections import Counter
|
|
from PIL import Image
|
|
from PIL.Image import Resampling
|
|
|
|
from classes.unet_segmenter import UNetSegmenter
|
|
|
|
REFUGE_ROOT = Path("REFUGE")
|
|
DEFAULT_MANIFEST = Path("manifest.csv")
|
|
OUTPUT_DIR = Path("temp/gt_test")
|
|
|
|
|
|
def to_mask_colors(disc: np.ndarray, cup: np.ndarray) -> Image.Image:
|
|
h, w = disc.shape
|
|
canvas = np.ones((h, w, 3), dtype=np.uint8) * 255
|
|
disc_mask = disc.astype(bool)
|
|
cup_mask = cup.astype(bool)
|
|
canvas[disc_mask] = [128, 128, 128]
|
|
canvas[cup_mask] = [0, 0, 0]
|
|
return Image.fromarray(canvas)
|
|
|
|
|
|
def overlay(
|
|
original: Image.Image, mask_rgb: Image.Image, alpha: float = 0.6
|
|
) -> Image.Image:
|
|
mask_rgba = mask_rgb.convert("RGBA")
|
|
updates = np.array(mask_rgba, dtype=np.float32)
|
|
updates[..., 3] = alpha * 255 * (updates[..., :3] != 255).any(axis=-1)
|
|
base = original.convert("RGBA")
|
|
return Image.alpha_composite(
|
|
base, Image.fromarray(updates.astype(np.uint8))
|
|
).convert("RGB")
|
|
|
|
|
|
def original_mask_to_rgb(mask_path: Path) -> Image.Image:
|
|
mask_img = Image.open(mask_path)
|
|
arr = np.asarray(mask_img)
|
|
h, w = arr.shape[:2]
|
|
canvas = np.ones((h, w, 3), dtype=np.uint8) * 255
|
|
|
|
if arr.ndim == 2:
|
|
border = np.concatenate([arr[0, :], arr[-1, :], arr[:, 0], arr[:, -1]])
|
|
bg_value = Counter(border.tolist()).most_common(1)[0][0]
|
|
disc_mask = arr != bg_value
|
|
fg_counts = Counter(arr[arr != bg_value].flatten())
|
|
if fg_counts:
|
|
# For REFUGE-style masks: cup should be the darkest (minimum value)
|
|
cup_value = min(fg_counts.keys())
|
|
cup_mask = arr == cup_value
|
|
else:
|
|
cup_mask = np.zeros_like(arr, dtype=bool)
|
|
else:
|
|
edges = np.concatenate(
|
|
[arr[0, :, :], arr[-1, :, :], arr[:, 0, :], arr[:, -1, :]], axis=0
|
|
)
|
|
bg_color = Counter(map(tuple, edges)).most_common(1)[0][0]
|
|
disc_mask = ~np.all(arr == bg_color, axis=-1)
|
|
color_counts = Counter(map(tuple, arr.reshape(-1, arr.shape[2])))
|
|
cup_mask = np.zeros((h, w), dtype=bool)
|
|
candidates = {}
|
|
for color, count in color_counts.items():
|
|
if color == bg_color:
|
|
continue
|
|
mask = np.all(arr == color, axis=-1)
|
|
candidates[color] = mask
|
|
if candidates:
|
|
# For REFUGE-style masks: cup should be the darkest color (closest to black)
|
|
cup_color = min(candidates.keys(), key=lambda color: sum(color))
|
|
cup_mask = candidates[cup_color]
|
|
disc_mask = disc_mask.astype(bool)
|
|
cup_mask = cup_mask & disc_mask
|
|
|
|
canvas[disc_mask] = [128, 128, 128]
|
|
canvas[cup_mask] = [0, 0, 0]
|
|
return Image.fromarray(canvas)
|
|
|
|
|
|
def process_entries(
|
|
segmenter: UNetSegmenter, entries, prefix: str, count: int, dest: Path
|
|
) -> None:
|
|
for entry in entries[:count]:
|
|
img_path = Path(entry.image_path)
|
|
if not img_path.exists():
|
|
continue
|
|
orig = Image.open(img_path).convert("RGB")
|
|
image = segmenter.preprocess_image(orig)
|
|
disc, cup = segmenter.load_masks(entry)
|
|
disc_coords = segmenter._mask_to_coords(disc)
|
|
cup_coords = segmenter._mask_to_coords(cup)
|
|
print(
|
|
f"{entry.sample_id}: disc coords {disc_coords.shape[0] if disc_coords is not None else 0}, "
|
|
f"cup coords {cup_coords.shape[0] if cup_coords is not None else 0}"
|
|
)
|
|
mask_rgb = to_mask_colors(disc, cup)
|
|
overlay_img = overlay(image, mask_rgb)
|
|
mask_rgb.save(dest / f"{prefix}_{entry.sample_id}_mask.png")
|
|
overlay_img.save(dest / f"{prefix}_{entry.sample_id}_overlay.png")
|
|
|
|
if entry.annotation_type_disc == "mask":
|
|
gt_mask_rgb = original_mask_to_rgb(entry.annotation_disc)
|
|
gt_overlay = overlay(
|
|
orig.resize(gt_mask_rgb.size, Resampling.BILINEAR), gt_mask_rgb
|
|
)
|
|
gt_mask_rgb.save(dest / f"{prefix}_{entry.sample_id}_gt_mask.png")
|
|
gt_overlay.save(dest / f"{prefix}_{entry.sample_id}_gt_overlay.png")
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(
|
|
description="Inspect ground-truth masks for REFUGE and Papila"
|
|
)
|
|
parser.add_argument("--manifest", type=Path, default=DEFAULT_MANIFEST)
|
|
parser.add_argument("--output", type=Path, default=OUTPUT_DIR)
|
|
parser.add_argument(
|
|
"--count", type=int, default=10, help="Number of samples per dataset"
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
args.output.mkdir(parents=True, exist_ok=True)
|
|
|
|
segmenter = UNetSegmenter(args.manifest)
|
|
refuge_entries = [e for e in segmenter._manifest if e.dataset == "refuge"]
|
|
papila_entries = [e for e in segmenter._manifest if e.dataset == "papila"]
|
|
|
|
process_entries(segmenter, refuge_entries, "refuge", args.count, args.output)
|
|
process_entries(segmenter, papila_entries, "papila", args.count, args.output)
|
|
print(f"Saved overlays to {args.output}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|