Spaces:
Running
on
Zero
Running
on
Zero
File size: 7,178 Bytes
fb9d4c3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 |
# 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)
|