| 25 | |
| 26 | |
| 27 | class Model(nn.Module): |
| 28 | def __init__(self, cfg): |
| 29 | super().__init__() |
| 30 | self.backbone = build_backbone(cfg) |
| 31 | self.sem_seg_head = SEM_SEG_HEADS_REGISTRY.get(cfg.MODEL.SEM_SEG_HEAD.NAME)(cfg, self.backbone.output_shape()) |
| 32 | self.text_backbone = BACKBONE_REGISTRY.get("CLIPLangEncoder")(cfg, None) |
| 33 | |
| 34 | self.proj_text_256 = torch.nn.Linear(768, 256) |
| 35 | |
| 36 | def forward(self, x, texts): |
| 37 | backboneOutput = self.backbone(x) |
| 38 | txt_embeds = self.text_backbone(texts) |
| 39 | txt_embeds = self.proj_text_256(txt_embeds) |
| 40 | |
| 41 | outputs = self.sem_seg_head(backboneOutput, txt_feats=txt_embeds) |
| 42 | |
| 43 | mask_cls_results = outputs["pred_logits"] |
| 44 | mask_pred_results = outputs["pred_masks"] |
| 45 | mask_box_results = outputs["pred_boxes"] |
| 46 | # upsample masks |
| 47 | mask_pred_results = F.interpolate( |
| 48 | mask_pred_results, |
| 49 | size=(x.shape[-2], x.shape[-1]), |
| 50 | mode="bilinear", |
| 51 | align_corners=False, |
| 52 | ) |
| 53 | |
| 54 | # bs 1 |
| 55 | for mask_cls_result, mask_pred_result, mask_box_result in zip(mask_cls_results, mask_pred_results, mask_box_results): |
| 56 | boxes = box_ops.box_cxcywh_to_xyxy(mask_box_result) |
| 57 | img_w = x.shape[-1] |
| 58 | img_h = x.shape[-2] |
| 59 | scale_fct = torch.tensor([img_w, img_h, img_w, img_h]) |
| 60 | scale_fct = scale_fct.to(mask_box_result) |
| 61 | mask_box_result = boxes * scale_fct |
| 62 | |
| 63 | scores = mask_cls_result.sigmoid() |
| 64 | num_classes = scores.shape[1] |
| 65 | labels = torch.arange(num_classes).unsqueeze(0).repeat(900, 1).flatten(0, 1).to(scores.device) |
| 66 | scores_per_image, topk_indices = scores.flatten(0, 1).topk(100, sorted=False) |
| 67 | |
| 68 | labels_per_image = labels[topk_indices] |
| 69 | topk_indices = topk_indices // num_classes |
| 70 | mask_pred = mask_pred_result[topk_indices] |
| 71 | mask_box_result = mask_box_result[topk_indices] |
| 72 | |
| 73 | pred_masks = (mask_pred > 0).float() |
| 74 | pred_boxes = mask_box_result |
| 75 | mask_scores_per_image = (mask_pred.sigmoid().flatten(1) * pred_masks.flatten(1)).sum(1) / (pred_masks.flatten(1).sum(1) + 1e-6) |
| 76 | scores = scores_per_image * mask_scores_per_image |
| 77 | pred_classes = labels_per_image |
| 78 | |
| 79 | return scores, pred_boxes, pred_masks, pred_classes |
| 80 | |
| 81 | |
| 82 | class Inference: |