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, {}
