"""Shared data / model code for the v2 study (strictly self-supervised MAE).""" import os, glob, random, xml.etree.ElementTree as ET import numpy as np, cv2, torch, torch.nn as nn, torch.nn.functional as F ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) DATA = os.path.join(ROOT, "NEU-DET-data", "NEU-DET") CLASSES = ["crazing", "inclusion", "patches", "pitted_surface", "rolled-in_scale", "scratches"] IMG, P = 224, 16 N = (IMG // P) ** 2 def set_seed(s): random.seed(s); np.random.seed(s); torch.manual_seed(s); torch.cuda.manual_seed_all(s) def load_split(split, cache=True): """Returns uint8 images (n,224,224) after CLAHE, labels (n,), boxes list (normalized xyxy). Labels/boxes are NEVER used in pretraining; only for evaluation.""" cp = os.path.join(ROOT, "v2", "cache", f"{split}.npz") if cache and os.path.exists(cp): d = np.load(cp, allow_pickle=True); return d["x"], d["y"], list(d["b"]) xs, ys, bs = [], [], [] clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)) for ci, c in enumerate(CLASSES): for f in sorted(glob.glob(os.path.join(DATA, split, "images", c, "*.jpg"))): im = cv2.imread(f, cv2.IMREAD_GRAYSCALE) h, w = im.shape im = cv2.resize(clahe.apply(im), (IMG, IMG), interpolation=cv2.INTER_LINEAR) stem = os.path.splitext(os.path.basename(f))[0] xp = glob.glob(os.path.join(DATA, "*", "annotations", stem + ".xml"))[0] # 1 xml is misfiled in NEU-DET r = ET.parse(xp).getroot() bx = [] for o in r.findall("object"): b = o.find("bndbox") x0, y0, x1, y1 = (float(b.find(t).text) for t in ("xmin", "ymin", "xmax", "ymax")) bx.append((x0 / w, y0 / h, x1 / w, y1 / h)) xs.append(im); ys.append(ci); bs.append(np.array(bx, dtype=np.float32)) x, y = np.stack(xs), np.array(ys) os.makedirs(os.path.dirname(cp), exist_ok=True) b = np.empty(len(bs), dtype=object); b[:] = bs np.savez(cp, x=x, y=y, b=b) return x, y, bs # ---------------- GPU augmentation (label-free, box-free) ---------------- def augment(x, strong=False): """x: (B,1,H,W) in [0,1] on GPU. Weak = flips/rot90/jitter/noise (as v1). Strong adds random-resized-crop, possible now that no boxes must be preserved.""" B = x.shape[0] if strong: scale = torch.empty(B, device=x.device).uniform_(0.35, 1.0).sqrt() tx = (torch.rand(B, device=x.device) * 2 - 1) * (1 - scale) ty = (torch.rand(B, device=x.device) * 2 - 1) * (1 - scale) th = torch.zeros(B, 2, 3, device=x.device) th[:, 0, 0] = scale; th[:, 1, 1] = scale; th[:, 0, 2] = tx; th[:, 1, 2] = ty x = F.grid_sample(x, F.affine_grid(th, x.shape, align_corners=False), align_corners=False) f = torch.rand(B, device=x.device) < 0.5; x[f] = x[f].flip(-1) f = torch.rand(B, device=x.device) < 0.5; x[f] = x[f].flip(-2) k = torch.randint(0, 4, (B,), device=x.device) x = torch.stack([torch.rot90(xi, int(ki), (-2, -1)) for xi, ki in zip(x, k)]) s = 0.4 if strong else 0.2 b = 1 + (torch.rand(B, 1, 1, 1, device=x.device) * 2 - 1) * s c = 1 + (torch.rand(B, 1, 1, 1, device=x.device) * 2 - 1) * s m = x.mean((2, 3), keepdim=True) x = ((x - m) * c + m) * b if strong: bl = torch.rand(B, device=x.device) < 0.3 if bl.any(): x[bl] = torch.from_numpy(np.stack([cv2.GaussianBlur(i[0].cpu().numpy(), (0, 0), random.uniform(0.3, 1.5)) for i in x[bl]]))[:, None].to(x.device) x = x + torch.randn_like(x) * 0.02 return x.clamp(0, 1) # ---------------- model (identical architecture to v1) ---------------- class PatchEmbed(nn.Module): def __init__(self, d=768): super().__init__(); self.proj = nn.Conv2d(1, d, P, P) def forward(self, x): return self.proj(x).flatten(2).transpose(1, 2) class Encoder(nn.Module): """ViT: 224x224 -> 196 patches, 6 pre-norm layers, 12 heads, 768-d. Same state-dict layout as v1.""" def __init__(self, d=768, depth=6, heads=12): super().__init__() self.patch_embed = PatchEmbed(d) self.cls_token = nn.Parameter(torch.zeros(1, 1, d)) self.pos_embed = nn.Parameter(torch.zeros(1, N + 1, d)) nn.init.trunc_normal_(self.pos_embed, std=.02); nn.init.trunc_normal_(self.cls_token, std=.02) layer = nn.TransformerEncoderLayer(d, heads, d * 4, dropout=0.0, batch_first=True, norm_first=True, activation="gelu") self.blocks = nn.TransformerEncoder(layer, depth, enable_nested_tensor=False) self.norm = nn.LayerNorm(d) self.apply(self._init) @staticmethod def _init(m): if isinstance(m, nn.Linear): nn.init.xavier_uniform_(m.weight) if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.LayerNorm): nn.init.ones_(m.weight); nn.init.zeros_(m.bias) def forward(self, x, ids_keep=None, return_layers=False): x = self.patch_embed(x) + self.pos_embed[:, 1:] if ids_keep is not None: x = torch.gather(x, 1, ids_keep[..., None].expand(-1, -1, x.shape[-1])) cls = (self.cls_token + self.pos_embed[:, :1]).expand(x.shape[0], -1, -1) x = torch.cat([cls, x], 1) outs = [] for blk in self.blocks.layers: x = blk(x); outs.append(x) x = self.norm(x) return (x, outs) if return_layers else x @torch.no_grad() def features(self, x, mode="cls"): z, outs = self.forward(x, return_layers=True) if mode == "mid": # mean-pooled patch tokens of blocks 4-6, each instance-normalised, averaged return torch.stack([F.layer_norm(o[:, 1:].float(), (o.shape[-1],)).mean(1) for o in outs[3:]]).mean(0) if mode == "cls": return z[:, 0] if mode == "mean": return z[:, 1:].mean(1) return torch.cat([z[:, 0], z[:, 1:].mean(1)], 1) class Decoder(nn.Module): def __init__(self, d=768, dd=256, depth=4, heads=8): super().__init__() self.embed = nn.Linear(d, dd) self.mask_token = nn.Parameter(torch.zeros(1, 1, dd)) self.pos_embed = nn.Parameter(torch.zeros(1, N + 1, dd)) nn.init.trunc_normal_(self.pos_embed, std=.02); nn.init.normal_(self.mask_token, std=.02) layer = nn.TransformerEncoderLayer(dd, heads, dd * 4, dropout=0.0, batch_first=True, norm_first=True, activation="gelu") self.blocks = nn.TransformerEncoder(layer, depth, enable_nested_tensor=False) self.norm = nn.LayerNorm(dd) self.pred = nn.Linear(dd, P * P) def forward(self, z, ids_restore): z = self.embed(z) B, L, D = z.shape mt = self.mask_token.expand(B, N + 1 - L, -1) x = torch.cat([z[:, 1:], mt], 1) x = torch.gather(x, 1, ids_restore[..., None].expand(-1, -1, D)) x = torch.cat([z[:, :1], x], 1) + self.pos_embed return self.pred(self.norm(self.blocks(x)))[:, 1:] def patchify(x): B = x.shape[0] return x.reshape(B, 1, 14, P, 14, P).permute(0, 2, 4, 1, 3, 5).reshape(B, N, P * P) def unpatchify(p): B = p.shape[0] return p.reshape(B, 14, 14, P, P).permute(0, 1, 3, 2, 4).reshape(B, 1, IMG, IMG) def random_masking(B, ratio, device): noise = torch.rand(B, N, device=device) ids_shuffle = noise.argsort(1); ids_restore = ids_shuffle.argsort(1) keep = N - int(ratio * N) ids_keep = ids_shuffle[:, :keep] mask = torch.ones(B, N, dtype=torch.bool, device=device); mask.scatter_(1, ids_keep, False) return ids_keep, ids_restore, mask def load_encoder(path, device="cuda"): enc = Encoder().to(device) s = torch.load(path, map_location=device, weights_only=False) if "encoder" in s: s = s["encoder"] enc.load_state_dict(s); enc.eval() return enc @torch.no_grad() def embed(enc, x_uint8, mode="cls", bs=64, device="cuda"): """mode '_tta': features averaged over the 8 dihedral transforms (4 rot90 x optional flip).""" enc.eval(); out = []; tta = mode.endswith("_tta"); base = mode[:-4] if tta else mode for i in range(0, len(x_uint8), bs): xb = torch.from_numpy(x_uint8[i:i + bs]).float().div(255)[:, None].to(device) with torch.autocast("cuda", dtype=torch.float16): if not tta: out.append(enc.features(xb, base).float().cpu()); continue f = 0 for fl in (False, True): for k in range(4): v = torch.rot90(xb.flip(-1) if fl else xb, k, (-2, -1)) f = f + enc.features(v, base).float() out.append((f / 8).cpu()) return torch.cat(out).numpy()