| """Real-IAD data loading with multi-view (sample) grouping. |
| |
| Official annotation JSONs define train/test. Filenames: |
| OK/S0003/usb_0003_OK_C1_20231019140654.jpg |
| NG/BX/S0001/usb_0001_NG_BX_C1_20231020105442.jpg |
| A physical sample = the image's parent directory; every sample has 5 views (C1..C5). |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import os |
| from dataclasses import dataclass, field |
|
|
| from PIL import Image |
|
|
| Image.MAX_IMAGE_PIXELS = None |
|
|
|
|
| @dataclass |
| class Sample: |
| """One physical object instance with its 5 camera views.""" |
| sample_id: str |
| category: str |
| label: int |
| anomaly_class: str |
| image_paths: list = field(default_factory=list) |
| mask_paths: list = field(default_factory=list) |
|
|
| @property |
| def n_views(self): |
| return len(self.image_paths) |
|
|
|
|
| def _load_split_json(category: str, json_root: str) -> dict: |
| path = os.path.join(json_root, f"{category}.json") |
| with open(path) as f: |
| return json.load(f) |
|
|
|
|
| def load_samples(category: str, image_root: str, json_root: str, split: str) -> list: |
| """Return samples for one category/split, grouped by parent directory.""" |
| d = _load_split_json(category, json_root) |
| base = os.path.join(image_root, category) |
| groups = {} |
| for rec in d[split]: |
| rel = rec["image_path"] |
| parent = os.path.dirname(rel) |
| key = f"{category}/{parent}" |
| g = groups.setdefault(key, Sample( |
| sample_id=key, category=category, |
| label=0 if rec["anomaly_class"] == "OK" else 1, |
| anomaly_class=rec["anomaly_class"], |
| )) |
| g.image_paths.append(os.path.join(base, rel)) |
| mp = rec.get("mask_path") |
| g.mask_paths.append(os.path.join(base, mp) if mp else None) |
| out = list(groups.values()) |
| for s in out: |
| order = sorted(range(s.n_views), key=lambda i: s.image_paths[i]) |
| s.image_paths = [s.image_paths[i] for i in order] |
| s.mask_paths = [s.mask_paths[i] for i in order] |
| return out |
|
|
|
|
| def subsample_samples(samples: list, n_ok: int, n_ng: int, seed: int = 0) -> list: |
| """Deterministic stratified pick of samples (keeps all 5 views).""" |
| import random |
| rng = random.Random(seed) |
| ok = [s for s in samples if s.label == 0] |
| ng = [s for s in samples if s.label == 1] |
| rng.shuffle(ok) |
| rng.shuffle(ng) |
| return ok[:n_ok] + ng[:n_ng] |
|
|
|
|
| def split_calibration(samples: list, calib_frac: float = 0.25, seed: int = 0): |
| """Stratified sample-level split into (calibration, eval).""" |
| import random |
| rng = random.Random(seed) |
| ok = [s for s in samples if s.label == 0] |
| ng = [s for s in samples if s.label == 1] |
| rng.shuffle(ok) |
| rng.shuffle(ng) |
| n_ok, n_ng = int(len(ok) * calib_frac), int(len(ng) * calib_frac) |
| calib = ok[:n_ok] + ng[:n_ng] |
| ev = ok[n_ok:] + ng[n_ng:] |
| return calib, ev |
|
|
