Download src/data_d3.py from Nabidnur/AIP-AD: direct link, hf CLI and curl.
- Browser
- Download file 1.96 kB
-
https://huggingface.co/Nabidnur/AIP-AD/resolve/main/src/data_d3.py
- Command line
-
hf download hf://Nabidnur/AIP-AD/src/data_d3.py
-
curl -L -o data_d3.py https://huggingface.co/Nabidnur/AIP-AD/resolve/main/src/data_d3.py
1.96 kB
| """Real-IAD D3 (CVPR 2025) adapter: RGB / PS / XYZ / PCD as a 4-agent swarm. | |
| Each D3 sample directory holds one file per modality: | |
| RGBL05 (jpg) - color image | |
| PS (jpg) - photometric-stereo surface map (pseudo-3D normals) | |
| XYZ (tiff) - surface height/coordinate map | |
| PCD (pcd) - registered point cloud | |
| Splits per category: train=200, test=150, validation=150 (all-labeled). | |
| The fusion layer treats the four modality scores exactly like camera views: | |
| agents that "lie" (miss the defect or cry wolf) are inverted into information. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| from dataclasses import dataclass, field | |
| MODALITIES = ("rgb", "ps", "xyz", "pcd") | |
| class D3Sample: | |
| sample_id: str | |
| category: str | |
| label: int | |
| anomaly_class: str | |
| paths: dict = field(default_factory=dict) # modality -> abs path | |
| def n_views(self): | |
| return len(self.paths) | |
| def load_d3_samples(category: str, image_root: str, json_root: str, split: str, | |
| modalities=MODALITIES) -> list: | |
| """modalities: subset of MODALITIES; pcd requires open3d at scoring time.""" | |
| path = os.path.join(json_root, f"{category}.json") | |
| with open(path) as f: | |
| d = json.load(f) | |
| base = os.path.join(image_root, category) | |
| keymap = {"rgb": "image_path", "ps": "ps_path", "xyz": "xyz_path", "pcd": "pcd_path"} | |
| samples = {} | |
| for rec in d[split]: | |
| parent = os.path.dirname(rec["image_path"]) | |
| key = f"{category}/{parent}" | |
| s = samples.setdefault(key, D3Sample(sample_id=key, category=category, | |
| label=0 if rec["anomaly_class"] == "OK" else 1, | |
| anomaly_class=rec["anomaly_class"])) | |
| for m in modalities: | |
| rel = rec.get(keymap[m]) | |
| if rel: | |
| s.paths[m] = os.path.join(base, rel) | |
| return list(samples.values()) | |