""" LTF V2: multi-branch (baseline + Level-1 MIL on patch tokens + Level-2 MIL on CLS + fusion). V1 code in LTF.py is unchanged; this module is a separate copy for version V2. """ from typing import Optional, Tuple import torch import torch.nn as nn from model.LTF import LTFCore from mil_module.mil_core import MILCore def _clip_vit_resblock_modules(transformer) -> list: """Return the OpenAI CLIP ViT residual blocks as a list.""" rb = transformer.resblocks if isinstance(rb, nn.Sequential): return list(rb.children()) return list(rb) def _count_clip_vit_depth(visual) -> int: return len(_clip_vit_resblock_modules(visual.transformer)) class BagTokenShuffle(nn.Module): """Randomly permute tokens within each bag during training with probability p_shuffle.""" def forward(self, x, p_shuffle: float = 0.0, p_drop: float = 0.0, key_padding_mask=None, **kwargs): del p_drop, key_padding_mask, kwargs if (not self.training) or p_shuffle <= 0: return x if float(torch.rand((), device=x.device)) >= float(p_shuffle): return x B, N, D = x.shape if N <= 1: return x perm = torch.rand(B, N, device=x.device).argsort(dim=1) return x.gather(1, perm.unsqueeze(-1).expand(B, N, D)) class LTFCoreV2(nn.Module): """Baseline and hierarchical MIL branches with shared frozen CLIP features. Level 1 pools patch tokens within local tiles and aggregates tile bags with a learned query. Level 2 pools CLS tokens across views. Cross-Shuffle is applied only during training. """ def __init__( self, embed_dim=1024, num_heads=8, num_classes=2, num_layers=1, dropout=0.1, tile_heads=2, tile_layers=2, use_global_image=True, use_patch_transformer=True, linear_classifier=False, mil_heads=8, fusion_hidden=512, l1_cross_shuffle_p: float = 0.9, l1_replace_ratio_range: Tuple[float, float] = (0.5, 0.9), l1_with_replacement_p: float = 0.5, # Level-2 CrossTokenShuffle l2_cross_shuffle_p: float = 0.5, l2_replace_ratio_range: Tuple[float, float] = (0.1, 0.5), l2_with_replacement_p: float = 0.5, bag_token_shuffle_p: float = 0.5, cross_shuffle_min_keep_original: int = 0, ): super().__init__() self.embed_dim = embed_dim self.use_global_image = use_global_image self.num_classes = num_classes self.l1_cross_shuffle_p = float(l1_cross_shuffle_p) self.l1_replace_ratio_range = (float(l1_replace_ratio_range[0]), float(l1_replace_ratio_range[1])) self.l1_with_replacement_p = float(l1_with_replacement_p) self.l2_cross_shuffle_p = float(l2_cross_shuffle_p) self.l2_replace_ratio_range = (float(l2_replace_ratio_range[0]), float(l2_replace_ratio_range[1])) self.l2_with_replacement_p = float(l2_with_replacement_p) self.bag_token_shuffle_p = float(bag_token_shuffle_p) self.cross_shuffle_min_keep_original = int(cross_shuffle_min_keep_original) self.baseline_core = LTFCore( embed_dim=embed_dim, num_heads=num_heads, num_classes=num_classes, num_layers=num_layers, dropout=dropout, tile_heads=tile_heads, tile_layers=tile_layers, use_global_image=use_global_image, use_patch_transformer=use_patch_transformer, linear_classifier=linear_classifier, ) mil_kw = dict( attn_hidden=min(256, embed_dim), attn_dropout=dropout, attn_bag_mode="embed_cls", ) self.mil_level1 = MILCore( d_in=embed_dim, num_classes=num_classes, mil_head_pooling="attn", mil_head_kwargs=mil_kw, use_cross_shuffle=True, trans=BagTokenShuffle(), allow_trans_token_drop=False, ) self.mil_level2 = MILCore( d_in=embed_dim, num_classes=num_classes, mil_head_pooling="attn", mil_head_kwargs=mil_kw, use_cross_shuffle=True, trans=BagTokenShuffle(), allow_trans_token_drop=False, ) self.level1_query = nn.Parameter(torch.zeros(1, 1, embed_dim)) mh = mil_heads if embed_dim % mil_heads == 0 else 8 while embed_dim % mh != 0: mh -= 1 if mh < 1: mh = 1 self.level1_aggregate_attn = nn.MultiheadAttention( embed_dim, mh, dropout=dropout, batch_first=True ) self.head_level1 = nn.Linear(embed_dim, num_classes) self.head_level2 = nn.Linear(embed_dim, num_classes) classifier_in_dim = embed_dim * 2 if use_global_image else embed_dim fusion_in = embed_dim + embed_dim + classifier_in_dim self.fusion_head = nn.Sequential( nn.Linear(fusion_in, fusion_hidden), nn.ReLU(), nn.Dropout(p=dropout), nn.Linear(fusion_hidden, num_classes), ) nn.init.trunc_normal_(self.level1_query, std=0.02) @staticmethod def _sample_replace_ratio(low: float, high: float, device: torch.device) -> float: lo, hi = min(low, high), max(low, high) lo = max(0.0, lo) hi = min(1.0, hi) if hi <= lo: return float(lo) u = torch.rand((), device=device) return float(lo + (hi - lo) * u) def _apply_patch_transformer(self, x): """x: (B, N, T, D) -> same shape after optional token-level encoder.""" B_2, N_2, T_2, D_2 = x.shape xf = x.reshape(B_2 * N_2, T_2, D_2) if self.baseline_core.patch_transformer is not None: xf = self.baseline_core.patch_transformer(xf) return xf.reshape(B_2, N_2, T_2, D_2) def _forward_baseline_from_xfeat(self, x_feat): """Run LTF classifier heads on processed tokens (same as LTFCore tail).""" B_2, N_2, T_2, D_2 = x_feat.shape bc = self.baseline_core new_cls = x_feat[:, :, 0, :] new_x = x_feat[:, :, 1:, :] x_flat = new_x.reshape(-1, D_2) x_result = bc.patch_classifier(x_flat) x_result = x_result.reshape(B_2, N_2, T_2 - 1, 1) main_feat = new_cls[:, 0] tile_feats = new_cls[:, 1:] cls_token = bc.cls_token.expand(B_2, 1, -1) feature_tile = torch.cat([cls_token, tile_feats], dim=1) feature_encoded = bc.tile_transformer(feature_tile) cls_out = feature_encoded[:, 0] if bc.use_global_image: concat_feat = torch.cat([main_feat, cls_out], dim=1) else: concat_feat = cls_out logits_b = bc.classifier(concat_feat) jpeg_b = torch.sigmoid(bc.jpeg_estimate(cls_out)) return logits_b, x_result, jpeg_b, concat_feat def forward( self, x, labels: Optional[torch.Tensor] = None, x_mil: Optional[torch.Tensor] = None, ): """Run the branches on token features shaped [B, N, T, D]. N must include the global view and at least one local tile. Optional x_mil provides alternate MIL tokens with the same shape. Binary labels enable training-time Cross-Shuffle; the baseline uses the original x features. """ x_feat = self._apply_patch_transformer(x) if x_mil is not None: if x_mil.shape != x.shape: raise ValueError(f"x_mil shape {tuple(x_mil.shape)} != x shape {tuple(x.shape)}") x_mil_feat = self._apply_patch_transformer(x_mil) else: x_mil_feat = x_feat logits_baseline, patch_logits, jpeg_pred, concat_feat = self._forward_baseline_from_xfeat( x_feat ) B, N, T, D = x_mil_feat.shape if N < 2: raise ValueError(f"LTFCoreV2 expects N>=2 (global + tiles), got N={N}") Nt = N - 1 train_shuffle = self.training and labels is not None p_bag_shuffle = self.bag_token_shuffle_p if train_shuffle else 0.0 patch_only = x_mil_feat[:, 1:, 1:, :].reshape(B * Nt, T - 1, D) do1 = ( train_shuffle and float(torch.rand((), device=x.device)) < self.l1_cross_shuffle_p ) rr1 = self._sample_replace_ratio(*self.l1_replace_ratio_range, x.device) wr1 = bool(float(torch.rand((), device=x.device)) < self.l1_with_replacement_p) y1 = labels.unsqueeze(1).expand(-1, Nt).reshape(-1) if labels is not None else None mk = self.cross_shuffle_min_keep_original mil1 = self.mil_level1( patch_only, y=y1, return_dict=True, do_cross_shuffle=do1, replace_ratio=rr1, with_replacement=wr1, detach_inserted=True, cross_shuffle_min_keep=mk, p_shuffle=p_bag_shuffle, ) be = mil1["mil_head"]["bag_embed"] if be is None: raise RuntimeError("MIL level1 expected bag_embed (use attn + embed_cls).") bag_seq = be.reshape(B, Nt, D) q = self.level1_query.expand(B, -1, -1) agg_out, _ = self.level1_aggregate_attn(q, bag_seq, bag_seq) level1_emb = agg_out.squeeze(1) cls_seq = x_mil_feat[:, :, 0, :] do2 = ( train_shuffle and float(torch.rand((), device=x.device)) < self.l2_cross_shuffle_p ) rr2 = self._sample_replace_ratio(*self.l2_replace_ratio_range, x.device) wr2 = bool(float(torch.rand((), device=x.device)) < self.l2_with_replacement_p) mil2 = self.mil_level2( cls_seq, y=labels, return_dict=True, do_cross_shuffle=do2, replace_ratio=rr2, with_replacement=wr2, detach_inserted=True, cross_shuffle_min_keep=mk, p_shuffle=p_bag_shuffle, ) be2 = mil2["mil_head"]["bag_embed"] if be2 is None: raise RuntimeError("MIL level2 expected bag_embed (use attn + embed_cls).") level2_emb = be2 logits_level1 = self.head_level1(level1_emb) logits_level2 = self.head_level2(level2_emb) fusion_in = torch.cat([level1_emb, level2_emb, concat_feat], dim=-1) logits_fusion = self.fusion_head(fusion_in) return { "logits_baseline": logits_baseline, "logits_level1": logits_level1, "logits_level2": logits_level2, "logits_fusion": logits_fusion, "patch_logits": patch_logits, "jpeg_pred": jpeg_pred, } class LTFNetV2(nn.Module): """Frozen CLIP visual + LTFCoreV2. Input: (B, N, 3, H, W).""" def __init__( self, clip_visual, num_heads=8, num_classes=2, num_layers=1, dropout=0.2, tile_heads=2, tile_layers=2, use_global_image=True, use_patch_transformer=True, linear_classifier=False, mil_heads=8, fusion_hidden=512, l1_cross_shuffle_p: float = 0.9, l1_replace_ratio_range: Tuple[float, float] = (0.5, 0.9), l1_with_replacement_p: float = 0.5, l2_cross_shuffle_p: float = 0.5, l2_replace_ratio_range: Tuple[float, float] = (0.1, 0.5), l2_with_replacement_p: float = 0.5, bag_token_shuffle_p: float = 0.5, cross_shuffle_min_keep_original: int = 0, use_clip_intermediate_for_mil: bool = False, clip_mil_stop_after_block_1based: int = 12, ): super().__init__() self.visual = clip_visual self.dtype = self.visual.conv1.weight.dtype for p in self.visual.parameters(): p.requires_grad = False self.visual.eval() self.input_resolution = self.visual.input_resolution self.embed_dim = self.visual.conv1.weight.shape[0] self.use_clip_intermediate_for_mil = bool(use_clip_intermediate_for_mil) self.clip_mil_stop_after_block_1based = int(clip_mil_stop_after_block_1based) if self.use_clip_intermediate_for_mil: n_bl = _count_clip_vit_depth(self.visual) stop = self.clip_mil_stop_after_block_1based if stop < 1 or stop > n_bl: raise ValueError( f"use_clip_intermediate_for_mil requires 1 <= v2_clip_mil_stop_after_block_1based " f"<= {n_bl} (ViT depth), got {stop}" ) self.core = LTFCoreV2( embed_dim=self.embed_dim, num_heads=num_heads, num_classes=num_classes, num_layers=num_layers, dropout=dropout, tile_heads=tile_heads, tile_layers=tile_layers, use_global_image=use_global_image, use_patch_transformer=use_patch_transformer, linear_classifier=linear_classifier, mil_heads=mil_heads, fusion_hidden=fusion_hidden, l1_cross_shuffle_p=l1_cross_shuffle_p, l1_replace_ratio_range=l1_replace_ratio_range, l1_with_replacement_p=l1_with_replacement_p, l2_cross_shuffle_p=l2_cross_shuffle_p, l2_replace_ratio_range=l2_replace_ratio_range, l2_with_replacement_p=l2_with_replacement_p, bag_token_shuffle_p=bag_token_shuffle_p, cross_shuffle_min_keep_original=cross_shuffle_min_keep_original, ) def _stem_to_lnd(self, x_bchw: torch.Tensor) -> torch.Tensor: """Convert images shaped [B, 3, H, W] to CLIP transformer input [L, N, D].""" x = x_bchw.type(self.dtype) x = self.visual.conv1(x) x = x.reshape(x.shape[0], x.shape[1], -1) x = x.permute(0, 2, 1) cls = self.visual.class_embedding.to(x.dtype) + torch.zeros( x.shape[0], 1, x.shape[2], dtype=x.dtype, device=x.device ) x = torch.cat([cls, x], dim=1) x = x + self.visual.positional_embedding.to(x.dtype) x = self.visual.ln_pre(x) return x.permute(1, 0, 2) def _visual_tokens_from_imgs( self, imgs_bchw: torch.Tensor ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """Return final-layer tokens and optional intermediate-layer MIL tokens. The second result is None when intermediate MIL features are disabled. """ x = self._stem_to_lnd(imgs_bchw) if not self.use_clip_intermediate_for_mil: x = self.visual.transformer(x) return x.permute(1, 0, 2), None blocks = _clip_vit_resblock_modules(self.visual.transformer) stop = self.clip_mil_stop_after_block_1based for i in range(stop): x = blocks[i](x) x_mil = x.permute(1, 0, 2).clone() for i in range(stop, len(blocks)): x = blocks[i](x) return x.permute(1, 0, 2), x_mil def get_tokens(self, x) -> torch.Tensor: """Return final ViT tokens shaped [B, T, D], without separating intermediate features.""" x = self._stem_to_lnd(x) x = self.visual.transformer(x) return x.permute(1, 0, 2) def encode_image(self, image) -> torch.Tensor: return self.visual(image.type(self.dtype)) def forward(self, x: torch.Tensor, labels: Optional[torch.Tensor] = None, chunk_size=256): chunk_size = getattr(self, "eval_chunk_size", chunk_size) B, N, C, H, W = x.shape imgs = x.reshape(B * N, C, H, W) feats_chunks = [] mil_chunks = [] with torch.no_grad(): for start in range(0, imgs.size(0), chunk_size): end = min(start + chunk_size, imgs.size(0)) f_tok, m_tok = self._visual_tokens_from_imgs(imgs[start:end]) feats_chunks.append(f_tok) if m_tok is not None: mil_chunks.append(m_tok) feats_token = torch.cat(feats_chunks, dim=0) T = feats_token.shape[-2] D = feats_token.shape[-1] feats = feats_token.reshape(B, N, T, D).float() if not mil_chunks: return self.core(feats, labels=labels, x_mil=None) mil_token = torch.cat(mil_chunks, dim=0).reshape(B, N, T, D).float() return self.core(feats, labels=labels, x_mil=mil_token)