| import os |
| import io |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import numpy as np |
| from PIL import Image |
| from einops import rearrange |
|
|
| class DropPath(nn.Module): |
| def __init__(self, drop_prob=0.0): |
| super().__init__() |
| self.drop_prob = drop_prob |
| def forward(self, x): |
| if self.drop_prob == 0.0 or not self.training: |
| return x |
| keep_prob = 1 - self.drop_prob |
| shape = (x.shape[0],) + (1,) * (x.ndim - 1) |
| random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device) |
| random_tensor.floor_() |
| return x.div(keep_prob) * random_tensor |
|
|
| class PatchEmbed2D(nn.Module): |
| def __init__(self, img_size, patch_size, in_chans=1, embed_dim=384): |
| super().__init__() |
| self.patch_size = patch_size |
| self.num_patches = (img_size[0] // patch_size[0]) * (img_size[1] // patch_size[1]) |
| self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) |
| def forward(self, x): |
| x = self.proj(x) |
| return rearrange(x, 'b e h w -> b (h w) e') |
|
|
| class TransformerBlock(nn.Module): |
| def __init__(self, embed_dim, num_heads, mlp_ratio=4.0, dropout=0.0, drop_path=0.0): |
| super().__init__() |
| self.norm1 = nn.LayerNorm(embed_dim) |
| self.attn = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout, batch_first=True) |
| self.norm2 = nn.LayerNorm(embed_dim) |
| mlp_dim = int(embed_dim * mlp_ratio) |
| self.mlp = nn.Sequential( |
| nn.Linear(embed_dim, mlp_dim), nn.GELU(), nn.Dropout(dropout), |
| nn.Linear(mlp_dim, embed_dim), nn.Dropout(dropout) |
| ) |
| self.drop_path = DropPath(drop_path) |
| def forward(self, x): |
| x_norm = self.norm1(x) |
| attn_out, attn_weights = self.attn(x_norm, x_norm, x_norm, need_weights=True, average_attn_weights=False) |
| x = x + self.drop_path(attn_out) |
| x = x + self.drop_path(self.mlp(self.norm2(x))) |
| return x, attn_weights |
|
|
| class MultiScaleViT2DEncoder(nn.Module): |
| def __init__(self, img_size, patch_size, embed_dim=384, depth=12, num_heads=6, |
| mlp_ratio=4.0, dropout=0.0, drop_path_rate=0.1): |
| super().__init__() |
| self.patch_embed = PatchEmbed2D(img_size, patch_size, 1, embed_dim) |
| num_patches = self.patch_embed.num_patches |
| self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) |
| self.pos_embed = nn.Parameter(torch.zeros(1, num_patches+1, embed_dim)) |
| dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] |
| self.blocks = nn.ModuleList([ |
| TransformerBlock(embed_dim, num_heads, mlp_ratio, dropout, dpr[i]) for i in range(depth) |
| ]) |
| self.norm = nn.LayerNorm(embed_dim) |
| self.scale_indices = [2, 5, 8, 11] |
| self.scale_fusion = nn.Sequential( |
| nn.Linear(embed_dim * len(self.scale_indices), embed_dim), nn.GELU(), nn.LayerNorm(embed_dim) |
| ) |
| nn.init.trunc_normal_(self.pos_embed, std=0.02) |
| nn.init.trunc_normal_(self.cls_token, std=0.02) |
| def forward(self, x, mask=None, return_all_tokens=False): |
| B = x.size(0) |
| tokens = self.patch_embed(x) |
| tokens = tokens + self.pos_embed[:, 1:tokens.shape[1]+1] |
| if mask is not None: |
| tokens = tokens[mask].reshape(B, -1, tokens.size(-1)) |
| cls = self.cls_token.expand(B, -1, -1) |
| tokens = torch.cat([cls, tokens], dim=1) |
| tokens = tokens + self.pos_embed[:, :tokens.shape[1]] |
| intermediate_features = [] |
| attn_maps = [] |
| for i, block in enumerate(self.blocks): |
| tokens, attn = block(tokens) |
| if i in self.scale_indices: |
| intermediate_features.append(tokens[:, 1:]) |
| attn_maps.append(attn) |
| tokens = self.norm(tokens) |
| if len(intermediate_features) > 0: |
| pooled_scales = [f.mean(dim=1) for f in intermediate_features] |
| multi_scale_vec = torch.cat(pooled_scales, dim=-1) |
| multi_scale_vec = self.scale_fusion(multi_scale_vec) |
| else: |
| multi_scale_vec = tokens[:, 0] |
| if return_all_tokens: |
| return tokens, multi_scale_vec, attn_maps |
| return tokens |
|
|
| class SecondOrderClassifier(nn.Module): |
| def __init__(self, encoder, num_classes=2, dropout=0.3, use_second_order=True): |
| super().__init__() |
| self.encoder = encoder |
| self.use_second_order = use_second_order |
| embed_dim = encoder.blocks[0].norm1.normalized_shape[0] |
| if use_second_order: |
| self.cov_proj = nn.Sequential( |
| nn.Linear(embed_dim * embed_dim, embed_dim * 4), nn.GELU(), nn.Dropout(dropout), |
| nn.Linear(embed_dim * 4, embed_dim) |
| ) |
| self.head = nn.Sequential( |
| nn.LayerNorm(embed_dim * 2), nn.Dropout(dropout), nn.Linear(embed_dim * 2, num_classes) |
| ) |
| else: |
| self.head = nn.Sequential(nn.LayerNorm(embed_dim), nn.Dropout(dropout), nn.Linear(embed_dim, num_classes)) |
| self.mc_dropout = nn.Dropout(dropout) |
| def forward(self, x, return_features=False, mc_dropout=False): |
| tokens, multi_scale, _ = self.encoder(x, return_all_tokens=True) |
| cls = tokens[:, 0] |
| if self.use_second_order: |
| patches = tokens[:, 1:] |
| mean_patches = patches.mean(dim=1, keepdim=True) |
| centered = patches - mean_patches |
| cov = torch.bmm(centered.transpose(1,2), centered) / centered.size(1) |
| cov_vec = cov.reshape(cov.size(0), -1) |
| second_order = self.cov_proj(cov_vec) |
| second_order = F.normalize(second_order, p=2, dim=-1) |
| features = torch.cat([cls, second_order], dim=-1) |
| else: |
| features = cls |
| if mc_dropout: |
| features = self.mc_dropout(features) |
| logits = self.head(features) |
| if return_features: |
| return logits, features, cls |
| return logits |
|
|
| def get_model(checkpoint_path, device): |
| """Load model once and transfer to target device.""" |
| print(f"Loading checkpoint from {checkpoint_path}...") |
| ckpt = torch.load(checkpoint_path, map_location='cpu', weights_only=False) |
| encoder = MultiScaleViT2DEncoder((224, 224), (16, 16), 384, 12, 6, drop_path_rate=0.0) |
| model = SecondOrderClassifier(encoder, num_classes=2, dropout=0.3, use_second_order=True) |
| model.load_state_dict(ckpt['model_state']) |
| model = model.to(device).eval() |
| print("Model loaded successfully!") |
| return model |
|
|
| def preprocess_image(image_path_or_bytes, target_size=(224, 224)): |
| """Convert and normalize image to Z-score tensor.""" |
| if isinstance(image_path_or_bytes, (str, bytes)) and isinstance(image_path_or_bytes, str): |
| img = Image.open(image_path_or_bytes).convert('L') |
| else: |
| img = Image.open(io.BytesIO(image_path_or_bytes)).convert('L') |
| |
| img = img.resize((target_size[1], target_size[0]), Image.BILINEAR) |
| tensor = torch.from_numpy(np.array(img, dtype=np.float32)).unsqueeze(0) / 255.0 |
| mu = tensor.mean() |
| std = tensor.std() + 1e-8 |
| tensor = (tensor - mu) / std |
| return tensor.unsqueeze(0), img |
|
|
| def attention_rollout(model, x, device): |
| """Compute the vision transformer attention maps using rollout.""" |
| model.eval() |
| with torch.no_grad(): |
| _, _, attn_maps = model.encoder(x.to(device), return_all_tokens=True) |
| B = x.size(0) |
| result = torch.eye(attn_maps[0].size(-1), device=device).unsqueeze(0).repeat(B, 1, 1) |
| for attn in attn_maps: |
| attn = attn.mean(dim=1) |
| attn = attn + torch.eye(attn.size(-1), device=device).unsqueeze(0) |
| attn = attn / attn.sum(dim=-1, keepdim=True) |
| result = torch.bmm(result, attn) |
| mask = result[:, 0, 1:] |
| H = W = int(np.sqrt(mask.size(1))) |
| mask = mask.reshape(B, H, W) |
| mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8) |
| return mask |
|
|
| def predict_with_uncertainty(model, tensor, device, temperature=1.6995, mc_samples=10): |
| """Run model inference, MC Dropout uncertainty estimation, and attention rollout.""" |
| tensor = tensor.to(device) |
| |
| |
| model.eval() |
| with torch.no_grad(): |
| logits = model(tensor) |
| raw_probs = torch.softmax(logits, dim=1) |
| calibrated_logits = logits / temperature |
| calibrated_probs = torch.softmax(calibrated_logits, dim=1) |
| |
| |
| mc_probs = [] |
| if mc_samples > 0: |
| model.train() |
| with torch.no_grad(): |
| for _ in range(mc_samples): |
| logit = model(tensor, mc_dropout=True) |
| prob = torch.softmax(logit, dim=1)[:, 1] |
| mc_probs.append(prob.item()) |
| model.eval() |
| mc_std = np.std(mc_probs) |
| else: |
| mc_std = 0.0 |
| |
| |
| attn_mask = attention_rollout(model, tensor, device) |
| attn_mask_np = attn_mask[0].cpu().numpy() |
| |
| return { |
| 'raw_probability': float(raw_probs[0, 1]), |
| 'calibrated_probability': float(calibrated_probs[0, 1]), |
| 'uncertainty': float(mc_std), |
| 'attention_map': attn_mask_np.tolist() |
| } |
|
|
