File size: 3,075 Bytes
acc1556 | 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 | """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 # e.g. "usb/OK/S0003" or "usb/NG/BX/S0001"
category: str
label: int # 0 = OK, 1 = anomaly
anomaly_class: str # "OK", "BX", ...
image_paths: list = field(default_factory=list) # absolute paths, sorted by camera
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"] # e.g. "OK/S0003/xxx_C1_t.jpg"
parent = os.path.dirname(rel) # "OK/S0003"
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
|