| | |
| |
|
| | from typing import Any, List |
| | import torch |
| | from torch.nn import functional as F |
| |
|
| | from detectron2.config import CfgNode |
| | from detectron2.structures import Instances |
| |
|
| | from .utils import resample_data |
| |
|
| |
|
| | class SegmentationLoss: |
| | """ |
| | Segmentation loss as cross-entropy for raw unnormalized scores given ground truth |
| | labels. Segmentation ground truth labels are defined for the bounding box of |
| | interest at some fixed resolution [S, S], where |
| | S = MODEL.ROI_DENSEPOSE_HEAD.HEATMAP_SIZE. |
| | """ |
| |
|
| | def __init__(self, cfg: CfgNode): |
| | """ |
| | Initialize segmentation loss from configuration options |
| | |
| | Args: |
| | cfg (CfgNode): configuration options |
| | """ |
| | self.heatmap_size = cfg.MODEL.ROI_DENSEPOSE_HEAD.HEATMAP_SIZE |
| | self.n_segm_chan = cfg.MODEL.ROI_DENSEPOSE_HEAD.NUM_COARSE_SEGM_CHANNELS |
| |
|
| | def __call__( |
| | self, |
| | proposals_with_gt: List[Instances], |
| | densepose_predictor_outputs: Any, |
| | packed_annotations: Any, |
| | ) -> torch.Tensor: |
| | """ |
| | Compute segmentation loss as cross-entropy on aligned segmentation |
| | ground truth and estimated scores. |
| | |
| | Args: |
| | proposals_with_gt (list of Instances): detections with associated ground truth data |
| | densepose_predictor_outputs: an object of a dataclass that contains predictor outputs |
| | with estimated values; assumed to have the following attributes: |
| | * coarse_segm - coarse segmentation estimates, tensor of shape [N, D, S, S] |
| | packed_annotations: packed annotations for efficient loss computation; |
| | the following attributes are used: |
| | - coarse_segm_gt |
| | - bbox_xywh_gt |
| | - bbox_xywh_est |
| | """ |
| | if packed_annotations.coarse_segm_gt is None: |
| | return self.fake_value(densepose_predictor_outputs) |
| | coarse_segm_est = densepose_predictor_outputs.coarse_segm[packed_annotations.bbox_indices] |
| | with torch.no_grad(): |
| | coarse_segm_gt = resample_data( |
| | packed_annotations.coarse_segm_gt.unsqueeze(1), |
| | packed_annotations.bbox_xywh_gt, |
| | packed_annotations.bbox_xywh_est, |
| | self.heatmap_size, |
| | self.heatmap_size, |
| | mode="nearest", |
| | padding_mode="zeros", |
| | ).squeeze(1) |
| | if self.n_segm_chan == 2: |
| | coarse_segm_gt = coarse_segm_gt > 0 |
| | return F.cross_entropy(coarse_segm_est, coarse_segm_gt.long()) |
| |
|
| | def fake_value(self, densepose_predictor_outputs: Any) -> torch.Tensor: |
| | """ |
| | Fake segmentation loss used when no suitable ground truth data |
| | was found in a batch. The loss has a value 0 and is primarily used to |
| | construct the computation graph, so that `DistributedDataParallel` |
| | has similar graphs on all GPUs and can perform reduction properly. |
| | |
| | Args: |
| | densepose_predictor_outputs: DensePose predictor outputs, an object |
| | of a dataclass that is assumed to have `coarse_segm` |
| | attribute |
| | Return: |
| | Zero value loss with proper computation graph |
| | """ |
| | return densepose_predictor_outputs.coarse_segm.sum() * 0 |
| |
|
