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