|
|
| 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
|
|
|
