Aggregate the context features according to the initial predicted probability distribution. Employ the soft-weighted method to aggregate the context. Output: The correlation of every class map with every feature map shape = [n, num_feats, num_cla
| 15 | |
| 16 | |
| 17 | class SpatialGather_Module(nn.Module): |
| 18 | """ |
| 19 | Aggregate the context features according to the initial |
| 20 | predicted probability distribution. |
| 21 | Employ the soft-weighted method to aggregate the context. |
| 22 | |
| 23 | Output: |
| 24 | The correlation of every class map with every feature map |
| 25 | shape = [n, num_feats, num_classes, 1] |
| 26 | |
| 27 | |
| 28 | """ |
| 29 | def __init__(self, cls_num=0, scale=1): |
| 30 | super(SpatialGather_Module, self).__init__() |
| 31 | self.cls_num = cls_num |
| 32 | self.scale = scale |
| 33 | |
| 34 | def forward(self, feats, probs): |
| 35 | batch_size, c, _, _ = probs.size(0), probs.size(1), probs.size(2), \ |
| 36 | probs.size(3) |
| 37 | |
| 38 | # each class image now a vector |
| 39 | probs = probs.view(batch_size, c, -1) |
| 40 | feats = feats.view(batch_size, feats.size(1), -1) |
| 41 | |
| 42 | feats = feats.permute(0, 2, 1) # batch x hw x c |
| 43 | probs = F.softmax(self.scale * probs, dim=2) # batch x k x hw |
| 44 | ocr_context = torch.matmul(probs, feats) |
| 45 | ocr_context = ocr_context.permute(0, 2, 1).unsqueeze(3) |
| 46 | return ocr_context |
| 47 | |
| 48 | |
| 49 | class ObjectAttentionBlock(nn.Module): |