Spaces:
Sleeping
Sleeping
| """ | |
| MultiSense-DF — Visual Branch | |
| EfficientNet-B4 CNN backbone + Temporal Transformer | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| from einops import rearrange | |
| import timm | |
| class TemporalTransformer(nn.Module): | |
| """6-layer Transformer over frame-level embeddings.""" | |
| def __init__(self, embed_dim=512, num_heads=8, num_layers=6, | |
| num_frames=125, dropout=0.1): | |
| super().__init__() | |
| self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) | |
| self.pos_embed = nn.Parameter(torch.zeros(1, num_frames + 1, embed_dim)) | |
| encoder_layer = nn.TransformerEncoderLayer( | |
| d_model=embed_dim, nhead=num_heads, | |
| dim_feedforward=embed_dim * 4, | |
| dropout=dropout, batch_first=True, norm_first=True | |
| ) | |
| self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers) | |
| self.norm = 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): | |
| # x: (B, T, D) | |
| B = x.size(0) | |
| cls = self.cls_token.expand(B, -1, -1) | |
| x = torch.cat([cls, x], dim=1) # (B, T+1, D) | |
| x = x + self.pos_embed[:, :x.size(1)] | |
| x = self.transformer(x) | |
| x = self.norm(x) | |
| return x[:, 0] # CLS token → (B, D) | |
| class VisualBranch(nn.Module): | |
| """ | |
| EfficientNet-B4 frame encoder + TemporalTransformer. | |
| Input : (B, T, 3, 224, 224) video clips | |
| Output: (B, 512) visual embedding + (B, 1) logit | |
| """ | |
| def __init__(self, embed_dim=512, num_frames=125, | |
| num_heads=8, num_layers=6, dropout=0.1, | |
| freeze_backbone_epochs=5): | |
| super().__init__() | |
| self.freeze_backbone_epochs = freeze_backbone_epochs | |
| # EfficientNet-B4 backbone (1792-d features) | |
| self.backbone = timm.create_model( | |
| 'efficientnet_b4', pretrained=True, num_classes=0 | |
| ) | |
| backbone_dim = self.backbone.num_features # 1792 | |
| # Project backbone → embed_dim | |
| self.proj = nn.Sequential( | |
| nn.Linear(backbone_dim, embed_dim), | |
| nn.LayerNorm(embed_dim), | |
| nn.GELU() | |
| ) | |
| # Temporal Transformer | |
| self.temporal = TemporalTransformer( | |
| embed_dim=embed_dim, | |
| num_heads=num_heads, | |
| num_layers=num_layers, | |
| num_frames=num_frames, | |
| dropout=dropout | |
| ) | |
| # Classification head | |
| self.classifier = nn.Sequential( | |
| nn.Linear(embed_dim, 256), | |
| nn.GELU(), | |
| nn.Dropout(dropout), | |
| nn.Linear(256, 1) | |
| ) | |
| def freeze_backbone(self): | |
| for p in self.backbone.parameters(): | |
| p.requires_grad = False | |
| def unfreeze_backbone(self): | |
| for p in self.backbone.parameters(): | |
| p.requires_grad = True | |
| def forward(self, x): | |
| # x: (B, T, C, H, W) | |
| B, T, C, H, W = x.shape | |
| frames = rearrange(x, 'b t c h w -> (b t) c h w') | |
| # Extract per-frame features | |
| feats = self.backbone(frames) # (B*T, 1792) | |
| feats = self.proj(feats) # (B*T, 512) | |
| feats = rearrange(feats, '(b t) d -> b t d', b=B, t=T) | |
| # Temporal aggregation | |
| embed = self.temporal(feats) # (B, 512) | |
| logit = self.classifier(embed) # (B, 1) | |
| return embed, logit | |
| if __name__ == '__main__': | |
| model = VisualBranch(num_frames=125) | |
| dummy = torch.randn(2, 125, 3, 224, 224) | |
| emb, logit = model(dummy) | |
| print(f'Visual embed: {emb.shape}, logit: {logit.shape}') | |