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_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