(self, bevs, trans_matrices=None, num_agent_tensor=None, batch_size=None)
| 424 | return fused_images |
| 425 | |
| 426 | def forward(self, bevs, trans_matrices=None, num_agent_tensor=None, batch_size=None): |
| 427 | bevs = bevs.permute(0, 1, 4, 2, 3) # (Batch, seq, z, h, w) |
| 428 | |
| 429 | x_8, x_7, x_6, x_5, x_3, x_2 = self.stpn(bevs) |
| 430 | x = x_8 |
| 431 | |
| 432 | # cls_preds, loc_preds, result = super().get_cls_loc_result(x) |
| 433 | |
| 434 | if self.train_completion: |
| 435 | # fuse reconstructed BEVs |
| 436 | ind_result = torch.argmax(torch.softmax(x, dim=1), dim=1) |
| 437 | result = self.ego_late_fusion(ind_result, bevs.squeeze(1), trans_matrices, num_agent_tensor, batch_size) |
| 438 | else: |
| 439 | # do detection |
| 440 | cls_preds, loc_preds, result = super().get_cls_loc_result(x) |
| 441 | |
| 442 | |
| 443 | if self.kd_flag == 1: |
| 444 | return result, x_8, x_7, x_6, x_5, x_3 |
| 445 | elif self.train_completion: |
| 446 | return result, x_8 |
| 447 | else: |
| 448 | return result |
| 449 | |
| 450 | |
| 451 | class Conv3D(nn.Module): |
nothing calls this directly
no test coverage detected