add Gait3D support

2022-12-05 21:47:24 +08:00
parent f79997bf8d
commit 77abcc2a10
11 changed files with 5525 additions and 13 deletions
@@ -0,0 +1,101 @@
+'''
+Modifed from https://github.com/Gait3D/Gait3D-Benchmark/blob/72beab994c137b902d826f4b9f9e95b107bebd78/lib/modeling/models/smplgait.py
+'''
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+
+from ..base_model import BaseModel
+from ..modules import SetBlockWrapper, HorizontalPoolingPyramid, PackSequenceWrapper, SeparateFCs, SeparateBNNecks
+
+
+class SMPLGait(BaseModel):
+    def __init__(self, cfgs, is_training):
+        super().__init__(cfgs, is_training)
+
+    def build_network(self, model_cfg):
+        # Baseline
+        self.Backbone = self.get_backbone(model_cfg['backbone_cfg'])
+        self.Backbone = SetBlockWrapper(self.Backbone)
+        self.FCs = SeparateFCs(**model_cfg['SeparateFCs'])
+        self.BNNecks = SeparateBNNecks(**model_cfg['SeparateBNNecks'])
+        self.TP = PackSequenceWrapper(torch.max)
+        self.HPP = HorizontalPoolingPyramid(bin_num=model_cfg['bin_num'])
+
+        # for SMPL
+        self.fc1 = nn.Linear(85, 128)
+        self.fc2 = nn.Linear(128, 256)
+        self.fc3 = nn.Linear(256, 256)
+        self.bn1 = nn.BatchNorm1d(128)
+        self.bn2 = nn.BatchNorm1d(256)
+        self.bn3 = nn.BatchNorm1d(256)
+        self.dropout2 = nn.Dropout(p=0.2)
+        self.dropout3 = nn.Dropout(p=0.2)
+
+    def forward(self, inputs):
+        ipts, labs, _, _, seqL = inputs
+
+        sils = ipts[0]    # [n, s, h, w]
+        smpls = ipts[1]   # [n, s, d]
+
+        # extract SMPL features
+        n, s, d = smpls.size()
+        sps = smpls.view(-1, d)
+        del smpls
+
+        sps = F.relu(self.bn1(self.fc1(sps)))
+        sps = F.relu(self.bn2(self.dropout2(self.fc2(sps))))  # (B, 256)
+        sps = F.relu(self.bn3(self.dropout3(self.fc3(sps))))  # (B, 256)
+        sps = sps.reshape(n, 1, s, 16, 16)
+        iden = Variable(torch.eye(16)).unsqueeze(
+            0).repeat(n, 1, s, 1, 1)   # [n, 1, s, 16, 16]
+        if sps.is_cuda:
+            iden = iden.cuda()
+        sps_trans = sps + iden   # [n, 1, s, 16, 16]
+
+        if len(sils.size()) == 4:
+            sils = sils.unsqueeze(1)
+
+        del ipts
+        outs = self.Backbone(sils)  # [n, c, s, h, w]
+        outs_n, outs_c, outs_s, outs_h, outs_w = outs.size()
+
+        zero_tensor = Variable(torch.zeros(
+            (outs_n, outs_c, outs_s, outs_h, outs_h-outs_w)))
+        if outs.is_cuda:
+            zero_tensor = zero_tensor.cuda()
+        # [n, s, c, h, h]  [n, s, c, 16, 16]
+        outs = torch.cat([outs, zero_tensor], -1)
+        outs = outs.reshape(outs_n*outs_c*outs_s, outs_h,
+                            outs_h)   # [n*c*s, 16, 16]
+
+        sps = sps_trans.repeat(1, outs_c, 1, 1, 1).reshape(
+            outs_n * outs_c * outs_s, 16, 16)
+
+        outs_trans = torch.bmm(outs, sps)
+        outs_trans = outs_trans.reshape(outs_n, outs_c, outs_s, outs_h, outs_h)
+
+        # Temporal Pooling, TP
+        outs_trans = self.TP(outs_trans, seqL, options={"dim": 2})[
+            0]  # [n, c, h, w]
+        # Horizontal Pooling Matching, HPM
+        feat = self.HPP(outs_trans)  # [n, c, p]
+        embed_1 = self.FCs(feat)  # [n, c, p]
+
+        embed_2, logits = self.BNNecks(embed_1)  # [n, c, p]
+
+        n, _, s, h, w = sils.size()
+        retval = {
+            'training_feat': {
+                'triplet': {'embeddings': embed_1, 'labels': labs},
+                'softmax': {'logits': logits, 'labels': labs}
+            },
+            'visual_summary': {
+                'image/sils': sils.view(n*s, 1, h, w)
+            },
+            'inference_feat': {
+                'embeddings': embed_1
+            }
+        }
+        return retval