| 7 | |
| 8 | |
| 9 | class bBbox(torch.nn.Module): |
| 10 | def __init__(self,fea_dim): |
| 11 | super(bBbox, self).__init__() |
| 12 | self.img_dim = 4096 |
| 13 | self.attention1 = Attention(dim=128,heads=4) |
| 14 | self.attention2 = Attention(dim=128,heads=4) |
| 15 | |
| 16 | self.linear_img = nn.Sequential(torch.nn.Linear(self.img_dim, fea_dim),torch.nn.ReLU()) |
| 17 | |
| 18 | self.classifier = nn.Linear(fea_dim,2) |
| 19 | |
| 20 | def forward(self, **kwargs): |
| 21 | frames=kwargs['bbox_vgg'] |
| 22 | fea_img = self.linear_img(frames) |
| 23 | fea_img = torch.reshape(fea_img, (-1, 45, 128)) |
| 24 | fea_img = self.attention1(fea_img) |
| 25 | fea_img = torch.mean(fea_img, -2) |
| 26 | fea_img = torch.reshape(fea_img, (-1, 83, 128)) |
| 27 | fea_img = self.attention2(fea_img) |
| 28 | fea_img = torch.mean(fea_img, -2) |
| 29 | output = self.classifier(fea_img) |
| 30 | return output, fea_img |
| 31 | |
| 32 | class bC3D(torch.nn.Module): |
| 33 | def __init__(self,fea_dim): |