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