def _forward_mask(self, features, instances): """ Forward logic of the mask prediction branch. Args: features (list[Tensor]): #level input features for mask prediction instances (list[Instances]): the per-image instances to train/predict masks. In training, they can be the proposals. In inference, they can be the predicted boxes. Returns: In training, a dict of losses. In inference, update `instances` with new fields "pred_masks" and return it. """ if not self.mask_on: return {} if self.training else instances if self.training: # The loss is only defined on positive proposals. proposals, _ = select_foreground_proposals(instances, self.num_classes) proposal_boxes = [x.proposal_boxes for x in proposals] mask_features = self.mask_pooler(features, proposal_boxes) mask_logits = self.mask_head.layers(mask_features) loss_mask, target_masks = mask_rcnn_loss(mask_logits, proposals) if self._vis: self._misc["target_masks"] = target_masks self._misc["fg_proposals"] = proposals return {"loss_mask": loss_mask} else: pred_boxes = [x.pred_boxes for x in instances] mask_features = self.mask_pooler(features, pred_boxes) return self.mask_head(mask_features, instances)
def _forward_z(self, features, instances): """ Forward logic of the z prediction branch. """ if not self.zpred_on: return {} if self.training else instances if self.training: # The loss is only defined on positive proposals. proposals, _ = select_foreground_proposals(instances, self.num_classes) proposal_boxes = [x.proposal_boxes for x in proposals] z_features = self.z_pooler(features, proposal_boxes) z_pred = self.z_head(z_features) src_boxes = cat([p.tensor for p in proposal_boxes]) loss_z_reg = z_rcnn_loss( z_pred, proposals, src_boxes, loss_weight=self.z_loss_weight, smooth_l1_beta=self.z_smooth_l1_beta, ) return {"loss_z_reg": loss_z_reg} else: pred_boxes = [x.pred_boxes for x in instances] z_features = self.z_pooler(features, pred_boxes) z_pred = self.z_head(z_features) z_rcnn_inference(z_pred, instances) return instances
def _forward_mask(self, features, instances): """ Forward logic of the mask prediction branch. Args: features (dict[str, Tensor]): #level input features for mask prediction instances (list[Instances]): the per-image instances to train/predict masks. In training, they can be the proposals. In inference, they can be the predicted boxes. Returns: In training, a dict of losses. In inference, update `instances` with new fields "pred_masks" and return it. """ if not self.mask_on: return {} if self.training else instances if self.training: proposals, _ = select_foreground_proposals(instances, self.num_classes) proposal_boxes = [x.proposal_boxes for x in proposals] mask_coarse_logits = self._forward_mask_coarse(features, proposal_boxes) losses = {"loss_mask": mask_rcnn_loss(mask_coarse_logits, proposals)} losses.update(self._forward_mask_point(features, mask_coarse_logits, proposals)) return losses else: pred_boxes = [x.pred_boxes for x in instances] mask_coarse_logits = self._forward_mask_coarse(features, pred_boxes) mask_logits = self._forward_mask_point(features, mask_coarse_logits, instances) mask_rcnn_inference(mask_logits, instances) return instances
def _forward_mask( self, features: Dict[str, torch.Tensor], instances: List[Instances] ) -> Union[Dict[str, torch.Tensor], List[Instances]]: """ Forward logic of the mask prediction branch. Args: features (dict[str, Tensor]): mapping from feature map names to tensor. Same as in :meth:`ROIHeads.forward`. instances (list[Instances]): the per-image instances to train/predict masks. In training, they can be the proposals. In inference, they can be the predicted boxes. Returns: In training, a dict of losses. In inference, update `instances` with new fields "pred_masks" and return it. """ if not self.mask_on: return {} if self.training else instances mask_features = [features[f] for f in self.mask_in_features] boundary_features = [features[f] for f in self.boundary_in_features] if self.training: # The loss is only defined on positive proposals. proposals, _ = select_foreground_proposals(instances, self.num_classes) proposal_boxes = [x.proposal_boxes for x in proposals] mask_features = self.mask_pooler(mask_features, proposal_boxes) boundary_features = self.boundary_pooler(boundary_features, proposal_boxes) return self.mask_head(mask_features, boundary_features, proposals) else: pred_boxes = [x.pred_boxes for x in instances] mask_features = self.mask_pooler(mask_features, pred_boxes) boundary_features = self.boundary_pooler(boundary_features, pred_boxes) return self.mask_head(mask_features, boundary_features, instances)
def _forward_keypoint(self, features, instances): """ Forward logic of the keypoint prediction branch. Args: features (list[Tensor]): #level input features for keypoint prediction instances (list[Instances]): the per-image instances to train/predict keypoints. In training, they can be the proposals. In inference, they can be the predicted boxes. Returns: In training, a dict of losses. In inference, update `instances` with new fields "pred_keypoints" and return it. """ if not self.keypoint_on: return {} if self.training else instances num_images = len(instances) if self.training: # The loss is defined on positive proposals with at >=1 visible keypoints. proposals, _ = select_foreground_proposals(instances, self.num_classes) proposals = select_proposals_with_visible_keypoints(proposals) proposal_boxes = [x.proposal_boxes for x in proposals] keypoint_features = self.keypoint_pooler(features, proposal_boxes) keypoint_logits = self.keypoint_head(keypoint_features) normalizer = (num_images * self.batch_size_per_image * self.positive_sample_fraction * keypoint_logits.shape[1]) if self.keypoint_loss_type == "FiberKeypointLoss": loss = fiber_keypoint_rcnn_loss(keypoint_logits, proposals) else: loss = keypoint_rcnn_loss( keypoint_logits, proposals, normalizer=None if self.normalize_loss_by_visible_keypoints else normalizer, ) return {"loss_keypoint": loss * self.keypoint_loss_weight} else: pred_boxes = [x.pred_boxes for x in instances] keypoint_features = self.keypoint_pooler(features, pred_boxes) keypoint_logits = self.keypoint_head(keypoint_features) keypoint_rcnn_inference(keypoint_logits, instances) return instances
def _forward_mask(self, features, instances): """ Forward logic of the mask prediction branch. Args: features (dict[str, Tensor]): #level input features for mask prediction instances (list[Instances]): the per-image instances to train/predict masks. In training, they can be the proposals. In inference, they can be the predicted boxes. Returns: In training, a dict of losses. In inference, update `instances` with new fields "pred_masks" and return it. """ if not self.mask_on: return {} if self.training else instances if self.training: instances, _ = select_foreground_proposals(instances, self.num_classes) if self.mask_pooler is not None: features = [features[f] for f in self.mask_in_features] boxes = [ x.proposal_boxes if self.training else x.pred_boxes for x in instances ] features = self.mask_pooler(features, boxes) else: features = dict([(f, features[f]) for f in self.mask_in_features]) mask_logits = self._forward_mask_rcnn(features, instances) if self.training: return self._forward_parallel_border(mask_logits, features, self.points_num, instances) else: # add "od_activated_map" attribute for instances, which is used for visualization od_activated_map, border_mask_logits = self._forward_parallel_border( mask_logits, features, self.points_num, instances) instance_num = 0 for instance in instances: instance.chosenP = od_activated_map[instance_num:instance_num + len(instance), :, :] instance_num += len(instance) mask_rcnn_inference(border_mask_logits, instances) return instances
def forward(self, images, features, proposals, targets=None): """ See :meth:`ROIHeads.forward`. """ del images if self.training: assert targets proposals = self.label_and_sample_proposals(proposals, targets) del targets proposal_boxes = [x.proposal_boxes for x in proposals] #gt_labels = [x.gt_classes for x in proposals] gt_labels = torch.cat([x.gt_classes for x in proposals], dim=0) #gt_labels = torch.stack([x.gt_classes for x in proposals]) for k, v in features.items(): features[k] = self.pre_pooler(v) box_features = self.pooler([features[f] for f in self.in_features], proposal_boxes) M, C, H, W = box_features.shape flatten_features = box_features.view(M, -1) if self.training: loss, acc, category_accuracy = self.classifier( [flatten_features, gt_labels]) storage = get_event_storage() storage.put_scalar("acc", acc * 100) del features if self.mask_on: proposals, fg_selection_masks = select_foreground_proposals( proposals, self.num_classes) # Since the ROI feature transform is shared between boxes and masks, # we don't need to recompute features. The mask loss is only defined # on foreground proposals, so we need to select out the foreground # features. mask_features = box_features[torch.cat(fg_selection_masks, dim=0)] del box_features return [], {"cls_loss": loss} else: return flatten_features
def _forward_densepose(self, features, instances): """ Forward logic of the densepose prediction branch. Args: features (list[Tensor]): #level input features for densepose prediction instances (list[Instances]): the per-image instances to train/predict densepose. In training, they can be the proposals. In inference, they can be the predicted boxes. Returns: In training, a dict of losses. In inference, update `instances` with new fields "densepose" and return it. """ if not self.densepose_on: return {} if self.training else instances if self.training: proposals, _ = select_foreground_proposals(instances, self.num_classes) proposals_dp = self.densepose_data_filter(proposals) if len(proposals_dp) > 0: proposal_boxes = [x.proposal_boxes for x in proposals_dp] features_dp = self.densepose_pooler(features, proposal_boxes) densepose_head_outputs = self.densepose_head(features_dp) densepose_outputs, _ = self.densepose_predictor( densepose_head_outputs) densepose_loss_dict = self.densepose_losses( proposals_dp, densepose_outputs) return densepose_loss_dict else: pred_boxes = [x.pred_boxes for x in instances] features_dp = self.densepose_pooler(features, pred_boxes) if len(features_dp) > 0: densepose_head_outputs = self.densepose_head(features_dp) densepose_outputs, _ = self.densepose_predictor( densepose_head_outputs) else: # If no detection occurred instances # set densepose_outputs to empty tensors empty_tensor = torch.zeros(size=(0, 0, 0, 0), device=features_dp.device) densepose_outputs = tuple([empty_tensor] * 4) densepose_inference(densepose_outputs, instances) return instances
def forward_dp_keypoint(self, features, instances): if not self.dp_keypoint_on: return {} if self.training else instances num_images = len(instances) if self.training: proposals, _ = select_foreground_proposals(instances, self.num_classes) proposals = select_proposals_with_visible_keypoints(proposals) # proposals = self.keypoint_data_filter(proposals) proposal_boxes = [x.proposal_boxes for x in proposals] # print('after:',len(proposal_boxes)) # if len(proposal_boxes) == 0: # return {"loss_keypoint": 0.} # if len(proposal_boxes) > 0: if self.use_mid: features = [self.mid_decoder(features)] keypoint_features = self.densepose_pooler(features, proposal_boxes) keypoint_output = self.densepose_head(keypoint_features) keypoint_logits = self.keypoint_predictor(keypoint_output) # The loss is defined on positive proposals with at >=1 visible keypoints. normalizer = ( num_images * self.batch_size_per_image * self.positive_sample_fraction * keypoint_logits.shape[1] ) # loss = dp_keypoint_rcnn_loss( # keypoint_logits, # instances, # normalizer=None if self.normalize_loss_by_visible_keypoints else normalizer, # ) loss = keypoint_rcnn_loss( keypoint_logits, proposals, normalizer=None if self.normalize_loss_by_visible_keypoints else normalizer, ) return {"loss_keypoint": loss * self.keypoint_loss_weight} else: if self.use_mid: features = [self.mid_decoder(features)] pred_boxes = [x.pred_boxes for x in instances] keypoint_features = self.densepose_pooler(features, pred_boxes) keypoint_output = self.densepose_head(keypoint_features) keypoint_logits = self.keypoint_predictor(keypoint_output) keypoint_rcnn_inference(keypoint_logits, instances) return instances
def forward_with_selected_boxes(self, features, instances, targets=None): assert self.training losses = {} if self._embedding_gt_box: proposals = targets else: proposals, _ = select_foreground_proposals(instances, self.num_classes) pred_instances = proposals if "roi_heads.mask_head" not in self._freeze: losses.update(self._forward_mask(features, pred_instances)) if "roi_heads.plane_head" not in self._freeze: plane_loss, pred_instances = self._forward_plane( features, pred_instances) losses.update(plane_loss) if "roi_heads.embedding_head" not in self._freeze: pred_instances = self._forward_embedding(features, pred_instances) return pred_instances, losses
def _forward_fiberlength(self, features, instances): """ Forward logic of the fiber length prediction branch. Args: features (list[Tensor]): #level input features for fiber length prediction instances (list[Instances]): the per-image instances to train/predict fiber lengths. In training, they can be the proposals. In inference, they can be the predicted boxes. Returns: In training, a dict of losses. In inference, update `instances` with new fields "pred_fiberlength" and return it. """ if not self.fiberlength_on: return {} if self.training else instances num_images = len(instances) if self.training: # The loss is defined on positive proposals with at >=1 visible keypoints. proposals, _ = select_foreground_proposals(instances, self.num_classes) proposals = select_proposals_with_visible_keypoints(proposals) proposal_boxes = [x.proposal_boxes for x in proposals] fiberlength_features = self.fiberlength_pooler( features, proposal_boxes) pred_fiberlengths = self.fiberlength_head(fiberlength_features) loss = fiberlength_loss(pred_fiberlengths, proposals) return {"loss_fiberlength": loss * self.fiberlength_loss_weight} else: pred_boxes = [x.pred_boxes for x in instances] fiberlength_features = self.fiberlength_pooler( features, pred_boxes) pred_fiberlengths = self.fiberlength_head(fiberlength_features) fiberlength_inference(pred_fiberlengths, instances) return instances
def _forward_keypoint( self, features: Dict[str, torch.Tensor], instances: List[Instances] ) -> Union[Dict[str, torch.Tensor], List[Instances]]: """ Forward logic of the keypoint prediction branch. Args: features (dict[str, Tensor]): mapping from feature map names to tensor. Same as in :meth:`ROIHeads.forward`. instances (list[Instances]): the per-image instances to train/predict keypoints. In training, they can be the proposals. In inference, they can be the boxes predicted by R-CNN box head. Returns: In training, a dict of losses. In inference, update `instances` with new fields "pred_keypoints" and return it. """ if not self.keypoint_on: return {} if self.training else instances if self.training: # head is only trained on positive proposals with >=1 visible keypoints. instances, _ = select_foreground_proposals(instances, self.num_classes) instances = select_proposals_with_visible_keypoints(instances) if self.keypoint_pooler is not None: features = [features[f] for f in self.keypoint_in_features] boxes = [ x.proposal_boxes if self.training else x.pred_boxes for x in instances ] features = self.keypoint_pooler(features, boxes) else: features = {f: features[f] for f in self.keypoint_in_features} return self.keypoint_head(features, instances)
def forward(self, batched_inputs): """ Args: batched_inputs: a list, batched outputs of :class:`DatasetMapper` . Each item in the list contains the inputs for one image. For now, each item in the list is a dict that contains: * image: Tensor, image in (C, H, W) format. * instances: Instances Other information that's included in the original dicts, such as: * "height", "width" (int): the output resolution of the model, used in inference. See :meth:`postprocess` for details. """ images, images_whwh = self.preprocess_image(batched_inputs) if isinstance(images, (list, torch.Tensor)): images = nested_tensor_from_tensor_list(images) # Feature Extraction. src = self.backbone(images.tensor) features = list() for f in self.in_features: feature = src[f] features.append(feature) # print('!!! pin1\n', len(features), features[0].size(), '\n!!!pin1') # Prepare Proposals. proposal_boxes = self.init_proposal_boxes.weight.clone() proposal_boxes = box_cxcywh_to_xyxy(proposal_boxes) proposal_boxes = proposal_boxes[None] * images_whwh[:, None, :] # proposal size: absolute size, x1y1, x2y2 # Prediction. outputs_class, outputs_coord, bboxes = self.head(features, proposal_boxes, self.init_proposal_features.weight) #TODO #3 mask forward # print('!!! pin3\n', bboxes.size(), '\n!!!pin3') output = {'pred_logits': outputs_class[-1], 'pred_boxes': outputs_coord[-1]} #TODO #3 mask forward # print('!!! pin2\n', mask_features.size(), '\n!!!pin2') proposal_list_instances = self.boxes2list_instances(bboxes, images.image_sizes) # print('!!! pin4\n', len(proposal_list_instances), len(proposal_list_instances[0]), '\n!!!pin4') if self.training: #TODO #3 mask forward gt_instances = [x["instances"].to(self.device) for x in batched_inputs] proposals_gt = self.label_and_sample_proposals(proposal_list_instances, gt_instances) # print('!!! pin5\n', len(proposals_gt), len(proposals_gt[0]), '\n!!!pin5') instances_fg, _ = select_foreground_proposals(proposals_gt, self.num_classes) # print('\ninstances_fg\n', len(proposals_gt), len(instances_fg), '\ninstances_fg\n') boxes_fg = [x.proposal_boxes for x in instances_fg] mask_features = self.mask_pooler(features, boxes_fg) # print('!!! pin6\n', len(instances_fg), len(instances_fg[0]), '\n!!!pin6') targets = self.prepare_targets(gt_instances) if self.deep_supervision: output['aux_outputs'] = [{'pred_logits': a, 'pred_boxes': b} for a, b in zip(outputs_class[:-1], outputs_coord[:-1])] loss_dict, match_indices = self.criterion(output, targets) # print('!!!pin1\n', loss_dict.keys(), '\n!!!pin1') #TODO #4 mask loss update loss_dict.update(self.mask_head(mask_features, instances_fg)) # print('!!!pin2\n', loss_dict.keys(), '\n!!!pin2') weight_dict = self.criterion.weight_dict for k in loss_dict.keys(): if k in weight_dict: loss_dict[k] *= weight_dict[k] return loss_dict else: box_cls = output["pred_logits"] box_pred = output["pred_boxes"] results = self.inference(box_cls, box_pred, images.image_sizes) #TODO #5 mask inference boxes_list = [] for boxes_per_image in bboxes: boxes_list.append(Boxes(boxes_per_image)) mask_features = self.mask_pooler(features, boxes_list) self.mask_head(mask_features, results) processed_results = [] for results_per_image, input_per_image, image_size in zip(results, batched_inputs, images.image_sizes): height = input_per_image.get("height", image_size[0]) width = input_per_image.get("width", image_size[1]) r = detector_postprocess(results_per_image, height, width) processed_results.append({"instances": r}) return processed_results
def _forward_shape(self, features, instances): """ Forward logic for the voxel and mesh refinement branch. Args: features (list[Tensor]): #level input features for voxel prediction instances (list[Instances]): the per-image instances to train/predict meshes. In training, they can be the proposals. In inference, they can be the predicted boxes. Returns: In training, a dict of losses. In inference, update `instances` with new fields "pred_voxels" & "pred_meshes" and return it. """ if not self.voxel_on and not self.mesh_on: return {} if self.training else instances if self.training: # The loss is only defined on positive proposals. proposals, _ = select_foreground_proposals(instances, self.num_classes) proposal_boxes = [x.proposal_boxes for x in proposals] losses = {} if self.voxel_on: voxel_features = self.voxel_pooler(features, proposal_boxes) voxel_logits = self.voxel_head(voxel_features) loss_voxel, target_voxels = voxel_rcnn_loss( voxel_logits, proposals, loss_weight=self.voxel_loss_weight) losses.update({"loss_voxel": loss_voxel}) if self._vis: self._misc["target_voxels"] = target_voxels if self.cls_agnostic_voxel: with torch.no_grad(): vox_in = voxel_logits.sigmoid().squeeze( 1) # (N, V, V, V) init_mesh = cubify(vox_in, self.cubify_thresh) # 1 else: raise ValueError( "No support for class specific predictions") if self.mesh_on: mesh_features = self.mesh_pooler(features, proposal_boxes) if not self.voxel_on: if mesh_features.shape[0] > 0: init_mesh = ico_sphere(self.ico_sphere_level, mesh_features.device) init_mesh = init_mesh.extend(mesh_features.shape[0]) else: init_mesh = Meshes(verts=[], faces=[]) pred_meshes = self.mesh_head(mesh_features, init_mesh) # loss weights loss_weights = { "chamfer": self.chamfer_loss_weight, "normals": self.normals_loss_weight, "edge": self.edge_loss_weight, } if not pred_meshes[0].isempty(): loss_chamfer, loss_normals, loss_edge, target_meshes = mesh_rcnn_loss( pred_meshes, proposals, loss_weights=loss_weights, gt_num_samples=self.gt_num_samples, pred_num_samples=self.pred_num_samples, gt_coord_thresh=self.gt_coord_thresh, ) if self._vis: self._misc["init_meshes"] = init_mesh self._misc["target_meshes"] = target_meshes else: loss_chamfer = sum( k.sum() for k in self.mesh_head.parameters()) * 0.0 loss_normals = sum( k.sum() for k in self.mesh_head.parameters()) * 0.0 loss_edge = sum(k.sum() for k in self.mesh_head.parameters()) * 0.0 losses.update({ "loss_chamfer": loss_chamfer, "loss_normals": loss_normals, "loss_edge": loss_edge, }) return losses else: pred_boxes = [x.pred_boxes for x in instances] if self.voxel_on: voxel_features = self.voxel_pooler(features, pred_boxes) voxel_logits = self.voxel_head(voxel_features) voxel_rcnn_inference(voxel_logits, instances) if self.cls_agnostic_voxel: with torch.no_grad(): vox_in = voxel_logits.sigmoid().squeeze( 1) # (N, V, V, V) init_mesh = cubify(vox_in, self.cubify_thresh) # 1 else: raise ValueError( "No support for class specific predictions") if self.mesh_on: mesh_features = self.mesh_pooler(features, pred_boxes) if not self.voxel_on: if mesh_features.shape[0] > 0: init_mesh = ico_sphere(self.ico_sphere_level, mesh_features.device) init_mesh = init_mesh.extend(mesh_features.shape[0]) else: init_mesh = Meshes(verts=[], faces=[]) pred_meshes = self.mesh_head(mesh_features, init_mesh) mesh_rcnn_inference(pred_meshes[-1], instances) else: assert self.voxel_on mesh_rcnn_inference(init_mesh, instances) return instances
def _forward_densepose(self, features, instances, extra=None): """ Forward logic of the densepose prediction branch. Args: features (list[Tensor]): #level input features for densepose prediction instances (list[Instances]): the per-image instances to train/predict densepose. In training, they can be the proposals. In inference, they can be the predicted boxes. Returns: In training, a dict of losses. In inference, update `instances` with new fields "densepose" and return it. """ if not self.densepose_on: return {} if self.training else instances bbox_locs_params = self.box_predictor.bbox_pred.weight bbox_cls_params = self.box_predictor.cls_score.weight with torch.no_grad(): bbox_params = torch.cat([bbox_cls_params, bbox_locs_params], dim=0) if self.training: proposals, _ = select_foreground_proposals(instances, self.num_classes) proposals_dp = self.densepose_data_filter(proposals) if len(proposals_dp) > 0: proposal_boxes = [x.proposal_boxes for x in proposals_dp] if self.use_mid: features = [self.mid_decoder(features)] if self.dp_semseg_on: segm_res, seg_losses, latent_features = self._forward_semsegs(features, instances, extra) else: if self.dp_semseg_on: segm_res, seg_losses, latent_features = self._forward_semsegs([features[-1]], instances, extra) features_dp = self.densepose_pooler(features, proposal_boxes) # print('roi pooler:',features_dp.size()) if self.inter_super_on: densepose_head_outputs, densepose_inter_head_outputs = self.densepose_head(features_dp) densepose_outputs, densepose_inter_outputs= self.densepose_predictor(densepose_head_outputs, densepose_inter_head_outputs) densepose_loss_dict = self.densepose_losses(proposals_dp, densepose_outputs, cls_emb_loss_on=True) # inter_loss_dict = self.densepose_losses(proposals_dp, densepose_inter_outputs, 'inter_') inter_loss_dict = self.densepose_inter_losses(proposals_dp, densepose_inter_outputs, 'inter_') for _, k in enumerate(inter_loss_dict.keys()): if k == 'inter_loss_densepose_M': densepose_loss_dict[k] = inter_loss_dict[k] * 1. else: densepose_loss_dict[k] = inter_loss_dict[k] * 0. else: densepose_head_outputs = self.densepose_head(features_dp) # densepose_outputs, _ = self.densepose_predictor(densepose_head_outputs, bbox_params) densepose_outputs, dp_outputs_from_kpt, dp_outputs_from_bbox = self.densepose_predictor(densepose_head_outputs, bbox_params) if self.dp_keypoint_on: keypoints_output = densepose_outputs[-1] densepose_outputs = densepose_outputs[:-1] densepose_loss_dict = self.densepose_losses(proposals_dp, densepose_outputs) if self.dp_keypoint_on: kpt_loss_dict = self._forward_dp_keypoint(keypoints_output, proposals_dp) for _, k in enumerate(kpt_loss_dict.keys()): densepose_loss_dict[k] = kpt_loss_dict[k] if self.dp_semseg_on: for _, k in enumerate(seg_losses.keys()): densepose_loss_dict[k] = seg_losses[k] if dp_outputs_from_kpt is not None: dp_kpt_loss_dict = self.densepose_losses(proposals_dp, dp_outputs_from_kpt) densepose_loss_dict['loss_densepose_I_from_kpt'] = dp_kpt_loss_dict['loss_densepose_I'] if dp_outputs_from_bbox is not None: dp_box_loss_dict = self.densepose_losses(proposals_dp, dp_outputs_from_bbox) densepose_loss_dict['loss_densepose_I_from_box'] = dp_box_loss_dict['loss_densepose_I'] return densepose_loss_dict else: pred_boxes = [x.pred_boxes for x in instances] if self.use_mid: features = [self.mid_decoder(features)] # if self.dp_semseg_on: # segm_res, seg_losses, latent_features = self._forward_semsegs(features, instances, extra) features_dp = self.densepose_pooler(features, pred_boxes) if len(features_dp) > 0: if self.inter_super_on: densepose_head_outputs, densepose_inter_head_outputs = self.densepose_head(features_dp) densepose_outputs, densepose_inter_outputs = self.densepose_predictor(densepose_head_outputs, densepose_inter_head_outputs) else: densepose_head_outputs = self.densepose_head(features_dp) densepose_outputs, _, _ = self.densepose_predictor(densepose_head_outputs,bbox_params) if self.dp_keypoint_on: keypoints_output = densepose_outputs[-1] densepose_outputs = densepose_outputs[:-1] instances = self._forward_dp_keypoint(keypoints_output, instances) else: # If no detection occured instances # set densepose_outputs to empty tensors empty_tensor = torch.zeros(size=(0, 0, 0, 0), device=features_dp.device) densepose_outputs = tuple([empty_tensor] * 5) densepose_inference(densepose_outputs, instances) return instances
def _forward_ama_densepose(self, features, instances, extra=None): """ Forward logic of the densepose prediction branch. Args: features (list[Tensor]): #level input features for densepose prediction instances (list[Instances]): the per-image instances to train/predict densepose. In training, they can be the proposals. In inference, they can be the predicted boxes. Returns: In training, a dict of losses. In inference, update `instances` with new fields "densepose" and return it. """ if not self.densepose_on: return {} if self.training else instances latent_features = None if self.dp_semseg_on: segm_res, seg_losses, latent_features = self._forward_semsegs(features, instances, extra) # features = latent_features if self.training: proposals, _ = select_foreground_proposals(instances, self.num_classes) proposals_dp = self.densepose_data_filter(proposals) proposal_boxes = [x.proposal_boxes for x in proposals_dp] features_dp = self.densepose_pooler(features, proposal_boxes) if len(proposals_dp) > 0: densepose_head_outputs, densepose_inter_head_outputs = self.densepose_head(features_dp, latent_features) bbox_locs_params = self.box_predictor.bbox_pred.weight bbox_cls_params = self.box_predictor.cls_score.weight bbox_params = torch.cat([bbox_cls_params, bbox_locs_params], dim=0) densepose_outputs, inter_output_1, inter_output_2 = self.densepose_predictor(densepose_head_outputs, densepose_inter_head_outputs, bbox_params) if self.dp_keypoint_on: keypoints_output = densepose_outputs[-1] densepose_outputs = densepose_outputs[:-1] densepose_loss_dict = self.densepose_losses(proposals_dp, densepose_outputs, cls_emb_loss_on=self.cfg.MODEL.ROI_DENSEPOSE_HEAD.IA_LOSS) inter_loss_dict = self.densepose_losses(proposals_dp, inter_output_1, prefix='inter_1_') for _, k in enumerate(inter_loss_dict.keys()): if 'loss_densepose_I' in k: densepose_loss_dict[k] = inter_loss_dict[k] else: densepose_loss_dict[k] = inter_loss_dict[k] * self.inter_weight inter_loss_dict = self.densepose_losses(proposals_dp, inter_output_2, prefix='inter_2_') for _, k in enumerate(inter_loss_dict.keys()): if 'loss_densepose_I' in k: densepose_loss_dict[k] = inter_loss_dict[k] else: densepose_loss_dict[k] = inter_loss_dict[k] * self.inter_weight if self.dp_semseg_on: for _, k in enumerate(seg_losses.keys()): densepose_loss_dict[k] = seg_losses[k] if self.dp_keypoint_on: kpt_loss_dict = self._forward_dp_keypoint(keypoints_output, proposals_dp) for _, k in enumerate(kpt_loss_dict.keys()): densepose_loss_dict[k] = kpt_loss_dict[k] # box_features = self.box_pooler(features, proposal_boxes) # box_features = self.box_head(box_features) return densepose_loss_dict else: pred_boxes = [x.pred_boxes for x in instances] features_dp = self.densepose_pooler(features, pred_boxes) # if len(features_dp[0]) == 0: # print('no proposal:', len(pred_boxes),pred_boxes) if len(features_dp[0]) > 0 and len(pred_boxes) > 0: densepose_head_outputs, densepose_inter_head_outputs = self.densepose_head(features_dp, latent_features) # densepose_head_outputs, densepose_inter_head_outputs = self.densepose_head(features_dp) bbox_locs_params = self.box_predictor.bbox_pred.weight bbox_cls_params = self.box_predictor.cls_score.weight bbox_params = torch.cat([bbox_cls_params, bbox_locs_params], dim=0) densepose_outputs, _, _ = self.densepose_predictor(densepose_head_outputs, densepose_inter_head_outputs, bbox_params) if self.dp_keypoint_on: keypoints_output = densepose_outputs[-1] densepose_outputs = densepose_outputs[:-1] instances=self._forward_dp_keypoint(keypoints_output, instances) else: # If no detection occured instances # set densepose_outputs to empty tensors empty_tensor = torch.zeros(size=(0, 0, 0, 0), device=features_dp[0].device) densepose_outputs = tuple([empty_tensor] * 5) densepose_inference(densepose_outputs, instances) # if self.dp_keypoint_on: # instances = self._forward_dp_keypoint(keypoints_output, instances) return instances
def _forward_ama_densepose(self, features, instances, extra=None): """ Forward logic of the densepose prediction branch. Args: features (list[Tensor]): #level input features for densepose prediction instances (list[Instances]): the per-image instances to train/predict densepose. In training, they can be the proposals. In inference, they can be the predicted boxes. Returns: In training, a dict of losses. In inference, update `instances` with new fields "densepose" and return it. """ if not self.densepose_on: return {} if self.training else instances latent_features = None if self.dp_semseg_on: segm_res, seg_losses, latent_features = self._forward_semsegs(features, instances, extra) # features = latent_features if self.training: proposals, _ = select_foreground_proposals(instances, self.num_classes) proposals_dp = self.densepose_data_filter(proposals) proposal_boxes = [x.proposal_boxes for x in proposals_dp] features_dp = self.densepose_pooler(features, proposal_boxes) # if len(proposals_dp) > 0 and len(features_dp[0]) > 0: if len(proposals_dp) > 0: if self.dp_semseg_on: # attention # attended_features = [] # for i_feat in range(len(features)): # if i_feat == 0: # i_segm_res = segm_res # else: # i_segm_res = F.interpolate(segm_res, # (int(segm_res.shape[2] / 2**i_feat), int(segm_res.shape[3] / 2**i_feat)), # mode="bilinear", # align_corners=False) # attended_features.append(features[i_feat] * i_segm_res) # latent_features = self.sematic_mask_pooler(attended_features, proposal_boxes) # att_mask = self.densepose_pooler([segm_res]*len(features_dp), proposal_boxes) latent_features = None #self.sematic_mask_pooler([latent_features], proposal_boxes)[0] # latent_features = torch.cat(latent_features, 1) # features_dp.append(latent_features) densepose_head_outputs, densepose_inter_head_outputs = self.densepose_head(features_dp, latent_features) densepose_outputs, inter_output_1, inter_output_2 = self.densepose_predictor(densepose_head_outputs, densepose_inter_head_outputs) if self.dp_keypoint_on: keypoints_output = densepose_outputs[-1] densepose_outputs = densepose_outputs[:-1] densepose_loss_dict = self.densepose_losses(proposals_dp, densepose_outputs, cls_emb_loss_on=self.cfg.MODEL.ROI_DENSEPOSE_HEAD.IA_LOSS) inter_loss_dict = self.densepose_losses(proposals_dp, inter_output_1, prefix='inter_1_') for _, k in enumerate(inter_loss_dict.keys()): densepose_loss_dict[k] = inter_loss_dict[k] * self.inter_weight inter_loss_dict = self.densepose_losses(proposals_dp, inter_output_2, prefix='inter_2_') for _, k in enumerate(inter_loss_dict.keys()): densepose_loss_dict[k] = inter_loss_dict[k] * self.inter_weight if self.dp_semseg_on: for _, k in enumerate(seg_losses.keys()): densepose_loss_dict[k] = seg_losses[k] if self.dp_keypoint_on: # keypoints_output, proposals_kp = select_proposals_idx_with_visible_keypoints(keypoints_output, proposals_dp) kpt_loss_dict = self._forward_dp_keypoint(keypoints_output, proposals_dp) for _, k in enumerate(kpt_loss_dict.keys()): densepose_loss_dict[k] = kpt_loss_dict[k] return densepose_loss_dict else: pred_boxes = [x.pred_boxes for x in instances] features_dp = self.densepose_pooler(features, pred_boxes) # if len(features_dp[0]) == 0: # print('no proposal:', len(pred_boxes),pred_boxes) if len(features_dp[0]) > 0 and len(pred_boxes) > 0: if self.dp_semseg_on: latent_features = None#self.sematic_mask_pooler([latent_features], pred_boxes)[0] # latent_features = torch.cat(latent_features, 1) # features_dp.append(latent_features) # if self.dp_semseg_on: # attended_features = [] # for i_feat in range(len(features)): # if i_feat == 0: # i_segm_res = segm_res # else: # i_segm_res = F.interpolate(segm_res, # (int(segm_res.shape[2] / 2**i_feat), int(segm_res.shape[3] / 2**i_feat)), # mode="bilinear", # align_corners=False) # attended_features.append(features[i_feat] * i_segm_res) # latent_features = self.sematic_mask_pooler(attended_features, pred_boxes) # if features_dp[0].shape[0] == 0: # print('no proposal:',len(pred_boxes)) densepose_head_outputs, densepose_inter_head_outputs = self.densepose_head(features_dp, latent_features) # densepose_head_outputs, densepose_inter_head_outputs = self.densepose_head(features_dp) densepose_outputs, _, _ = self.densepose_predictor(densepose_head_outputs, densepose_inter_head_outputs) if self.dp_keypoint_on: keypoints_output = densepose_outputs[-1] densepose_outputs = densepose_outputs[:-1] instances=self._forward_dp_keypoint(keypoints_output, instances) else: # If no detection occured instances # set densepose_outputs to empty tensors empty_tensor = torch.zeros(size=(0, 0, 0, 0), device=features_dp[0].device) densepose_outputs = tuple([empty_tensor] * 5) densepose_inference(densepose_outputs, instances) # if self.dp_keypoint_on: # instances = self._forward_dp_keypoint(keypoints_output, instances) return instances
def _forward_densepose(self, features, instances): """ Forward logic of the densepose prediction branch. Args: features (list[Tensor]): #level input features for densepose prediction instances (list[Instances]): the per-image instances to train/predict densepose. In training, they can be the proposals. In inference, they can be the predicted boxes. Returns: In training, a dict of losses. In inference, update `instances` with new fields "densepose" and return it. """ if not self.densepose_on: return {} if self.training else instances if self.training: proposals, _ = select_foreground_proposals(instances, self.num_classes) proposals_dp = self.densepose_data_filter(proposals) if len(proposals_dp) > 0: proposal_boxes = [x.proposal_boxes for x in proposals_dp] if self.use_mid: features = [self.mid_decoder(features)] features_dp = self.densepose_pooler(features, proposal_boxes) # print('roi pooler:',features_dp.size()) if self.inter_super_on: densepose_head_outputs, densepose_inter_head_outputs = self.densepose_head(features_dp) densepose_outputs, densepose_inter_outputs= self.densepose_predictor(densepose_head_outputs, densepose_inter_head_outputs) densepose_loss_dict = self.densepose_losses(proposals_dp, densepose_outputs, cls_emb_loss_on=True) # inter_loss_dict = self.densepose_losses(proposals_dp, densepose_inter_outputs, 'inter_') inter_loss_dict = self.densepose_inter_losses(proposals_dp, densepose_inter_outputs, 'inter_') for _, k in enumerate(inter_loss_dict.keys()): if k == 'inter_loss_densepose_M': densepose_loss_dict[k] = inter_loss_dict[k] * 1. else: densepose_loss_dict[k] = inter_loss_dict[k] * 0. else: densepose_head_outputs = self.densepose_head(features_dp) densepose_outputs, _ = self.densepose_predictor(densepose_head_outputs) if self.dp_keypoint_on: keypoints_output = densepose_outputs[-1] densepose_outputs = densepose_outputs[:-1] densepose_loss_dict = self.densepose_losses(proposals_dp, densepose_outputs) if self.dp_keypoint_on: kpt_loss_dict = self._forward_dp_keypoint(keypoints_output, proposals_dp) for _, k in enumerate(kpt_loss_dict.keys()): densepose_loss_dict[k] = kpt_loss_dict[k] return densepose_loss_dict else: pred_boxes = [x.pred_boxes for x in instances] if self.use_mid: features = [self.mid_decoder(features)] features_dp = self.densepose_pooler(features, pred_boxes) if len(features_dp) > 0: if self.inter_super_on: densepose_head_outputs, densepose_inter_head_outputs = self.densepose_head(features_dp) densepose_outputs, densepose_inter_outputs = self.densepose_predictor(densepose_head_outputs, densepose_inter_head_outputs) else: densepose_head_outputs = self.densepose_head(features_dp) densepose_outputs, _ = self.densepose_predictor(densepose_head_outputs) if self.dp_keypoint_on: keypoints_output = densepose_outputs[-1] densepose_outputs = densepose_outputs[:-1] instances = self._forward_dp_keypoint(keypoints_output, instances) else: # If no detection occured instances # set densepose_outputs to empty tensors empty_tensor = torch.zeros(size=(0, 0, 0, 0), device=features_dp.device) densepose_outputs = tuple([empty_tensor] * 5) densepose_inference(densepose_outputs, instances) return instances
def forward(self, images, features, proposals, targets=None): """ See :class:`ROIHeads.forward`. """ del images if self.training: proposals = self.label_and_sample_proposals(proposals, targets) del targets if self.training: losses = self._forward_box( [features[f] for f in self.in_features_box], proposals) # During training the proposals used by the box head are # used by the mask, keypoint (and densepose) heads. foreground_proposals, _ = select_foreground_proposals( proposals, self.num_classes) shared_features = self._forward_shared_block( [features[f] for f in self.in_features_shared], [x.proposal_boxes for x in foreground_proposals]) losses.update( self._forward_mask(shared_features, foreground_proposals)) losses.update( self._forward_keypoint(shared_features, foreground_proposals)) losses.update( self._forward_densepose(shared_features, foreground_proposals)) return proposals, losses else: has_adaptive_pool = detectron2.modeling.poolers.ROIAlign.__name__ == "AdaptivePool" or detectron2.modeling.poolers.RoIPool.__name__ == "AdaptivePool" if has_adaptive_pool: # check how it is done in apply deltas # try do it with tensors, avoiding loops level_assignments = assign_boxes_to_levels( [x.proposal_boxes for x in proposals], self.box_min_level, self.box_max_level, self.box_canonical_box_size, self.box_canonical_level) inds = [] for level in range(len(self.box_pooler_scales)): inds.append( torch.nonzero(level_assignments == level).squeeze(1)) i = 0 for image_idx in range(len(proposals)): boxes = proposals[image_idx].proposal_boxes.tensor for level, spatial_scale in enumerate( self.box_pooler_scales): inds_ = inds[level][i:i + len(boxes)] boxes[inds_, 0] = ( (boxes[inds_, 0] * spatial_scale - 0.5).floor() + 0.5) / spatial_scale boxes[inds_, 1] = ( (boxes[inds_, 1] * spatial_scale - 0.5).floor() + 0.5) / spatial_scale boxes[inds_, 2] = ( (boxes[inds_, 2] * spatial_scale - 0.5).ceil() + 0.5) / spatial_scale boxes[inds_, 3] = ( (boxes[inds_, 3] * spatial_scale - 0.5).ceil() + 0.5) / spatial_scale proposals[image_idx].proposal_boxes.tensor = boxes i += len(boxes) pred_instances = self._forward_box( [features[f] for f in self.in_features_box], proposals) if has_adaptive_pool: level_assignments = assign_boxes_to_levels( [x.pred_boxes for x in pred_instances], self.shared_min_level, self.shared_max_level, self.shared_canonical_box_size, self.shared_canonical_level) inds = [] for level in range(len(self.shared_pooler_scales)): inds.append( torch.nonzero(level_assignments == level).squeeze(1)) i = 0 for image_idx in range(len(pred_instances)): boxes = pred_instances[image_idx].pred_boxes.tensor for level, spatial_scale in enumerate( self.shared_pooler_scales): inds_ = inds[level][i:i + len(boxes)] boxes[inds_, 0] = ( (boxes[inds_, 0] * spatial_scale - 0.5).floor() + 0.5) / spatial_scale boxes[inds_, 1] = ( (boxes[inds_, 1] * spatial_scale - 0.5).floor() + 0.5) / spatial_scale boxes[inds_, 2] = ( (boxes[inds_, 2] * spatial_scale - 0.5).ceil() + 0.5) / spatial_scale boxes[inds_, 3] = ( (boxes[inds_, 3] * spatial_scale - 0.5).ceil() + 0.5) / spatial_scale pred_instances[image_idx].pred_boxes.tensor = boxes i += len(boxes) # During inference cascaded prediction is used: the mask and keypoints heads are only # applied to the top scoring box detections. shared_features = self._forward_shared_block( [features[f] for f in self.in_features_shared], [x.pred_boxes for x in pred_instances]) pred_instances = self.forward_with_given_boxes( shared_features, pred_instances) return pred_instances, {}