| """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] |
| 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 |
|
|
|
|
| |
| 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) |
|
|
|
|
| |
| 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": |
| 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 '<base>_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() |
|
|
