| 233 | |
| 234 | |
| 235 | class SparseGraphConvolution(nn.Module): |
| 236 | |
| 237 | def __init__(self, in_dims=16, embedding_dims=16, dropout=0): |
| 238 | super(SparseGraphConvolution, self).__init__() |
| 239 | |
| 240 | self.dropout = dropout |
| 241 | |
| 242 | self.spatial_temporal_sparse_gcn = nn.ModuleList() |
| 243 | self.temporal_spatial_sparse_gcn = nn.ModuleList() |
| 244 | |
| 245 | self.spatial_temporal_sparse_gcn.append(GraphConvolution(in_dims, embedding_dims)) |
| 246 | self.spatial_temporal_sparse_gcn.append(GraphConvolution(embedding_dims, embedding_dims)) |
| 247 | |
| 248 | self.temporal_spatial_sparse_gcn.append(GraphConvolution(in_dims, embedding_dims)) |
| 249 | self.temporal_spatial_sparse_gcn.append(GraphConvolution(embedding_dims, embedding_dims)) |
| 250 | |
| 251 | def forward(self, graph, normalized_spatial_adjacency_matrix, normalized_temporal_adjacency_matrix): |
| 252 | |
| 253 | # graph [1 seq_len num_pedestrians 3] |
| 254 | # _matrix [batch num_heads seq_len seq_len] |
| 255 | |
| 256 | graph = graph[:, :, :, 1:] |
| 257 | spa_graph = graph.permute(1, 0, 2, 3) # (seq_len 1 num_p 2) |
| 258 | tem_graph = spa_graph.permute(2, 1, 0, 3) # (num_p 1 seq_len 2) |
| 259 | |
| 260 | gcn_spatial_features = self.spatial_temporal_sparse_gcn[0](spa_graph, normalized_spatial_adjacency_matrix) |
| 261 | gcn_spatial_features = gcn_spatial_features.permute(2, 1, 0, 3) |
| 262 | |
| 263 | # [num_p num_heads seq_len d] |
| 264 | gcn_spatial_temporal_features = self.spatial_temporal_sparse_gcn[1](gcn_spatial_features, normalized_temporal_adjacency_matrix) |
| 265 | |
| 266 | gcn_temporal_features = self.temporal_spatial_sparse_gcn[0](tem_graph, |
| 267 | normalized_temporal_adjacency_matrix) |
| 268 | gcn_temporal_features = gcn_temporal_features.permute(2, 1, 0, 3) |
| 269 | gcn_temporal_spatial_features = self.temporal_spatial_sparse_gcn[1](gcn_temporal_features, |
| 270 | normalized_spatial_adjacency_matrix) |
| 271 | |
| 272 | return gcn_spatial_temporal_features, gcn_temporal_spatial_features.permute(2, 1, 0, 3) |
| 273 | |
| 274 | |
| 275 | class TrajectoryModel(nn.Module): |