|
|
import torch
|
|
|
import torch.nn as nn
|
|
|
import torch.nn.functional as F
|
|
|
from mmengine.model import BaseModule, Sequential
|
|
|
|
|
|
from mmaction.models.utils import Graph
|
|
|
from mmaction.registry import MODELS
|
|
|
from .msg3d_utils import MSGCN, MSTCN, MW_MSG3DBlock
|
|
|
|
|
|
|
|
|
@MODELS.register_module()
|
|
|
class MSG3D(BaseModule):
|
|
|
|
|
|
def __init__(self,
|
|
|
graph_cfg,
|
|
|
in_channels=3,
|
|
|
base_channels=96,
|
|
|
num_gcn_scales=13,
|
|
|
num_g3d_scales=6,
|
|
|
num_person=2,
|
|
|
tcn_dropout=0):
|
|
|
super().__init__()
|
|
|
|
|
|
self.graph = Graph(**graph_cfg)
|
|
|
|
|
|
A = torch.tensor(
|
|
|
self.graph.A[0], dtype=torch.float32, requires_grad=False)
|
|
|
self.register_buffer('A', A)
|
|
|
self.num_point = A.shape[-1]
|
|
|
self.in_channels = in_channels
|
|
|
self.base_channels = base_channels
|
|
|
|
|
|
self.data_bn = nn.BatchNorm1d(self.num_point * in_channels *
|
|
|
num_person)
|
|
|
c1, c2, c3 = base_channels, base_channels * 2, base_channels * 4
|
|
|
|
|
|
|
|
|
self.gcn3d1 = MW_MSG3DBlock(3, c1, A, num_g3d_scales, window_stride=1)
|
|
|
self.sgcn1 = Sequential(
|
|
|
MSGCN(num_gcn_scales, 3, c1, A), MSTCN(c1, c1), MSTCN(c1, c1))
|
|
|
self.sgcn1[-1].act = nn.Identity()
|
|
|
self.tcn1 = MSTCN(c1, c1, tcn_dropout=tcn_dropout)
|
|
|
|
|
|
self.gcn3d2 = MW_MSG3DBlock(c1, c2, A, num_g3d_scales, window_stride=2)
|
|
|
self.sgcn2 = Sequential(
|
|
|
MSGCN(num_gcn_scales, c1, c1, A), MSTCN(c1, c2, stride=2),
|
|
|
MSTCN(c2, c2))
|
|
|
self.sgcn2[-1].act = nn.Identity()
|
|
|
self.tcn2 = MSTCN(c2, c2, tcn_dropout=tcn_dropout)
|
|
|
|
|
|
self.gcn3d3 = MW_MSG3DBlock(c2, c3, A, num_g3d_scales, window_stride=2)
|
|
|
self.sgcn3 = Sequential(
|
|
|
MSGCN(num_gcn_scales, c2, c2, A), MSTCN(c2, c3, stride=2),
|
|
|
MSTCN(c3, c3))
|
|
|
self.sgcn3[-1].act = nn.Identity()
|
|
|
self.tcn3 = MSTCN(c3, c3, tcn_dropout=tcn_dropout)
|
|
|
|
|
|
def forward(self, x):
|
|
|
N, M, T, V, C = x.size()
|
|
|
x = x.permute(0, 1, 3, 4, 2).contiguous().reshape(N, M * V * C, T)
|
|
|
x = self.data_bn(x)
|
|
|
x = x.reshape(N * M, V, C, T).permute(0, 2, 3, 1).contiguous()
|
|
|
|
|
|
|
|
|
x = F.relu(self.sgcn1(x) + self.gcn3d1(x), inplace=True)
|
|
|
x = self.tcn1(x)
|
|
|
|
|
|
x = F.relu(self.sgcn2(x) + self.gcn3d2(x), inplace=True)
|
|
|
x = self.tcn2(x)
|
|
|
|
|
|
x = F.relu(self.sgcn3(x) + self.gcn3d3(x), inplace=True)
|
|
|
x = self.tcn3(x)
|
|
|
|
|
|
|
|
|
return x.reshape((N, M) + x.shape[1:])
|
|
|
|
