File size: 6,951 Bytes
b8fae22 1a18f22 b8fae22 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 | """Dataset reader for the standardized `processed_unified` layout.
Expected layout (see dataset/SEGMENTATION_WORKSPACE_README.md):
<data_root>/<dataset>/<protocol>/<split>/images/*.png
<data_root>/<dataset>/<protocol>/<split>/masks/*.png
<data_root>/<dataset>/metadata.json (optional, preferred)
<data_root>/<dataset>/manifest.jsonl (optional)
Returns per item: {"image": FloatTensor[C,H,W], "mask": LongTensor[H,W], "name": str}.
Binary and multi-class masks are both supported: masks keep their integer class
ids (0..C-1). Auto-detection of in_channels / num_classes falls back to scanning
files when metadata is absent, so the loader is robust to missing metadata.
"""
from __future__ import annotations
import json
import os
from glob import glob
from typing import Optional, Callable, List, Tuple
import numpy as np
import cv2
from torch.utils.data import Dataset
_MODALITY_CHANNELS = { # hint table; only used when metadata lacks in_channels
"rgb": 3, "fundus": 3, "colonoscopy": 3, "endoscopy": 3, "histopathology": 3,
"ultrasound": 1, "mri": 1, "ct": 1, "grayscale": 1,
}
# Documented class counts (incl. background). metadata.json on the server has no
# num_classes field, so this table is the fast, reliable primary source; unknown
# datasets fall back to a FULL scan of the mask set (accurate but slower).
_KNOWN_NUM_CLASSES = {
"cvc_clinicdb": 2, "kvasir_seg": 2, "fives": 2, "busi": 2,
"refuge2": 3, "acdc_png": 4,
"idridd_segmentation": 6, "pannuke_semantic": 6,
}
def _read_metadata(data_root: str, dataset: str) -> dict:
path = os.path.join(data_root, dataset, "metadata.json")
if os.path.isfile(path):
try:
with open(path) as f:
return json.load(f)
except Exception:
return {}
return {}
def _pair_from_manifest(split_dir: str, manifest: str) -> Optional[List[Tuple[str, str]]]:
if not os.path.isfile(manifest):
return None
pairs = []
base = os.path.dirname(manifest)
with open(manifest) as f:
for line in f:
line = line.strip()
if not line:
continue
rec = json.loads(line)
img = rec.get("image") or rec.get("image_path") or rec.get("img")
msk = rec.get("mask") or rec.get("mask_path") or rec.get("label")
if img is None or msk is None:
return None
# manifest paths may be relative to dataset root or absolute
ip = img if os.path.isabs(img) else os.path.join(base, img)
mp = msk if os.path.isabs(msk) else os.path.join(base, msk)
# only keep entries that fall under this split dir
if os.path.normpath(split_dir) in os.path.normpath(ip):
pairs.append((ip, mp))
return pairs or None
def _pair_by_glob(split_dir: str) -> List[Tuple[str, str]]:
img_dir = os.path.join(split_dir, "images")
msk_dir = os.path.join(split_dir, "masks")
imgs = sorted(glob(os.path.join(img_dir, "*")))
pairs = []
for ip in imgs:
stem = os.path.splitext(os.path.basename(ip))[0]
# mask may share extension or be .png
cands = glob(os.path.join(msk_dir, stem + ".*"))
if not cands:
continue
pairs.append((ip, cands[0]))
return pairs
def detect_in_channels(meta: dict, sample_img: Optional[str]) -> int:
if meta.get("in_channels"):
return int(meta["in_channels"])
mod = str(meta.get("modality", "")).lower()
for k, v in _MODALITY_CHANNELS.items():
if k in mod:
return v
if sample_img and os.path.isfile(sample_img):
im = cv2.imread(sample_img, cv2.IMREAD_UNCHANGED)
if im is not None and im.ndim == 3 and im.shape[2] >= 3:
return 3
return 1
def detect_num_classes(meta: dict, mask_paths: List[str], dataset: str = "") -> int:
if dataset in _KNOWN_NUM_CLASSES:
return _KNOWN_NUM_CLASSES[dataset]
if meta.get("num_classes"):
return int(meta["num_classes"])
# unknown dataset: scan ALL masks so a rare class is never missed
vals = set()
for mp in mask_paths:
m = cv2.imread(mp, cv2.IMREAD_GRAYSCALE)
if m is not None:
vals.update(np.unique(m).tolist())
if not vals:
return 2
maxv = max(vals)
return int(maxv) + 1 if maxv >= 1 else 2
class UnifiedSegDataset(Dataset):
def __init__(self, data_root: str, dataset: str, protocol: str, split: str,
transform: Optional[Callable] = None,
in_channels: int = 0, num_classes: int = 0,
synth_dir: str = ""):
self.data_root = data_root
self.dataset = dataset
self.split = split
self.transform = transform
split_dir = os.path.join(data_root, dataset, protocol, split)
if not os.path.isdir(split_dir):
raise FileNotFoundError(
f"split dir not found: {split_dir}\n"
f"(data is prepared separately; see dataset/ scripts)")
manifest = os.path.join(data_root, dataset, "manifest.jsonl")
pairs = _pair_from_manifest(split_dir, manifest) or _pair_by_glob(split_dir)
if not pairs:
raise RuntimeError(f"no (image,mask) pairs found in {split_dir}")
# optionally merge synthetic (image,mask) pairs into the (train) split
if synth_dir and os.path.isdir(synth_dir):
sp = _pair_by_glob(synth_dir if os.path.isdir(os.path.join(synth_dir, "images"))
else os.path.dirname(synth_dir))
pairs = pairs + sp
self.pairs = pairs
meta = _read_metadata(data_root, dataset)
self.in_channels = in_channels or detect_in_channels(meta, pairs[0][0])
self.num_classes = num_classes or detect_num_classes(meta, [p[1] for p in pairs], dataset)
def __len__(self) -> int:
return len(self.pairs)
def _load_image(self, path: str) -> np.ndarray:
if self.in_channels == 1:
im = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
if im is None:
raise IOError(f"cannot read image {path}")
return im[:, :, None] # H,W,1
im = cv2.imread(path, cv2.IMREAD_COLOR) # BGR
if im is None:
raise IOError(f"cannot read image {path}")
return cv2.cvtColor(im, cv2.COLOR_BGR2RGB) # H,W,3
def __getitem__(self, idx: int):
ip, mp = self.pairs[idx]
image = self._load_image(ip)
mask = cv2.imread(mp, cv2.IMREAD_GRAYSCALE)
if mask is None:
raise IOError(f"cannot read mask {mp}")
mask = mask.astype(np.int64)
if self.transform is not None:
image, mask = self.transform(image, mask)
return {"image": image, "mask": mask,
"name": os.path.splitext(os.path.basename(ip))[0]}
|