Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn.functional as F | |
| from mmcv.runner import BaseModule | |
| from .models import build_model | |
| from .models.util.misc import NestedTensor, inverse_sigmoid | |
| class HDetrWrapper(BaseModule): | |
| def __init__(self, | |
| args=None, | |
| init_cfg=None): | |
| super(HDetrWrapper, self).__init__(init_cfg) | |
| model, box_postprocessor = build_model(args) | |
| self.model = model | |
| self.box_postprocessor = box_postprocessor | |
| self.model.num_queries = self.model.num_queries_one2one | |
| self.model.transformer.two_stage_num_proposals = self.model.num_queries | |
| self.cls_index = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, | |
| 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, | |
| 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, | |
| 82, 84, 85, 86, 87, 88, 89, 90] | |
| def forward(self, | |
| img, | |
| img_metas): | |
| """Forward function for training mode. | |
| Args: | |
| img (Tensor): of shape (N, C, H, W) encoding input images. | |
| Typically these should be mean centered and std scaled. | |
| img_metas (list[dict]): Meta information of each image, e.g., | |
| image size, scaling factor, etc. | |
| """ | |
| input_img_h, input_img_w = img_metas[0]["batch_input_shape"] | |
| batch_size = img.size(0) | |
| img_masks = img.new_ones((batch_size, input_img_h, input_img_w), | |
| dtype=torch.bool) | |
| for img_id in range(batch_size): | |
| img_h, img_w, _ = img_metas[img_id]["img_shape"] | |
| img_masks[img_id, :img_h, :img_w] = False | |
| samples = NestedTensor(tensors=img, mask=img_masks) | |
| features, pos = self.model.backbone(samples) | |
| srcs = [] | |
| masks = [] | |
| for l, feat in enumerate(features): | |
| src, mask = feat.decompose() | |
| srcs.append(self.model.input_proj[l](src)) | |
| masks.append(mask) | |
| assert mask is not None | |
| if self.model.num_feature_levels > len(srcs): | |
| _len_srcs = len(srcs) | |
| for l in range(_len_srcs, self.model.num_feature_levels): | |
| if l == _len_srcs: | |
| src = self.model.input_proj[l](features[-1].tensors) | |
| else: | |
| src = self.model.input_proj[l](srcs[-1]) | |
| m = samples.mask | |
| mask = F.interpolate(m[None].float(), size=src.shape[-2:]).to( | |
| torch.bool | |
| )[0] | |
| pos_l = self.model.backbone[1](NestedTensor(src, mask)).to(src.dtype) | |
| srcs.append(src) | |
| masks.append(mask) | |
| pos.append(pos_l) | |
| query_embeds = None | |
| if not self.model.two_stage or self.model.mixed_selection: | |
| query_embeds = self.model.query_embed.weight[0: self.model.num_queries, :] | |
| # make attn mask | |
| """ attention mask to prevent information leakage | |
| """ | |
| self_attn_mask = ( | |
| torch.zeros([self.model.num_queries, self.model.num_queries, ]).bool().to(src.device) | |
| ) | |
| self_attn_mask[self.model.num_queries_one2one:, 0: self.model.num_queries_one2one, ] = True | |
| self_attn_mask[0: self.model.num_queries_one2one, self.model.num_queries_one2one:, ] = True | |
| ( | |
| hs, | |
| init_reference, | |
| inter_references, | |
| enc_outputs_class, | |
| enc_outputs_coord_unact, | |
| ) = self.model.transformer(srcs, masks, pos, query_embeds, self_attn_mask) | |
| outputs_classes_one2one = [] | |
| outputs_coords_one2one = [] | |
| outputs_classes_one2many = [] | |
| outputs_coords_one2many = [] | |
| for lvl in range(hs.shape[0]): | |
| if lvl == 0: | |
| reference = init_reference | |
| else: | |
| reference = inter_references[lvl - 1] | |
| reference = inverse_sigmoid(reference) | |
| outputs_class = self.model.class_embed[lvl](hs[lvl]) | |
| tmp = self.model.bbox_embed[lvl](hs[lvl]) | |
| if reference.shape[-1] == 4: | |
| tmp += reference | |
| else: | |
| assert reference.shape[-1] == 2 | |
| tmp[..., :2] += reference | |
| outputs_coord = tmp.sigmoid() | |
| outputs_classes_one2one.append( | |
| outputs_class[:, 0: self.model.num_queries_one2one] | |
| ) | |
| outputs_classes_one2many.append( | |
| outputs_class[:, self.model.num_queries_one2one:] | |
| ) | |
| outputs_coords_one2one.append( | |
| outputs_coord[:, 0: self.model.num_queries_one2one] | |
| ) | |
| outputs_coords_one2many.append(outputs_coord[:, self.model.num_queries_one2one:]) | |
| outputs_classes_one2one = torch.stack(outputs_classes_one2one) | |
| outputs_coords_one2one = torch.stack(outputs_coords_one2one) | |
| sampled_logits = outputs_classes_one2one[-1][:, :, self.cls_index] | |
| out = { | |
| "pred_logits": sampled_logits, | |
| "pred_boxes": outputs_coords_one2one[-1], | |
| } | |
| return out | |
| def simple_test(self, img, img_metas, rescale=False): | |
| # out: dict | |
| out = self(img, img_metas) | |
| if rescale: | |
| ori_target_sizes = [meta_info['ori_shape'][:2] for meta_info in img_metas] | |
| else: | |
| ori_target_sizes = [meta_info['img_shape'][:2] for meta_info in img_metas] | |
| ori_target_sizes = out['pred_logits'].new_tensor(ori_target_sizes, dtype=torch.int64) | |
| # results: List[dict(scores, labels, boxes)] | |
| results = self.box_postprocessor(out, ori_target_sizes) | |
| return results | |