def forward(self, feat, proposals, img_shape): batch_size = proposals.shape[0] num_proposals = proposals.shape[1] rois_pad = mm2trt_util.arange_by_input(proposals, 0).unsqueeze(1) rois_pad = rois_pad.repeat(1, num_proposals).view(-1, 1) proposals = proposals.view(-1, 4) rois = torch.cat([rois_pad, proposals], dim=1) # rcnn bbox_results = self._bbox_forward(feat, rois) cls_score = bbox_results['cls_score'] bbox_pred = bbox_results['bbox_pred'] bbox_head_outputs = self.bbox_head.get_bboxes(rois, cls_score, bbox_pred, img_shape, batch_size, num_proposals, self.test_cfg) num_detections, det_boxes, det_scores, det_classes = bbox_head_outputs result = [num_detections, det_boxes, det_scores, det_classes] if self.enable_mask: # mask roi input num_mask_proposals = det_boxes.size(1) rois_pad = mm2trt_util.arange_by_input(det_boxes, 0).unsqueeze(1) rois_pad = rois_pad.repeat(1, num_mask_proposals).view(-1, 1) mask_proposals = det_boxes.view(-1, 4) mask_rois = torch.cat([rois_pad, mask_proposals], dim=1) mask_results = self._mask_forward(feat, mask_rois) mask_pred = mask_results['mask_pred'] mc, mh, mw = mask_pred.shape[1:] mask_pred = mask_pred.reshape(batch_size, -1, mc, mh, mw).sigmoid() if not self.module.mask_head.class_agnostic: det_index = det_classes.unsqueeze(-1).long() det_index = det_index + 1 mask_pad = mask_pred[:, :, 0:1, ...] * 0 mask_pred = torch.cat([mask_pad, mask_pred], dim=2) mask_pred = mm2trt_util.gather_topk(mask_pred, dim=2, index=det_index) mask_pred = mask_pred.squeeze(2) result += [mask_pred] return result
def forward(self, feat, proposals): batch_size = proposals.shape[0] num_proposals = proposals.shape[1] rois_pad = mm2trt_util.arange_by_input(proposals, 0).unsqueeze(1) rois_pad = rois_pad.repeat(1, num_proposals).view(-1, 1) proposals = proposals.view(-1, 4) rois = torch.cat([rois_pad, proposals], dim=1) roi_feats = self.bbox_roi_extractor( feat[:len(self.bbox_roi_extractor.featmap_strides)], rois) if self.shared_head is not None: roi_feats = self.shared_head(roi_feats) # rcnn cls_score, bbox_pred = self.bbox_head(roi_feats) if isinstance(cls_score, list): cls_score = sum(cls_score) / float(len(cls_score)) scores = F.softmax(cls_score, dim=1) bboxes = delta2bbox(proposals, bbox_pred, self.bbox_head.bbox_coder.means, self.bbox_head.bbox_coder.stds) scores = scores.view(batch_size, num_proposals, -1) bboxes = bboxes.view(batch_size, num_proposals, -1, 4) num_bboxes = bboxes.shape[1] bboxes_ext = bboxes[:, :, 0:1, :] * 0 bboxes = torch.cat([bboxes, bboxes_ext], 2) num_detections, det_boxes, det_scores, det_classes = self.rcnn_nms( scores, bboxes, num_bboxes, self.test_cfg.max_per_img) return num_detections, det_boxes, det_scores, det_classes
def bbox2roi(proposals): num_proposals = proposals.shape[1] rois_pad = mm2trt_util.arange_by_input(proposals, 0).unsqueeze(1) rois_pad = rois_pad.repeat(1, num_proposals).view(-1, 1) proposals = proposals.view(-1, 4) rois = torch.cat([rois_pad, proposals], dim=1) return rois
def forward(self, feat, proposals, img_shape): batch_size = proposals.shape[0] num_proposals = proposals.shape[1] rois_pad = mm2trt_util.arange_by_input(proposals, 0).unsqueeze(1) rois_pad = rois_pad.repeat(1, num_proposals).view(-1, 1) proposals = proposals.view(-1, 4) rois = torch.cat([rois_pad, proposals], dim=1) # rcnn bbox_results = self._bbox_forward(feat, rois) cls_score = bbox_results['cls_score'] bbox_pred = bbox_results['bbox_pred'] bbox_head_outputs = self.bbox_head.get_bboxes(rois, cls_score, bbox_pred, img_shape, batch_size, num_proposals, self.test_cfg) num_detections, det_boxes, det_scores, det_classes = bbox_head_outputs grid_rois = bbox2roi(det_boxes) grid_feats = self.grid_roi_extractor( feat[:len(self.grid_roi_extractor.featmap_strides)], grid_rois) grid_pred = self.grid_head(grid_feats) det_scores, det_boxes = self.grid_head.get_bboxes( det_scores.view(-1), det_boxes.view(-1, 4), grid_pred['fused'], img_shape) det_scores = det_scores.view(batch_size, -1) det_boxes = det_boxes.view(batch_size, -1, 4) return num_detections, det_boxes, det_scores, det_classes
def forward(self, feat ,proposals, img_shape): ms_scores = [] batch_size = proposals.shape[0] num_proposals = proposals.shape[1] rois_pad = mm2trt_util.arange_by_input(proposals, 0).unsqueeze(1) rois_pad = rois_pad.repeat(1, num_proposals).view(-1, 1) proposals = proposals.view(-1, 4) rois = proposals for i in range(self.num_stages): bbox_results = self._bbox_forward(i, feat, torch.cat([rois_pad, rois], dim=1)) ms_scores.append(bbox_results['cls_score']) bbox_pred = bbox_results['bbox_pred'] if i < self.num_stages - 1: bbox_label = bbox_results['cls_score'].argmax(dim=1) rois = self.bbox_head[i].regress_by_class(rois, bbox_label, bbox_pred, img_shape) rois = torch.cat([rois_pad, rois], dim=1) ### bbox_head.get_boxes cls_score = bbox_results['cls_score'] bbox_pred = bbox_results['bbox_pred'] num_detections, det_boxes, det_scores, det_classes = self.bbox_head[-1].get_bboxes( rois, cls_score, bbox_pred, img_shape, batch_size, num_proposals, self.test_cfg ) return num_detections, det_boxes, det_scores, det_classes
def forward(self, feat, proposals, img_shape): batch_size = proposals.shape[0] num_proposals = proposals.shape[1] rois_pad = mm2trt_util.arange_by_input(proposals, 0).unsqueeze(1) rois_pad = rois_pad.repeat(1, num_proposals).view(-1, 1) proposals = proposals.view(-1, 4) rois = torch.cat([rois_pad, proposals], dim=1) bbox_results = self._bbox_forward(feat, rois) cls_score = bbox_results['cls_score'] bbox_pred = bbox_results['bbox_pred'] num_detections, det_boxes, det_scores, det_classes = self.bbox_head.get_bboxes( rois, cls_score, bbox_pred, img_shape, batch_size, num_proposals, self.test_cfg) return num_detections, det_boxes, det_scores, det_classes
def forward(self, feat, proposals): ms_scores = [] batch_size = proposals.shape[0] num_proposals = proposals.shape[1] rois_pad = mm2trt_util.arange_by_input(proposals, 0).unsqueeze(1) rois_pad = rois_pad.repeat(1, num_proposals).view(-1, 1) proposals = proposals.view(-1, 4) rois = proposals for i in range(self.num_stages): bbox_results = self._bbox_forward( i, feat, torch.cat([rois_pad, rois], dim=1)) ms_scores.append(bbox_results['cls_score']) bbox_pred = bbox_results['bbox_pred'] if i < self.num_stages - 1: bbox_label = bbox_results['cls_score'].argmax(dim=1) rois = self.regress_by_class(i, rois, bbox_label, bbox_pred) ### bbox_head.get_boxes cls_score = sum(ms_scores) / self.num_stages if isinstance(cls_score, list): cls_score = sum(cls_score) / float(len(cls_score)) scores = F.softmax(cls_score, dim=1) bbox_head = self.bbox_head[-1] bboxes = delta2bbox(rois, bbox_pred, bbox_head.bbox_coder.means, bbox_head.bbox_coder.stds) scores = scores.view(batch_size, num_proposals, -1) bboxes = bboxes.view(batch_size, num_proposals, -1, 4) bboxes = bboxes.repeat(1, 1, bbox_head.num_classes, 1) bboxes_ext = bboxes[:, :, 0:1, :] * 0 bboxes = torch.cat([bboxes, bboxes_ext], 2) num_bboxes = bboxes.shape[1] num_detections, det_boxes, det_scores, det_classes = self.rcnn_nms( scores, bboxes, num_bboxes, self.test_cfg.max_per_img) return num_detections, det_boxes, det_scores, det_classes
def forward(self, feat, proposals, img_shape): ms_scores = [] batch_size = proposals.shape[0] num_proposals = proposals.shape[1] rois_pad = mm2trt_util.arange_by_input(proposals, 0).unsqueeze(1) rois_pad = rois_pad.repeat(1, num_proposals).view(-1, 1) proposals = proposals.view(-1, 4) rois = proposals if self.module.with_semantic: _, semantic_feat = self.semantic_head(feat) else: semantic_feat = None for i in range(self.num_stages): bbox_results = self._bbox_forward( i, feat, torch.cat([rois_pad, rois], dim=1), semantic_feat=semantic_feat) ms_scores.append(bbox_results['cls_score']) bbox_pred = bbox_results['bbox_pred'] if i < self.num_stages - 1: bbox_label = bbox_results['cls_score'].argmax(dim=1) rois = self.bbox_head[i].regress_by_class( rois, bbox_label, bbox_pred, img_shape) rois = torch.cat([rois_pad, rois], dim=1) # bbox_head.get_boxes cls_score = bbox_results['cls_score'] bbox_pred = bbox_results['bbox_pred'] num_detections, det_boxes, det_scores, det_classes = self.bbox_head[ -1].get_bboxes(rois, cls_score, bbox_pred, img_shape, batch_size, num_proposals, self.test_cfg) result = [num_detections, det_boxes, det_scores, det_classes] if self.enable_mask: # mask roi input num_mask_proposals = det_boxes.size(1) rois_pad = mm2trt_util.arange_by_input(det_boxes, 0).unsqueeze(1) rois_pad = rois_pad.repeat(1, num_mask_proposals).view(-1, 1) mask_proposals = det_boxes.view(-1, 4) mask_rois = torch.cat([rois_pad, mask_proposals], dim=1) mask_roi_extractor = self.mask_roi_extractor[-1] mask_feats = mask_roi_extractor( feat[:mask_roi_extractor.num_inputs], mask_rois) if self.module.with_semantic and ('mask' in self.module.semantic_fusion): mask_semantic_feat = self.semantic_roi_extractor( [semantic_feat], mask_rois) mask_feats += mask_semantic_feat last_feat = None aug_masks = [] for i in range(self.num_stages): mask_head = self.mask_head[i] if self.module.mask_info_flow: mask_pred, last_feat = mask_head(mask_feats, last_feat) else: mask_pred = mask_head(mask_feats) mask_pred = mask_pred.sigmoid() aug_masks.append(mask_pred) mask_pred = merge_aug_masks(aug_masks, self.test_cfg) mc, mh, mw = mask_pred.shape[1:] mask_pred = mask_pred.reshape(batch_size, -1, mc, mh, mw) if not self.module.mask_head[-1].class_agnostic: det_index = det_classes.unsqueeze(-1).long() det_index = det_index + 1 mask_pad = mask_pred[:, :, 0:1, ...] * 0 mask_pred = torch.cat([mask_pad, mask_pred], dim=2) mask_pred = mm2trt_util.gather_topk( mask_pred, dim=2, index=det_index) mask_pred = mask_pred.squeeze(2) result += [mask_pred] return result