|
|
|
import numpy as np |
|
from typing import Iterable, Optional, Tuple |
|
import cv2 |
|
|
|
from densepose.structures import DensePoseDataRelative |
|
|
|
from .base import Boxes, Image, MatrixVisualizer, PointsVisualizer |
|
|
|
|
|
class DensePoseDataCoarseSegmentationVisualizer: |
|
""" |
|
Visualizer for ground truth segmentation |
|
""" |
|
|
|
def __init__(self, inplace=True, cmap=cv2.COLORMAP_PARULA, alpha=0.7, **kwargs): |
|
self.mask_visualizer = MatrixVisualizer( |
|
inplace=inplace, |
|
cmap=cmap, |
|
val_scale=255.0 / DensePoseDataRelative.N_BODY_PARTS, |
|
alpha=alpha, |
|
) |
|
|
|
def visualize( |
|
self, |
|
image_bgr: Image, |
|
bbox_densepose_datas: Optional[Tuple[Iterable[Boxes], Iterable[DensePoseDataRelative]]], |
|
) -> Image: |
|
if bbox_densepose_datas is None: |
|
return image_bgr |
|
for bbox_xywh, densepose_data in zip(*bbox_densepose_datas): |
|
matrix = densepose_data.segm.numpy() |
|
mask = np.zeros(matrix.shape, dtype=np.uint8) |
|
mask[matrix > 0] = 1 |
|
image_bgr = self.mask_visualizer.visualize(image_bgr, mask, matrix, bbox_xywh.numpy()) |
|
return image_bgr |
|
|
|
|
|
class DensePoseDataPointsVisualizer: |
|
def __init__(self, densepose_data_to_value_fn=None, cmap=cv2.COLORMAP_PARULA, **kwargs): |
|
self.points_visualizer = PointsVisualizer() |
|
self.densepose_data_to_value_fn = densepose_data_to_value_fn |
|
self.cmap = cmap |
|
|
|
def visualize( |
|
self, |
|
image_bgr: Image, |
|
bbox_densepose_datas: Optional[Tuple[Iterable[Boxes], Iterable[DensePoseDataRelative]]], |
|
) -> Image: |
|
if bbox_densepose_datas is None: |
|
return image_bgr |
|
for bbox_xywh, densepose_data in zip(*bbox_densepose_datas): |
|
x0, y0, w, h = bbox_xywh.numpy() |
|
x = densepose_data.x.numpy() * w / 255.0 + x0 |
|
y = densepose_data.y.numpy() * h / 255.0 + y0 |
|
pts_xy = zip(x, y) |
|
if self.densepose_data_to_value_fn is None: |
|
image_bgr = self.points_visualizer.visualize(image_bgr, pts_xy) |
|
else: |
|
v = self.densepose_data_to_value_fn(densepose_data) |
|
img_colors_bgr = cv2.applyColorMap(v, self.cmap) |
|
colors_bgr = [ |
|
[int(v) for v in img_color_bgr.ravel()] for img_color_bgr in img_colors_bgr |
|
] |
|
image_bgr = self.points_visualizer.visualize(image_bgr, pts_xy, colors_bgr) |
|
return image_bgr |
|
|
|
|
|
def _densepose_data_u_for_cmap(densepose_data): |
|
u = np.clip(densepose_data.u.numpy(), 0, 1) * 255.0 |
|
return u.astype(np.uint8) |
|
|
|
|
|
def _densepose_data_v_for_cmap(densepose_data): |
|
v = np.clip(densepose_data.v.numpy(), 0, 1) * 255.0 |
|
return v.astype(np.uint8) |
|
|
|
|
|
def _densepose_data_i_for_cmap(densepose_data): |
|
i = ( |
|
np.clip(densepose_data.i.numpy(), 0.0, DensePoseDataRelative.N_PART_LABELS) |
|
* 255.0 |
|
/ DensePoseDataRelative.N_PART_LABELS |
|
) |
|
return i.astype(np.uint8) |
|
|
|
|
|
class DensePoseDataPointsUVisualizer(DensePoseDataPointsVisualizer): |
|
def __init__(self, **kwargs): |
|
super(DensePoseDataPointsUVisualizer, self).__init__( |
|
densepose_data_to_value_fn=_densepose_data_u_for_cmap, **kwargs |
|
) |
|
|
|
|
|
class DensePoseDataPointsVVisualizer(DensePoseDataPointsVisualizer): |
|
def __init__(self, **kwargs): |
|
super(DensePoseDataPointsVVisualizer, self).__init__( |
|
densepose_data_to_value_fn=_densepose_data_v_for_cmap, **kwargs |
|
) |
|
|
|
|
|
class DensePoseDataPointsIVisualizer(DensePoseDataPointsVisualizer): |
|
def __init__(self, **kwargs): |
|
super(DensePoseDataPointsIVisualizer, self).__init__( |
|
densepose_data_to_value_fn=_densepose_data_i_for_cmap, **kwargs |
|
) |
|
|
