|
|
|
from dataclasses import fields
|
|
import torch
|
|
|
|
from densepose.structures import DensePoseChartPredictorOutput, DensePoseTransformData
|
|
|
|
|
|
def densepose_chart_predictor_output_hflip(
|
|
densepose_predictor_output: DensePoseChartPredictorOutput,
|
|
transform_data: DensePoseTransformData,
|
|
) -> DensePoseChartPredictorOutput:
|
|
"""
|
|
Change to take into account a Horizontal flip.
|
|
"""
|
|
if len(densepose_predictor_output) > 0:
|
|
|
|
PredictorOutput = type(densepose_predictor_output)
|
|
output_dict = {}
|
|
|
|
for field in fields(densepose_predictor_output):
|
|
field_value = getattr(densepose_predictor_output, field.name)
|
|
|
|
if isinstance(field_value, torch.Tensor):
|
|
setattr(densepose_predictor_output, field.name, torch.flip(field_value, [3]))
|
|
|
|
densepose_predictor_output = _flip_iuv_semantics_tensor(
|
|
densepose_predictor_output, transform_data
|
|
)
|
|
densepose_predictor_output = _flip_segm_semantics_tensor(
|
|
densepose_predictor_output, transform_data
|
|
)
|
|
|
|
for field in fields(densepose_predictor_output):
|
|
output_dict[field.name] = getattr(densepose_predictor_output, field.name)
|
|
|
|
return PredictorOutput(**output_dict)
|
|
else:
|
|
return densepose_predictor_output
|
|
|
|
|
|
def _flip_iuv_semantics_tensor(
|
|
densepose_predictor_output: DensePoseChartPredictorOutput,
|
|
dp_transform_data: DensePoseTransformData,
|
|
) -> DensePoseChartPredictorOutput:
|
|
point_label_symmetries = dp_transform_data.point_label_symmetries
|
|
uv_symmetries = dp_transform_data.uv_symmetries
|
|
|
|
N, C, H, W = densepose_predictor_output.u.shape
|
|
u_loc = (densepose_predictor_output.u[:, 1:, :, :].clamp(0, 1) * 255).long()
|
|
v_loc = (densepose_predictor_output.v[:, 1:, :, :].clamp(0, 1) * 255).long()
|
|
Iindex = torch.arange(C - 1, device=densepose_predictor_output.u.device)[
|
|
None, :, None, None
|
|
].expand(N, C - 1, H, W)
|
|
densepose_predictor_output.u[:, 1:, :, :] = uv_symmetries["U_transforms"][Iindex, v_loc, u_loc]
|
|
densepose_predictor_output.v[:, 1:, :, :] = uv_symmetries["V_transforms"][Iindex, v_loc, u_loc]
|
|
|
|
for el in ["fine_segm", "u", "v"]:
|
|
densepose_predictor_output.__dict__[el] = densepose_predictor_output.__dict__[el][
|
|
:, point_label_symmetries, :, :
|
|
]
|
|
return densepose_predictor_output
|
|
|
|
|
|
def _flip_segm_semantics_tensor(
|
|
densepose_predictor_output: DensePoseChartPredictorOutput, dp_transform_data
|
|
):
|
|
if densepose_predictor_output.coarse_segm.shape[1] > 2:
|
|
densepose_predictor_output.coarse_segm = densepose_predictor_output.coarse_segm[
|
|
:, dp_transform_data.mask_label_symmetries, :, :
|
|
]
|
|
return densepose_predictor_output
|
|
|
