|
|
|
|
| from dataclasses import dataclass
|
| from typing import Any, Optional, Tuple
|
| import torch
|
|
|
|
|
| @dataclass
|
| class DensePoseChartResult:
|
| """
|
| DensePose results for chart-based methods represented by labels and inner
|
| coordinates (U, V) of individual charts. Each chart is a 2D manifold
|
| that has an associated label and is parameterized by two coordinates U and V.
|
| Both U and V take values in [0, 1].
|
| Thus the results are represented by two tensors:
|
| - labels (tensor [H, W] of long): contains estimated label for each pixel of
|
| the detection bounding box of size (H, W)
|
| - uv (tensor [2, H, W] of float): contains estimated U and V coordinates
|
| for each pixel of the detection bounding box of size (H, W)
|
| """
|
|
|
| labels: torch.Tensor
|
| uv: torch.Tensor
|
|
|
| def to(self, device: torch.device):
|
| """
|
| Transfers all tensors to the given device
|
| """
|
| labels = self.labels.to(device)
|
| uv = self.uv.to(device)
|
| return DensePoseChartResult(labels=labels, uv=uv)
|
|
|
|
|
| @dataclass
|
| class DensePoseChartResultWithConfidences:
|
| """
|
| We add confidence values to DensePoseChartResult
|
| Thus the results are represented by two tensors:
|
| - labels (tensor [H, W] of long): contains estimated label for each pixel of
|
| the detection bounding box of size (H, W)
|
| - uv (tensor [2, H, W] of float): contains estimated U and V coordinates
|
| for each pixel of the detection bounding box of size (H, W)
|
| Plus one [H, W] tensor of float for each confidence type
|
| """
|
|
|
| labels: torch.Tensor
|
| uv: torch.Tensor
|
| sigma_1: Optional[torch.Tensor] = None
|
| sigma_2: Optional[torch.Tensor] = None
|
| kappa_u: Optional[torch.Tensor] = None
|
| kappa_v: Optional[torch.Tensor] = None
|
| fine_segm_confidence: Optional[torch.Tensor] = None
|
| coarse_segm_confidence: Optional[torch.Tensor] = None
|
|
|
| def to(self, device: torch.device):
|
| """
|
| Transfers all tensors to the given device, except if their value is None
|
| """
|
|
|
| def to_device_if_tensor(var: Any):
|
| if isinstance(var, torch.Tensor):
|
| return var.to(device)
|
| return var
|
|
|
| return DensePoseChartResultWithConfidences(
|
| labels=self.labels.to(device),
|
| uv=self.uv.to(device),
|
| sigma_1=to_device_if_tensor(self.sigma_1),
|
| sigma_2=to_device_if_tensor(self.sigma_2),
|
| kappa_u=to_device_if_tensor(self.kappa_u),
|
| kappa_v=to_device_if_tensor(self.kappa_v),
|
| fine_segm_confidence=to_device_if_tensor(self.fine_segm_confidence),
|
| coarse_segm_confidence=to_device_if_tensor(self.coarse_segm_confidence),
|
| )
|
|
|
|
|
| @dataclass
|
| class DensePoseChartResultQuantized:
|
| """
|
| DensePose results for chart-based methods represented by labels and quantized
|
| inner coordinates (U, V) of individual charts. Each chart is a 2D manifold
|
| that has an associated label and is parameterized by two coordinates U and V.
|
| Both U and V take values in [0, 1].
|
| Quantized coordinates Uq and Vq have uint8 values which are obtained as:
|
| Uq = U * 255 (hence 0 <= Uq <= 255)
|
| Vq = V * 255 (hence 0 <= Vq <= 255)
|
| Thus the results are represented by one tensor:
|
| - labels_uv_uint8 (tensor [3, H, W] of uint8): contains estimated label
|
| and quantized coordinates Uq and Vq for each pixel of the detection
|
| bounding box of size (H, W)
|
| """
|
|
|
| labels_uv_uint8: torch.Tensor
|
|
|
| def to(self, device: torch.device):
|
| """
|
| Transfers all tensors to the given device
|
| """
|
| labels_uv_uint8 = self.labels_uv_uint8.to(device)
|
| return DensePoseChartResultQuantized(labels_uv_uint8=labels_uv_uint8)
|
|
|
|
|
| @dataclass
|
| class DensePoseChartResultCompressed:
|
| """
|
| DensePose results for chart-based methods represented by a PNG-encoded string.
|
| The tensor of quantized DensePose results of size [3, H, W] is considered
|
| as an image with 3 color channels. PNG compression is applied and the result
|
| is stored as a Base64-encoded string. The following attributes are defined:
|
| - shape_chw (tuple of 3 int): contains shape of the result tensor
|
| (number of channels, height, width)
|
| - labels_uv_str (str): contains Base64-encoded results tensor of size
|
| [3, H, W] compressed with PNG compression methods
|
| """
|
|
|
| shape_chw: Tuple[int, int, int]
|
| labels_uv_str: str
|
|
|
|
|
| def quantize_densepose_chart_result(result: DensePoseChartResult) -> DensePoseChartResultQuantized:
|
| """
|
| Applies quantization to DensePose chart-based result.
|
|
|
| Args:
|
| result (DensePoseChartResult): DensePose chart-based result
|
| Return:
|
| Quantized DensePose chart-based result (DensePoseChartResultQuantized)
|
| """
|
| h, w = result.labels.shape
|
| labels_uv_uint8 = torch.zeros([3, h, w], dtype=torch.uint8, device=result.labels.device)
|
| labels_uv_uint8[0] = result.labels
|
| labels_uv_uint8[1:] = (result.uv * 255).clamp(0, 255).byte()
|
| return DensePoseChartResultQuantized(labels_uv_uint8=labels_uv_uint8)
|
|
|
|
|
| def compress_quantized_densepose_chart_result(
|
| result: DensePoseChartResultQuantized,
|
| ) -> DensePoseChartResultCompressed:
|
| """
|
| Compresses quantized DensePose chart-based result
|
|
|
| Args:
|
| result (DensePoseChartResultQuantized): quantized DensePose chart-based result
|
| Return:
|
| Compressed DensePose chart-based result (DensePoseChartResultCompressed)
|
| """
|
| import base64
|
| import numpy as np
|
| from io import BytesIO
|
| from PIL import Image
|
|
|
| labels_uv_uint8_np_chw = result.labels_uv_uint8.cpu().numpy()
|
| labels_uv_uint8_np_hwc = np.moveaxis(labels_uv_uint8_np_chw, 0, -1)
|
| im = Image.fromarray(labels_uv_uint8_np_hwc)
|
| fstream = BytesIO()
|
| im.save(fstream, format="png", optimize=True)
|
| labels_uv_str = base64.encodebytes(fstream.getvalue()).decode()
|
| shape_chw = labels_uv_uint8_np_chw.shape
|
| return DensePoseChartResultCompressed(labels_uv_str=labels_uv_str, shape_chw=shape_chw)
|
|
|
|
|
| def decompress_compressed_densepose_chart_result(
|
| result: DensePoseChartResultCompressed,
|
| ) -> DensePoseChartResultQuantized:
|
| """
|
| Decompresses DensePose chart-based result encoded into a base64 string
|
|
|
| Args:
|
| result (DensePoseChartResultCompressed): compressed DensePose chart result
|
| Return:
|
| Quantized DensePose chart-based result (DensePoseChartResultQuantized)
|
| """
|
| import base64
|
| import numpy as np
|
| from io import BytesIO
|
| from PIL import Image
|
|
|
| fstream = BytesIO(base64.decodebytes(result.labels_uv_str.encode()))
|
| im = Image.open(fstream)
|
| labels_uv_uint8_np_chw = np.moveaxis(np.array(im, dtype=np.uint8), -1, 0)
|
| return DensePoseChartResultQuantized(
|
| labels_uv_uint8=torch.from_numpy(labels_uv_uint8_np_chw.reshape(result.shape_chw))
|
| )
|
|
|
