| """PV-PRISM: complete reference-guided, multi-scale defect segmentation.""" |
| from collections import OrderedDict |
| import copy |
| import hashlib |
| import json |
| import os |
| from pathlib import Path |
| import sys |
|
|
| ROOT = Path(__file__).resolve().parent |
| sys.dont_write_bytecode = True |
| os.environ.setdefault("MKL_THREADING_LAYER", "GNU") |
| os.environ.setdefault("USE_TF", "0") |
| os.environ.setdefault("HF_HUB_OFFLINE", "1") |
| sys.path.insert(0, str(ROOT)) |
|
|
| import torch |
| import torch._dynamo |
| if torch.__version__.split("+")[0] != "2.7.1": |
| raise RuntimeError("Use the verified PyTorch 2.7.1 environment; see launch.py and requirements.txt") |
| from torch.nn import functional as F |
| from prism.pipeline.common import Dataset, flips |
| from prism.pipeline.evidence import Evidence |
| from prism.pipeline.components import CachedReferenceSegmenter, LocalExpert, FeatureCache, RegionVerifier, probabilities |
| from prism.models.encoder import FrozenEncoder |
| from prism.models.reference_bank import NormalBank |
| from prism.training.source_pretraining import load_heads |
| from prism.data.geometry import native_logits |
| from prism.data.geometry import read_image, make_canvas |
| from prism.data.loading import load_record |
| from prism.models.reference_segmenter import ReferenceSegmenter |
|
|
|
|
| def read(path): |
| return json.loads(Path(path).read_text()) |
|
|
|
|
| def sha(path): |
| h = hashlib.sha256() |
| with Path(path).open("rb") as f: |
| for block in iter(lambda: f.read(8 << 20), b""): |
| h.update(block) |
| return h.hexdigest() |
|
|
|
|
| def asset(relative): |
| p = (ROOT / relative).resolve() |
| if not p.is_relative_to(ROOT) or not p.is_file(): |
| raise ValueError(f"Missing or external asset: {relative}") |
| return p |
|
|
|
|
| class SupportDataset(Dataset): |
| """Query images and annotated training supports have separate interfaces.""" |
| def __init__(self, dataset, config): |
| self.name, self.spec = dataset, config["specification"] |
| self.train = [] |
| for entry in config["supports"]: |
| row = dict(entry) |
| row["image_path"] = str(asset(row["image_path"])) |
| row["mask_path"] = str(asset(row["mask_path"])) |
| self.train.append(row) |
| self.dev = [] |
| self.by_id = {r["sample_id"]: r for r in self.train} |
| self.image_cache, self.mask_cache = OrderedDict(), OrderedDict() |
| self.limit, self.checked = 48, set() |
| self.support_ids = set(self.by_id) |
| self.known_groups = config["known_image_groups"] |
|
|
| def target(self, row): |
| if row["sample_id"] not in self.support_ids: |
| raise ValueError("Query labels are inaccessible to PV-PRISM inference") |
| return super().target(row) |
|
|
| def query(self, path): |
| path = Path(path).resolve() |
| digest = sha(path) |
| return dict(sample_id="query-" + digest, image_path=str(path), sha256=digest, |
| group_id=self.known_groups.get(digest, "external-" + digest)) |
|
|
|
|
| class ReferenceEvidence(Evidence): |
| """Uses the unchanged registered evidence computation with local assets.""" |
| def __init__(self, dataset, device, registry): |
| cfg = registry["datasets"][dataset] |
| self.data, self.device = SupportDataset(dataset, cfg), device |
| config = copy.deepcopy(registry["segmenter_config"]) |
| for key in ("weights", "metadata"): |
| config["encoder"][key] = str(asset(config["encoder"][key])) |
| config["encoder"]["vendor"] = str((ROOT / config["encoder"]["vendor"]).resolve()) |
| config["encoder"]["global_long_side"] = 1024 |
| config["pv"]["classes"] = cfg["specification"]["classes"][1:] |
| self.core = CachedReferenceSegmenter(FrozenEncoder(config["encoder"]), config).to(device) |
| load_heads(self.core, asset(cfg["weights"]["segmenter"])) |
| self.core.eval().requires_grad_(False) |
| obj = torch.load(asset(registry["normal_bank"]), map_location="cpu", weights_only=True) |
| self.normal = NormalBank(**obj).validate(8).to(device) |
| self.features = FeatureCache(96) |
| self.expert = LocalExpert(len(self.data.spec["classes"]), False).to(device).eval().requires_grad_(False) |
| cp = torch.load(asset(cfg["weights"]["expert"]), map_location="cpu", weights_only=True) |
| self.expert.load_state_dict(cp["model"], strict=True) |
| self.expert_calls, self.support_items = 0, {} |
| self.support_selection = cfg["support_selection"] |
|
|
| def supports(self, c, query_group): |
| rows = [self.data.by_id[i] for i in self.support_selection[str(c)] |
| if self.data.by_id[i]["group_id"] != query_group][:5] |
| if len(rows) != 5 or len({r["group_id"] for r in rows}) != 5: |
| raise ValueError("Five distinct non-query training support groups are required") |
| out = [] |
| for row in rows: |
| key = row["sample_id"], c |
| if key not in self.support_items: |
| self.support_items[key] = self.data.item(row, c, device=self.device) |
| out.append(self.support_items[key]) |
| return out |
|
|
|
|
| class FrozenSupportSegmenter(ReferenceSegmenter): |
| """Memoize immutable support prototypes for the contact refinement branch.""" |
| def __init__(self, encoder, config): |
| super().__init__(encoder, config) |
| self.memo = {} |
|
|
| def _support(self, item, bank, preencoded=None): |
| if self.training or any(p.requires_grad for p in self.parameters()) or preencoded is not None: |
| raise ValueError("Support memoization requires a completely frozen model") |
| key = item["image_id"] |
| if key not in self.memo: |
| self.memo[key] = super()._support(item, bank) |
| return self.memo[key] |
|
|
|
|
| class DevelopmentEvidence: |
| """The original two-channel task preserves canonical crack and any heads.""" |
| def __init__(self, cfg, registry, device): |
| self.device, self.config = device, cfg |
| config = copy.deepcopy(cfg["segmenter_config"]) |
| for key in ("weights", "metadata"): |
| config["encoder"][key] = str(asset(config["encoder"][key])) |
| config["encoder"]["vendor"] = str(ROOT / config["encoder"]["vendor"]) |
| encoder = FrozenEncoder(config["encoder"]) |
| self.canonical = ReferenceSegmenter(encoder, copy.deepcopy(config)).to(device).eval().requires_grad_(False) |
| load_heads(self.canonical, asset(cfg["weights"]["segmenter"])) |
| config["encoder"]["global_long_side"] = 1024 |
| self.high = FrozenSupportSegmenter(encoder, config).to(device).eval().requires_grad_(False) |
| load_heads(self.high, asset(cfg["weights"]["segmenter"])) |
| self.normal = NormalBank(**torch.load(asset(registry["normal_bank"]), map_location="cpu", weights_only=True)).validate(8).to(device) |
| self.expert = LocalExpert(2, False).to(device).eval().requires_grad_(False) |
| self.expert.load_state_dict(torch.load(asset(cfg["weights"]["expert"]), map_location="cpu", weights_only=True)["model"], strict=True) |
| self.supports = {} |
| for category, records in cfg["supports"].items(): |
| self.supports[category] = [] |
| for record in records: |
| row = dict(record) |
| row["image_path"], row["mask_path"] = str(asset(row["image_path"])), str(asset(row["mask_path"])) |
| self.supports[category].append(load_record(row, category, device)) |
|
|
| def query(self, path, view=0): |
| digest = sha(path) |
| group = self.config["known_image_groups"].get(digest, "external-" + digest) |
| if group in {s["group_id"] for rr in self.supports.values() for s in rr}: |
| raise ValueError("A fixed training support cannot be its own query group") |
| row = dict(image_id="anonymous", group_id=group, image_path=str(path)) |
| item = load_record(row, device=self.device, augmentation=dict(horizontal=bool(view & 1), vertical=bool(view & 2), gain=1., gamma=1.)) |
| if set(item) != {"image", "geometry_valid", "geometry", "group_id", "image_id"}: |
| raise ValueError("Query loader exposed annotation") |
| return item |
|
|
| def expert_evidence(self, path, view): |
| image, bits = read_image(path) |
| image = flips(image, view) |
| h, w = image.shape[-2:] |
| scale = 768 / max(h, w) |
| nh, nw = max(1, round(h * scale)), max(1, round(w * scale)) |
| x = F.interpolate(image[None], (nh, nw), mode="bilinear", align_corners=False)[0] |
| x = F.pad(x, (0, 768 - nw, 0, 768 - nh)) |
| _, _, geometry = make_canvas(image, bits) |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| z = self.expert(x[None].to(self.device)) |
| z = F.interpolate(z, x.shape[-2:], mode="bilinear", align_corners=False)[:, :1].float() |
| content = F.interpolate(z[:, :, :nh, :nw], (geometry.content_h, geometry.content_w), mode="bilinear", align_corners=False) |
| return F.pad(content, (0, geometry.canvas_w - geometry.content_w, 0, geometry.canvas_h - geometry.content_h))[0], geometry |
|
|
| @torch.no_grad() |
| def predict(self, path, refiners): |
| q = self.query(path) |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| out = self.canonical(q["image"], q["geometry_valid"], [self.supports[c] for c in ("contact", "crack")], self.normal, subbank_predictions=False) |
| protected = native_logits(out["final"].float(), q["geometry"])[0] |
| any_logit = native_logits(out["any"].float(), q["geometry"])[0] |
| members = [] |
| for view in (0, 1, 2, 3): |
| item = self.query(path, view) |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| out = self.high(item["image"], item["geometry_valid"], [self.supports["contact"]], self.normal, subbank_predictions=False) |
| evidence = torch.cat((out["final"][:, :1].float(), out["absolute"][:, :1].float(), out["any"].float(), out["gate"][:, :1].float()), 1)[0] |
| independent, geometry = self.expert_evidence(path, view) |
| if geometry != item["geometry"]: |
| raise ValueError("Expert alignment changed") |
| evidence = torch.cat((evidence, independent), 0) |
| for refiner in refiners: |
| delta = refiner.canvas_delta(item["image"], evidence, item["geometry_valid"]) |
| z = native_logits((evidence[:1] + delta)[None], geometry)[0, 0] |
| members.append(flips(z.float(), view)) |
| logits = torch.stack((torch.stack(members).mean(0), protected[1])) |
| if not torch.isfinite(logits).all() or not torch.isfinite(any_logit).all(): |
| raise FloatingPointError("Nonfinite prediction") |
| return logits.sigmoid().cpu().numpy(), any_logit.sigmoid().cpu().numpy(), q |
|
|
|
|
| class PVPRISM: |
| """One released architecture: high resolution, four views, two refiners.""" |
| def __init__(self, dataset, device="cuda:0", memory_fraction=.5): |
| registry = read(ROOT / "configs/model.json") |
| if dataset not in registry["datasets"]: |
| raise ValueError(dataset) |
| self.dataset, self.device = dataset, torch.device(device) |
| self.config = registry["datasets"][dataset] |
| if self.device.type != "cuda" or not torch.cuda.is_available(): |
| raise ValueError("The verified full model requires CUDA BF16; CPU parity is not claimed") |
| if not 0 < memory_fraction <= 1: |
| raise ValueError("memory_fraction must be in (0,1]") |
| torch.set_num_threads(2) |
| torch.manual_seed(17007) |
| torch.cuda.set_device(self.device) |
| torch.cuda.set_per_process_memory_fraction(memory_fraction, self.device) |
| torch.backends.cudnn.benchmark = False |
| torch.backends.cuda.matmul.allow_tf32 = False |
| torch.backends.cudnn.allow_tf32 = False |
| self.engine = (DevelopmentEvidence(self.config, registry, self.device) if dataset == "development" |
| else ReferenceEvidence(dataset, self.device, registry)) |
| self.last_any_probability = None |
| self.refiners = [] |
| for relative in self.config["weights"]["refiners"]: |
| model = RegionVerifier().to(self.device).eval().requires_grad_(False) |
| checkpoint = torch.load(asset(relative), map_location="cpu", weights_only=True) |
| model.load_state_dict(checkpoint["editor"], strict=True) |
| self.refiners.append(model) |
|
|
| @torch.no_grad() |
| def predict(self, image): |
| if self.dataset == "development": |
| p, self.last_any_probability, q = self.engine.predict(image, self.refiners) |
| return p, dict(model="PV-PRISM", dataset=self.dataset, classes=["contact", "crack"], |
| image_sha256=sha(image), group_id=q["group_id"], geometry=q["geometry"].to_dict(), |
| contact_global_resolution=1024, protected_global_resolution=768, |
| native_window=512, views=["I", "H", "V", "HV"], refiners=2, |
| query_gt_used=False, teacher_and_references_required=True, |
| inverse_map="native continuous logits -> independent sigmoid -> threshold 0.5", |
| any_output="independent canonical head, not a union of contact/crack") |
| row = self.engine.data.query(image) |
| classes = list(range(1, len(self.engine.data.spec["classes"]))) |
| total = None |
| for view in (0, 1, 2, 3): |
| q, evidence = self.engine.encode(row, classes, view) |
| for refiner in self.refiners: |
| logits = [] |
| for c, e in evidence.items(): |
| delta = refiner.canvas_delta(q["image"], e, q["geometry_valid"]) |
| z = native_logits((e[:1] + delta)[None], q["geometry"])[0] |
| logits.append(flips(z, view)) |
| z = torch.cat(logits, 0).float() / 8 |
| total = z if total is None else total + z |
| p = probabilities(total).cpu().numpy() |
| meta = dict(model="PV-PRISM", dataset=self.dataset, image_sha256=row["sha256"], |
| group_id=row["group_id"], classes=self.engine.data.spec["classes"], |
| global_resolution=1024, native_window=512, views=["I", "H", "V", "HV"], |
| refiners=2, geometry=q["geometry"].to_dict(), query_gt_used=False, |
| teacher_and_references_required=True, |
| inverse_map="native continuous logits -> softmax([0,foreground]) -> argmax") |
| return p, meta |
|
|
