Spaces:
Paused
Paused
| # Copyright (c) Facebook, Inc. and its affiliates. | |
| 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 | |
| ) | |