def _forward_mask( self, features: Dict[str, torch.Tensor], box_features: torch.Tensor, instances: List[Instances], ) -> Union[Dict[str, torch.Tensor], List[Instances]]: if not self.mask_on: return {} if self.training else instances features = [features[f] for f in self.in_features] if self.training: # The loss is only defined on positive proposals. proposals, fg_selection_masks = select_foreground_proposals( instances, self.num_classes) proposal_boxes = [x.proposal_boxes for x in proposals] mask_features = self.mask_pooler(features, proposal_boxes) box_features = box_features[torch.cat(fg_selection_masks, dim=0)] cam_raw_batch = self.calculate_cams(mask_features, box_features, proposals, self.training) return self.mask_head(mask_features, cam_raw_batch, proposals) else: pred_boxes = [x.pred_boxes for x in instances] mask_features = self.mask_pooler(features, pred_boxes) cam_raw_batch = self.calculate_cams(mask_features, box_features, instances, self.training) return self.mask_head(mask_features, cam_raw_batch, instances)
def _forward_mask(self, features: Dict[str, torch.Tensor], instances: 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 boxes predicted by R-CNN box head. Returns: In training, a dict of losses. In inference, update `instances` with new fields "pred_masks" and return it. """ if self.training: # head is only trained on positive proposals. 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) if features.shape[0] == 0: return None else: features = {f: features[f] for f in self.mask_in_features} if features.shape[0] == 0: return None return self.mask_head(features, instances)
def forward(self, x, proposal_boxes=None): if x.dim() > 2: x = torch.flatten(x, start_dim=1) scores = self.cls_score(x) proposal_deltas = self.bbox_pred(x) if self.attr_on: if self.training: assert proposal_boxes is not None, "Proposals are None while attr=True" proposals, fg_selection_atrributes = select_foreground_proposals( proposal_boxes, self.num_classes) attribute_features = x[torch.cat(fg_selection_atrributes, dim=0)] cls_labels = torch.cat([prop.gt_classes for prop in proposals]) else: # get labels and indices of proposals with foreground cls_labels = torch.argmax(scores, dim=1) attribute_features = x # get embeddings of indices using gt cls labels cls_embed_out = self.cls_embed(cls_labels) # concat with fc7 feats concat_attr = cat([attribute_features, cls_embed_out], dim=1) # pass through attr head layers fc_attr = self.attr_linear1(concat_attr) attr_score = self.attr_linear2(F.relu(fc_attr)) return scores, proposal_deltas, attr_score, cat( [p.gt_attributes for p in proposals], dim=0) if self.training else None return scores, proposal_deltas
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 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 features = [features[f] for f in self.in_features] 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) return self.keypoint_head(keypoint_features, proposals) else: pred_boxes = [x.pred_boxes for x in instances] keypoint_features = self.keypoint_pooler(features, pred_boxes) return self.keypoint_head(keypoint_features, 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 features = [features[f] for f in self.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(features, proposal_boxes) return self.mask_head(mask_features, proposals) else: pred_boxes = [x.pred_boxes for x in instances] second_stream_outputs = pred_boxes return second_stream_outputs
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 (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 features = [features[f] for f in self.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(features, proposal_boxes) return self.mask_head(mask_features, proposals) else: pred_boxes = [x.pred_boxes for x in instances] mask_features = self.mask_pooler(features, pred_boxes) mask_logits = self.mask_head.layers(mask_features) mask_rcnn_inference(mask_logits, instances) return instances
def _forward_contact(self, second_stream_outputs, features, instances): """ Forward logic of the contact_head branch. Args: features (list[Tensor]): #level input features for contact prediction instances (list[Instances]): the per-image instances to train/predict contact. In training, they can be proposals. In inference, they can be predicted boxes. Returns: In training, a dict of losses. In inference, update `instances` with new field "contact" and return it. """ if not self.contact_head_on: return {} if self.training else instances features = [features[f] for f in self.in_features] if self.training: proposals, _ = select_foreground_proposals(instances, self.num_classes) proposal_boxes = [x.proposal_boxes for x in proposals] hand_features, object_features, hand_object_features = self.get_pairwise_features( features, proposal_boxes, second_stream_outputs, use_intersection=self.use_intersection ) return self.contact_head(hand_features, hand_object_features, proposals) else: pred_boxes = [x.pred_boxes for x in instances] hand_features, object_features, hand_object_features = self.get_pairwise_features( features, pred_boxes, second_stream_outputs, use_intersection=self.use_intersection ) return self.contact_head(hand_features, hand_object_features, instances)
def forward(self, images, features, proposals, targets=None): 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] box_features = self._shared_roi_transform( [features[f] for f in self.in_features], proposal_boxes) feature_pooled = box_features.mean(dim=[2, 3]) predictions = self.box_predictor(feature_pooled) if self.training: del features losses = self.box_predictor.losses(predictions, proposals) if self.mask_on: proposals, fg_selection_masks = select_foreground_proposals( proposals, self.num_classes) mask_features = box_features[torch.cat(fg_selection_masks, dim=0)] del box_features losses.update(self.mask_head(mask_features, proposals)) if self.attribute_on: losses.update( self.attribute_head.forward_attribute_loss( proposals, feature_pooled)) return [], losses else: pred_instances, _ = self.box_predictor.inference( predictions, proposals) pred_instances = self.forward_with_given_boxes( features, pred_instances) return pred_instances, {}
def _forward_prop( self, features: Dict[str, torch.Tensor], instances: List[Instances] ) -> Union[Dict[str, torch.Tensor], List[Instances]]: """Forward logic of the box prediction branch. If `self.train_on_pred_boxes is True`, the function puts predicted boxes in the `proposal_boxes` field of `proposals` argument. Args: features (dict[str, Tensor]): mapping from feature map names to tensor. Same as in :meth:`ROIHeads.forward`. proposals (list[Instances]): the per-image object proposals with their matching ground truth. Each has fields "proposal_boxes", and "objectness_logits", "gt_classes", "gt_boxes". Returns: In training, a dict of losses. In inference, a list of `Instances`, the predicted instances. """ features = [features[f] for f in self.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] prop_features = self.prop_pooler(features, proposal_boxes) # logger.info("forward_prop feat {}, foreground proposals {}".format(prop_features.shape, proposals)) return self.prop_head(prop_features, proposals) else: pred_boxes = [x.pred_boxes for x in instances] prop_features = self.prop_pooler(features, pred_boxes) return self.prop_head(prop_features, instances)
def _forward_densepose(self, features: Dict[str, torch.Tensor], instances: List[Instances]): """ Forward logic of the densepose prediction branch. Args: features (dict[str, Tensor]): input data as a mapping from feature map name to tensor. Axis 0 represents the number of images `N` in the input data; axes 1-3 are channels, height, and width, which may vary between feature maps (e.g., if a feature pyramid is used). instances (list[Instances]): length `N` list of `Instances`. The i-th `Instances` contains instances for the i-th input image, 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 features = [features[f] for f in self.in_features] if self.training: proposals, _ = select_foreground_proposals(instances, self.num_classes) features, proposals = self.densepose_data_filter( features, proposals) if len(proposals) > 0: proposal_boxes = [x.proposal_boxes for x in proposals] if self.use_decoder: features = [self.decoder(features)] features_dp = self.densepose_pooler(features, proposal_boxes) densepose_head_outputs = self.densepose_head(features_dp) densepose_predictor_outputs = self.densepose_predictor( densepose_head_outputs) densepose_loss_dict = self.densepose_losses( proposals, densepose_predictor_outputs, embedder=self.embedder) return densepose_loss_dict else: pred_boxes = [x.pred_boxes for x in instances] if self.use_decoder: features = [self.decoder(features)] features_dp = self.densepose_pooler(features, pred_boxes) if len(features_dp) > 0: densepose_head_outputs = self.densepose_head(features_dp) densepose_predictor_outputs = self.densepose_predictor( densepose_head_outputs) else: densepose_predictor_outputs = None densepose_inference(densepose_predictor_outputs, instances) return instances
def _forward_densepose(self, features: List[torch.Tensor], instances: List[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 features = [features[f] for f in self.in_features] if self.training: proposals, _ = select_foreground_proposals(instances, self.num_classes) proposals_dp = self.densepose_data_filter(proposals) if len(proposals_dp) > 0: # NOTE may deadlock in DDP if certain workers have empty proposals_dp proposal_boxes = [x.proposal_boxes for x in proposals_dp] if self.use_decoder: features = [self.decoder(features)] features_dp = self.densepose_pooler(features, proposal_boxes) densepose_head_outputs = self.densepose_head(features_dp) densepose_outputs, _, confidences, _ = self.densepose_predictor( densepose_head_outputs) densepose_loss_dict = self.densepose_losses( proposals_dp, densepose_outputs, confidences) return densepose_loss_dict else: pred_boxes = [x.pred_boxes for x in instances] if self.use_decoder: features = [self.decoder(features)] features_dp = self.densepose_pooler(features, pred_boxes) if len(features_dp) > 0: densepose_head_outputs = self.densepose_head(features_dp) densepose_outputs, _, confidences, _ = 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) confidences = tuple([empty_tensor] * 4) densepose_inference(densepose_outputs, confidences, instances) return instances
def forward_attr(self, proposals, box_features): proposals, fg_selection_attributes = select_foreground_proposals( proposals, self.num_classes ) attribute_features = box_features[torch.cat(fg_selection_attributes, dim=0)] obj_labels = torch.cat([p.gt_classes for p in proposals]) attribute_labels = torch.cat([p.gt_attributes for p in proposals], dim=0) attribute_scores = self.attribute_predictor(attribute_features, obj_labels) return attribute_scores
def forward(self, images, features, proposals, targets=None): 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] box_features = self._shared_roi_transform( [features[f] for f in self.in_features], proposal_boxes ) feature_pooled = box_features.mean(dim=[2, 3]) predictions = self.box_predictor(feature_pooled) if self.training: del features losses = self.box_predictor.losses(predictions, proposals) if self.mask_on: proposals, fg_selection_masks = select_foreground_proposals( proposals, self.num_classes ) mask_features = box_features[torch.cat(fg_selection_masks, dim=0)] del box_features losses.update(self.mask_head(mask_features, proposals)) if self.attribute_on: losses.update(self.forward_attribute_loss(proposals, feature_pooled)) return [], losses elif self.extract_on: pred_class_logits, pred_proposal_deltas = predictions # pred_class_logits = pred_class_logits[:, :-1] # background is last cls_lables = torch.argmax(pred_class_logits, dim=1) num_preds_per_image = [len(p) for p in proposals] if self.extractor_mode == 1 or self.extractor_mode == 3: if self.attribute_on: attr_scores = self.forward_attribute_score(feature_pooled, cls_lables) return proposal_boxes, self.predict_probs(pred_class_logits, num_preds_per_image), feature_pooled.split(num_preds_per_image, dim=0), attr_scores.split(num_preds_per_image, dim=0) else: return proposal_boxes, self.predict_probs(pred_class_logits, num_preds_per_image), feature_pooled.split(num_preds_per_image, dim=0) elif self.extractor_mode == 2: return self.predict_boxes(proposals, pred_proposal_deltas, num_preds_per_image), self.predict_probs(pred_class_logits, num_preds_per_image) else: raise ValueError('BUA.EXTRATOR.MODE ERROR') else: pred_instances, _ = self.box_predictor.inference(predictions, proposals) pred_instances = self.forward_with_given_boxes(features, pred_instances) return pred_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] box_features = self._shared_roi_transform( [features[f] for f in self.in_features], proposal_boxes ) # @Will Lee 2020/7/22 Res5ROIHeads中直接使用box_feature.mean(dim=[2,3])作为predictor的输入,这里做了变动 # predictions = self.box_predictor(box_features.mean(dim=[2, 3])) predictions = self.box_predictor(box_features) if self.training: del features losses = self.box_predictor.losses(predictions, proposals) 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 losses.update(self.mask_head(mask_features, proposals)) return [], losses else: pred_instances, _ = self.box_predictor.inference(predictions, proposals) pred_instances = self.forward_with_given_boxes(features, pred_instances) return pred_instances, {}