| import torch.nn as nn |
| import torch |
| import torchvision.models as models |
| import torch.nn.functional as F |
|
|
| class SignBartMod(nn.Module): |
| def __init__(self, num_classes, num_keypoints, num_face_landmarks, num_emotions, |
| d_model=256, nhead=16, num_layers=2, |
| dim_feedforward=1024, dropout=0.1): |
| super().__init__() |
|
|
| self.n_pose = 7 |
| self.n_pose_feat = 7 |
| self.n_hand = 21 |
| self.n_face = num_face_landmarks |
|
|
| self.proj_pose_x = nn.Linear(self.n_pose_feat, d_model) |
| self.proj_lhand_x = nn.Linear(self.n_hand, d_model) |
| self.proj_rhand_x = nn.Linear(self.n_hand, d_model) |
| self.proj_face_x = nn.Linear(self.n_face, d_model) |
|
|
| self.proj_pose_y = nn.Linear(self.n_pose_feat, d_model) |
| self.proj_lhand_y = nn.Linear(self.n_hand, d_model) |
| self.proj_rhand_y = nn.Linear(self.n_hand, d_model) |
| self.proj_face_y = nn.Linear(self.n_face, d_model) |
|
|
| self.fusion_x = nn.Linear(4 * d_model, d_model) |
| self.fusion_y = nn.Linear(4 * d_model, d_model) |
| self.norm_x = nn.LayerNorm(d_model) |
| self.norm_y = nn.LayerNorm(d_model) |
|
|
| self.face_x_projector = nn.Linear(self.n_face, d_model) |
| self.face_y_projector = nn.Linear(self.n_face, d_model) |
| self.pos_encoder = nn.Parameter(torch.randn(1, 500, d_model)) |
|
|
| sign_enc_layer = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout, batch_first=True) |
| self.encoder = nn.TransformerEncoder(sign_enc_layer, num_layers) |
|
|
| sign_dec_layer = nn.TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout, batch_first=True) |
| self.decoder = nn.TransformerDecoder(sign_dec_layer, num_layers) |
|
|
| emo_enc_layer = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout, batch_first=True) |
| self.emotion_encoder = nn.TransformerEncoder(emo_enc_layer, num_layers=1) |
|
|
| emo_dec_layer = nn.TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout, batch_first=True) |
| self.emotion_decoder = nn.TransformerDecoder(emo_dec_layer, num_layers=1) |
|
|
| self.classifier = nn.Linear(d_model, num_classes) |
| self.emotion_head = nn.Sequential( |
| nn.Linear(d_model, d_model // 2), |
| nn.ReLU(), |
| nn.Dropout(dropout), |
| nn.Linear(d_model // 2, num_emotions) |
| ) |
|
|
| def forward(self, x, saliency_mask=None): |
| B, T, K, _ = x.shape |
| p_feat_end = self.n_pose_feat |
| lh_end = p_feat_end + self.n_hand |
| rh_end = lh_end + self.n_hand |
|
|
| pose_x = self.proj_pose_x (x[:, :, :p_feat_end, 0]) |
| lhand_x = self.proj_lhand_x(x[:, :, p_feat_end:lh_end, 0]) |
| rhand_x = self.proj_rhand_x(x[:, :, lh_end:rh_end, 0]) |
| face_x = self.proj_face_x (x[:, :, rh_end:, 0]) |
|
|
| x_emb = self.norm_x(self.fusion_x(torch.cat([pose_x, lhand_x, rhand_x, face_x], dim=-1))) |
| x_emb = x_emb + self.pos_encoder[:, :T, :] |
| memory = self.encoder(x_emb) |
|
|
| pose_y = self.proj_pose_y (x[:, :, :p_feat_end, 1]) |
| lhand_y = self.proj_lhand_y(x[:, :, p_feat_end:lh_end, 1]) |
| rhand_y = self.proj_rhand_y(x[:, :, lh_end:rh_end, 1]) |
| face_y = self.proj_face_y (x[:, :, rh_end:, 1]) |
|
|
| y_emb = self.norm_y(self.fusion_y(torch.cat([pose_y, lhand_y, rhand_y, face_y], dim=-1))) |
| y_emb = y_emb + self.pos_encoder[:, :T, :] |
|
|
| tgt_mask = torch.triu(torch.ones(T, T, device=x.device), diagonal=1).bool() |
| output = self.decoder(y_emb, memory, tgt_mask=tgt_mask) |
| logits = self.classifier(output.mean(dim=1)) |
|
|
| fx_emb = self.face_x_projector(x[:, :, rh_end:, 0]) + self.pos_encoder[:, :T, :] |
| fy_emb = self.face_y_projector(x[:, :, rh_end:, 1]) + self.pos_encoder[:, :T, :] |
| f_memory = self.emotion_encoder(fx_emb) |
| f_output = self.emotion_decoder(fy_emb, f_memory, tgt_mask=tgt_mask) |
| emotion_logits = self.emotion_head(f_output.mean(dim=1)) |
|
|
| saliency_logits = None |
| if saliency_mask is not None and self.training: |
| mask = saliency_mask.unsqueeze(-1) |
| foreground_feat = (output * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-6) |
| saliency_logits = self.classifier(foreground_feat) |
|
|
| return logits, emotion_logits, saliency_logits |
