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