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
| import torch import torch.nn as nn
class STGCN_Plus(nn.Module): """ 时空图卷积网络增强版 用于驾驶员姿态行为识别 """ def __init__(self, num_joints=16, num_classes=10): super().__init__() self.st_gcn = nn.ModuleList([ STGCNBlock(in_channels=3, out_channels=64, A=None), STGCNBlock(in_channels=64, out_channels=128, A=None), STGCNBlock(in_channels=128, out_channels=256, A=None), ]) self.classifier = nn.Sequential( nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(256, num_classes) ) def forward(self, x): """ Args: x: (B, C, T, V, M) B: batch size C: 3 (x, y, z) T: temporal frames V: 16 joints M: 1 (single person) Returns: logits: (B, num_classes) """ for stgcn in self.st_gcn: x = stgcn(x) return self.classifier(x)
class STGCNBlock(nn.Module): """时空图卷积块""" def __init__(self, in_channels, out_channels, A): super().__init__() self.gcn = GraphConvolution(in_channels, out_channels, A) self.tcn = nn.Sequential( nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True), nn.Conv2d(out_channels, out_channels, (9, 1), padding=(4, 0)), nn.BatchNorm2d(out_channels), ) self.residual = nn.Conv2d(in_channels, out_channels, 1) def forward(self, x): return self.tcn(self.gcn(x)) + self.residual(x)
ADJACENCY_MATRIX = [ [0, 1], [1, 2], [2, 3], [3, 4], [1, 5], [5, 6], [6, 7], [1, 8], [8, 9], [9, 10], [4, 11], [11, 12], [12, 13], [4, 14], [14, 15], ]
|