File size: 8,073 Bytes
acb2ce3 0d8aa44 acb2ce3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 | import torch
import torch.nn as nn
class LTFCore(nn.Module):
"""
Core module.
Input: x (B, N, T, D)
Output: logits (B, num_classes), token_scores (B, N, T-1, 1), jpeg_score (B, 1)
"""
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,
):
super().__init__()
self.use_global_image = use_global_image
self.use_patch_transformer = use_patch_transformer
# Learnable CLS for tile-level aggregation
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
# Tile-level encoder (across tiles)
encoder_layer = nn.TransformerEncoderLayer(
d_model=embed_dim,
dropout=dropout,
nhead=tile_heads,
batch_first=True
)
self.tile_transformer = nn.TransformerEncoder(encoder_layer, num_layers=tile_layers)
if use_patch_transformer:
token_layer = nn.TransformerEncoderLayer(
d_model=embed_dim,
dropout=dropout,
nhead=num_heads,
batch_first=True
)
self.patch_transformer = nn.TransformerEncoder(token_layer, num_layers=num_layers)
else:
self.patch_transformer = None
# Image-level classifier: with global (main tile + aggregated tiles) or only aggregated tiles (ablation)
classifier_in_dim = embed_dim * 2 if use_global_image else embed_dim
if linear_classifier:
self.classifier = nn.Linear(classifier_in_dim, num_classes)
else:
self.classifier = nn.Sequential(
nn.Linear(classifier_in_dim, 512),
nn.ReLU(),
nn.Dropout(p=dropout),
nn.Linear(512, num_classes)
)
# Per-token head (for non-CLS tokens)
self.patch_classifier = nn.Sequential(
nn.Linear(embed_dim, 512),
nn.ReLU(),
nn.Dropout(p=dropout),
nn.Linear(512, 1)
)
# Scalar head (e.g., JPEG-related score)
self.jpeg_estimate = nn.Sequential(
nn.Linear(embed_dim, 256),
nn.ReLU(),
nn.Linear(256, 128),
nn.ReLU(),
nn.Linear(128, 128),
nn.ReLU(),
nn.Linear(128, 1),
)
nn.init.trunc_normal_(self.cls_token, std=0.02)
def forward(self, x):
"""
x: (B, N, T, D)
returns: (B, num_classes), (B, N, T-1, 1), (B, 1)
"""
B_2, N_2, T_2, D_2 = x.shape
# Token-level encoding per tile (optional)
x = x.reshape(B_2*N_2, T_2, D_2) # [B*N, T, D]
if self.patch_transformer is not None:
x = self.patch_transformer(x)
x = x.reshape(B_2, N_2, T_2, D_2)
# Split CLS vs non-CLS tokens
new_cls = x[:, :, 0, :]
new_x = x[:, :, 1:, :]
# Per-token head on non-CLS tokens
x_flat = new_x.reshape(-1, D_2) # [B*N*T, D]
x_result = self.patch_classifier(x_flat) # [B*N*(T-1), 1]
x_result = x_result.reshape(B_2, N_2, T_2 - 1, 1) # [B, N, T-1, 1]
main_feat = new_cls[:, 0] # (B, D)
tile_feats = new_cls[:, 1:] # (B, N-1, D)
# print(main_feat.shape, tile_feats.shape)
# Tile-level aggregation with a learnable CLS
cls_token = self.cls_token.expand(B_2, 1, -1) # (B, 1, D)
feature_tile = torch.cat([cls_token, tile_feats], dim=1) # (B, N, D)
feature_encoded = self.tile_transformer(feature_tile) # (B, N, D)
cls_out = feature_encoded[:, 0] # (B, D)
# Image-level classification: with or without global (main tile) feature (ablation v1_1)
if self.use_global_image:
concat_feat = torch.cat([main_feat, cls_out], dim=1) # (B, 2D)
else:
concat_feat = cls_out # (B, D)
return self.classifier(concat_feat), x_result, torch.sigmoid(self.jpeg_estimate(cls_out)) # (B, num_classes) (B)
class LTFNet(nn.Module):
"""
Wrapper with a frozen visual backbone (e.g., CLIP visual) + LTFCore.
Input: images x (B, N, 3, H, W)
Output: same as LTFCore
"""
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,
):
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
# Token feature dim = transformer hidden size (ViT-L -> 1024). Use conv1 out_channels;
# do NOT use output_dim which is the projection-head dim (often 768) and would cause
# "expecting 768 but got 1024" in patch_transformer.
self.embed_dim = self.visual.conv1.weight.shape[0]
self.core = LTFCore(
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,
)
def get_tokens(self, x) -> torch.Tensor:
"""
Extract token sequence (with CLS) from the visual backbone.
Assumes CLIP ViT structure: conv1 -> patch embed -> class_embedding + pos_embed -> ln_pre -> transformer.
x: (B, 3, H, W) -> tokens: (B, T, D)
"""
x = x.type(self.dtype)
x = self.visual.conv1(x) # [B, C, H, W]
x = x.reshape(x.shape[0], x.shape[1], -1) # [B, C, HW]
x = x.permute(0, 2, 1) # [B, HW, C]
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) # [B, 1+HW, C]
x = x + self.visual.positional_embedding.to(x.dtype)
x = self.visual.ln_pre(x)
x = x.permute(1, 0, 2) # [sequence_len, B, C]
x = self.visual.transformer(x)
x = x.permute(1, 0, 2) # [B, sequence_len, C]
return x
def encode_image(self, image) -> torch.Tensor:
# image: (B, 3, H, W) -> feature: (B, D)
return self.visual(image.type(self.dtype))
def forward(self, x: torch.Tensor, chunk_size=256):
# Tensor[B, 5, 3, 224, 224]
"""
x: Tensor of shape (B, N, 3, H, W)
chunk_size: number of tiles to pass through the visual backbone per batch. Same tiles are
computed and concatenated; only speed/memory change, no effect on outputs.
Eval can set model.eval_chunk_size (e.g. 512) in config for faster run.
"""
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 = []
with torch.no_grad():
for start in range(0, imgs.size(0), chunk_size):
end = min(start + chunk_size, imgs.size(0))
feats_chunk = self.get_tokens(imgs[start:end]) # (chunk, T, D)
feats_chunks.append(feats_chunk)
feats_token = torch.cat(feats_chunks, dim=0)
feats = feats_token.reshape(B, N, feats_token.shape[-2], -1) # (B, N, T, D)
feats = feats.float()
return self.core(feats) # logits: (B, num_classes) (B) |