(self, img_s_list, mask_s_list, img_q)
| 11 | self.backbone = resnet.__dict__[backbone](pretrained=True) |
| 12 | |
| 13 | def forward(self, img_s_list, mask_s_list, img_q): |
| 14 | h, w = img_q.shape[-2:] |
| 15 | |
| 16 | # feature maps of support images |
| 17 | feature_s_list = [] |
| 18 | for k in range(len(img_s_list)): |
| 19 | feature_s_list.append(self.backbone.base_forward(img_s_list[k])) |
| 20 | # feature map of query image |
| 21 | feature_q = self.backbone.base_forward(img_q) |
| 22 | |
| 23 | # foreground(target class) and background prototypes pooled from K support features |
| 24 | feature_fg_list = [] |
| 25 | feature_bg_list = [] |
| 26 | for k in range(len(img_s_list)): |
| 27 | feature_fg_list.append(self.masked_average_pooling(feature_s_list[k], |
| 28 | (mask_s_list[k] == 1).float())[None, :]) |
| 29 | feature_bg_list.append(self.masked_average_pooling(feature_s_list[k], |
| 30 | (mask_s_list[k] == 0).float())[None, :]) |
| 31 | # average K foreground prototypes and K background prototypes |
| 32 | feature_fg = torch.mean(torch.cat(feature_fg_list, dim=0), dim=0) |
| 33 | feature_bg = torch.mean(torch.cat(feature_bg_list, dim=0), dim=0) |
| 34 | |
| 35 | # measure the similarity of query features to fg/bg prototypes |
| 36 | similarity_fg = F.cosine_similarity(feature_q, feature_fg[..., None, None], dim=1) |
| 37 | similarity_bg = F.cosine_similarity(feature_q, feature_bg[..., None, None], dim=1) |
| 38 | |
| 39 | out = torch.cat((similarity_bg[:, None, ...], similarity_fg[:, None, ...]), dim=1) * 10.0 |
| 40 | out = F.interpolate(out, size=(h, w), mode="bilinear", align_corners=True) |
| 41 | |
| 42 | return out |
| 43 | |
| 44 | def masked_average_pooling(self, feature, mask): |
| 45 | feature = F.interpolate(feature, size=mask.shape[-2:], mode="bilinear", align_corners=True) |
nothing calls this directly
no test coverage detected