Spaces:
Running
on
Zero
Running
on
Zero
| # Copyright (c) Facebook, Inc. and its affiliates. | |
| from typing import Dict | |
| import torch | |
| from torch.nn import functional as F | |
| from detectron2.structures.boxes import Boxes, BoxMode | |
| from ..structures import ( | |
| DensePoseChartPredictorOutput, | |
| DensePoseChartResult, | |
| DensePoseChartResultWithConfidences, | |
| ) | |
| from . import resample_fine_and_coarse_segm_to_bbox | |
| from .base import IntTupleBox, make_int_box | |
| def resample_uv_tensors_to_bbox( | |
| u: torch.Tensor, | |
| v: torch.Tensor, | |
| labels: torch.Tensor, | |
| box_xywh_abs: IntTupleBox, | |
| ) -> torch.Tensor: | |
| """ | |
| Resamples U and V coordinate estimates for the given bounding box | |
| Args: | |
| u (tensor [1, C, H, W] of float): U coordinates | |
| v (tensor [1, C, H, W] of float): V coordinates | |
| labels (tensor [H, W] of long): labels obtained by resampling segmentation | |
| outputs for the given bounding box | |
| box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs | |
| Return: | |
| Resampled U and V coordinates - a tensor [2, H, W] of float | |
| """ | |
| x, y, w, h = box_xywh_abs | |
| w = max(int(w), 1) | |
| h = max(int(h), 1) | |
| u_bbox = F.interpolate(u, (h, w), mode="bilinear", align_corners=False) | |
| v_bbox = F.interpolate(v, (h, w), mode="bilinear", align_corners=False) | |
| uv = torch.zeros([2, h, w], dtype=torch.float32, device=u.device) | |
| for part_id in range(1, u_bbox.size(1)): | |
| uv[0][labels == part_id] = u_bbox[0, part_id][labels == part_id] | |
| uv[1][labels == part_id] = v_bbox[0, part_id][labels == part_id] | |
| return uv | |
| def resample_uv_to_bbox( | |
| predictor_output: DensePoseChartPredictorOutput, | |
| labels: torch.Tensor, | |
| box_xywh_abs: IntTupleBox, | |
| ) -> torch.Tensor: | |
| """ | |
| Resamples U and V coordinate estimates for the given bounding box | |
| Args: | |
| predictor_output (DensePoseChartPredictorOutput): DensePose predictor | |
| output to be resampled | |
| labels (tensor [H, W] of long): labels obtained by resampling segmentation | |
| outputs for the given bounding box | |
| box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs | |
| Return: | |
| Resampled U and V coordinates - a tensor [2, H, W] of float | |
| """ | |
| return resample_uv_tensors_to_bbox( | |
| predictor_output.u, | |
| predictor_output.v, | |
| labels, | |
| box_xywh_abs, | |
| ) | |
| def densepose_chart_predictor_output_to_result( | |
| predictor_output: DensePoseChartPredictorOutput, boxes: Boxes | |
| ) -> DensePoseChartResult: | |
| """ | |
| Convert densepose chart predictor outputs to results | |
| Args: | |
| predictor_output (DensePoseChartPredictorOutput): DensePose predictor | |
| output to be converted to results, must contain only 1 output | |
| boxes (Boxes): bounding box that corresponds to the predictor output, | |
| must contain only 1 bounding box | |
| Return: | |
| DensePose chart-based result (DensePoseChartResult) | |
| """ | |
| assert len(predictor_output) == 1 and len(boxes) == 1, ( | |
| f"Predictor output to result conversion can operate only single outputs" | |
| f", got {len(predictor_output)} predictor outputs and {len(boxes)} boxes" | |
| ) | |
| boxes_xyxy_abs = boxes.tensor.clone() | |
| boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS) | |
| box_xywh = make_int_box(boxes_xywh_abs[0]) | |
| labels = resample_fine_and_coarse_segm_to_bbox(predictor_output, box_xywh).squeeze(0) | |
| uv = resample_uv_to_bbox(predictor_output, labels, box_xywh) | |
| return DensePoseChartResult(labels=labels, uv=uv) | |
| def resample_confidences_to_bbox( | |
| predictor_output: DensePoseChartPredictorOutput, | |
| labels: torch.Tensor, | |
| box_xywh_abs: IntTupleBox, | |
| ) -> Dict[str, torch.Tensor]: | |
| """ | |
| Resamples confidences for the given bounding box | |
| Args: | |
| predictor_output (DensePoseChartPredictorOutput): DensePose predictor | |
| output to be resampled | |
| labels (tensor [H, W] of long): labels obtained by resampling segmentation | |
| outputs for the given bounding box | |
| box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs | |
| Return: | |
| Resampled confidences - a dict of [H, W] tensors of float | |
| """ | |
| x, y, w, h = box_xywh_abs | |
| w = max(int(w), 1) | |
| h = max(int(h), 1) | |
| confidence_names = [ | |
| "sigma_1", | |
| "sigma_2", | |
| "kappa_u", | |
| "kappa_v", | |
| "fine_segm_confidence", | |
| "coarse_segm_confidence", | |
| ] | |
| confidence_results = {key: None for key in confidence_names} | |
| confidence_names = [ | |
| key for key in confidence_names if getattr(predictor_output, key) is not None | |
| ] | |
| confidence_base = torch.zeros([h, w], dtype=torch.float32, device=predictor_output.u.device) | |
| # assign data from channels that correspond to the labels | |
| for key in confidence_names: | |
| resampled_confidence = F.interpolate( | |
| getattr(predictor_output, key), | |
| (h, w), | |
| mode="bilinear", | |
| align_corners=False, | |
| ) | |
| result = confidence_base.clone() | |
| for part_id in range(1, predictor_output.u.size(1)): | |
| if resampled_confidence.size(1) != predictor_output.u.size(1): | |
| # confidence is not part-based, don't try to fill it part by part | |
| continue | |
| result[labels == part_id] = resampled_confidence[0, part_id][labels == part_id] | |
| if resampled_confidence.size(1) != predictor_output.u.size(1): | |
| # confidence is not part-based, fill the data with the first channel | |
| # (targeted for segmentation confidences that have only 1 channel) | |
| result = resampled_confidence[0, 0] | |
| confidence_results[key] = result | |
| return confidence_results # pyre-ignore[7] | |
| def densepose_chart_predictor_output_to_result_with_confidences( | |
| predictor_output: DensePoseChartPredictorOutput, boxes: Boxes | |
| ) -> DensePoseChartResultWithConfidences: | |
| """ | |
| Convert densepose chart predictor outputs to results | |
| Args: | |
| predictor_output (DensePoseChartPredictorOutput): DensePose predictor | |
| output with confidences to be converted to results, must contain only 1 output | |
| boxes (Boxes): bounding box that corresponds to the predictor output, | |
| must contain only 1 bounding box | |
| Return: | |
| DensePose chart-based result with confidences (DensePoseChartResultWithConfidences) | |
| """ | |
| assert len(predictor_output) == 1 and len(boxes) == 1, ( | |
| f"Predictor output to result conversion can operate only single outputs" | |
| f", got {len(predictor_output)} predictor outputs and {len(boxes)} boxes" | |
| ) | |
| boxes_xyxy_abs = boxes.tensor.clone() | |
| boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS) | |
| box_xywh = make_int_box(boxes_xywh_abs[0]) | |
| labels = resample_fine_and_coarse_segm_to_bbox(predictor_output, box_xywh).squeeze(0) | |
| uv = resample_uv_to_bbox(predictor_output, labels, box_xywh) | |
| confidences = resample_confidences_to_bbox(predictor_output, labels, box_xywh) | |
| return DensePoseChartResultWithConfidences(labels=labels, uv=uv, **confidences) | |