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