Spaces:
Running on Zero
Running on Zero
File size: 2,797 Bytes
a269f8c | 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 |
import timm
import torch
import torch.nn as nn
EFFICIENTNET_DIM = 1280
class TemporalAttention(nn.Module):
"""
Input : (B, N_FRAMES, EFFICIENTNET_DIM)
Output : (pooled (B, EFFICIENTNET_DIM), attn_weights (B, N_FRAMES))
"""
def __init__(self, input_dim: int = EFFICIENTNET_DIM):
super().__init__()
self.attention = nn.Sequential(
nn.Linear(input_dim, 128),
nn.Tanh(),
nn.Linear(128, 1),
)
def forward(self, x: torch.Tensor):
scores = self.attention(x) # (B, 16, 1)
weights = torch.softmax(scores, dim=1) # (B, 16, 1)
pooled = (weights * x).sum(dim=1) # (B, 1280)
return pooled, weights.squeeze(-1) # (B, 1280), (B, 16)
class FallDetector(nn.Module):
"""
EfficientNet-Lite0 + Temporal Attention + MLP classifier.
Forward:
(B, 16, 3, 224, 224)
-> per-frame EfficientNet-Lite0 -> (B, 16, 1280)
-> TemporalAttention -> (B, 1280), (B, 16)
-> MLP classifier -> (B, 1) logit
Sigmoid is applied explicitly at inference, not inside the module.
"""
def __init__(self, pretrained_backbone: bool = False):
super().__init__()
backbone = timm.create_model("efficientnet_lite0", pretrained=pretrained_backbone)
self.conv_stem = backbone.conv_stem
self.bn1 = backbone.bn1
self.blocks = backbone.blocks
self.conv_head = backbone.conv_head
self.bn2 = backbone.bn2
self.global_pool = backbone.global_pool
self.temporal_attention = TemporalAttention(EFFICIENTNET_DIM)
self.classifier = nn.Sequential(
nn.Linear(EFFICIENTNET_DIM, 512),
nn.ReLU(inplace=True),
nn.Dropout(0.2),
nn.Linear(512, 128),
nn.ReLU(inplace=True),
nn.Dropout(0.2),
nn.Linear(128, 1),
)
def forward(self, x: torch.Tensor):
B, T, C, H, W = x.shape
x_flat = x.view(B * T, C, H, W)
f = self.conv_stem(x_flat)
f = self.bn1(f)
f = self.blocks(f)
f = self.conv_head(f)
f = self.bn2(f)
f = self.global_pool(f) # (B*16, 1280)
f = f.view(B, T, EFFICIENTNET_DIM) # (B, 16, 1280)
pooled, attn_weights = self.temporal_attention(f)
logit = self.classifier(pooled) # (B, 1)
return logit, attn_weights
def count_parameters(self):
total = sum(p.numel() for p in self.parameters())
trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
return total, trainable
|