| | |
| | import numpy as np |
| | from typing import Optional, Tuple |
| | import cv2 |
| |
|
| | from densepose.structures import DensePoseDataRelative |
| |
|
| | from ..structures import DensePoseChartPredictorOutput |
| | from .base import Boxes, Image, MatrixVisualizer |
| |
|
| |
|
| | class DensePoseOutputsVisualizer: |
| | def __init__( |
| | self, inplace=True, cmap=cv2.COLORMAP_PARULA, alpha=0.7, to_visualize=None, **kwargs |
| | ): |
| | assert to_visualize in "IUV", "can only visualize IUV" |
| | self.to_visualize = to_visualize |
| |
|
| | if self.to_visualize == "I": |
| | val_scale = 255.0 / DensePoseDataRelative.N_PART_LABELS |
| | else: |
| | val_scale = 1.0 |
| | self.mask_visualizer = MatrixVisualizer( |
| | inplace=inplace, cmap=cmap, val_scale=val_scale, alpha=alpha |
| | ) |
| |
|
| | def visualize( |
| | self, |
| | image_bgr: Image, |
| | dp_output_with_bboxes: Tuple[Optional[DensePoseChartPredictorOutput], Optional[Boxes]], |
| | ) -> Image: |
| | densepose_output, bboxes_xywh = dp_output_with_bboxes |
| | if densepose_output is None or bboxes_xywh is None: |
| | return image_bgr |
| |
|
| | assert isinstance( |
| | densepose_output, DensePoseChartPredictorOutput |
| | ), "DensePoseChartPredictorOutput expected, {} encountered".format(type(densepose_output)) |
| |
|
| | S = densepose_output.coarse_segm |
| | I = densepose_output.fine_segm |
| | U = densepose_output.u |
| | V = densepose_output.v |
| | N = S.size(0) |
| | assert N == I.size( |
| | 0 |
| | ), "densepose outputs S {} and I {}" " should have equal first dim size".format( |
| | S.size(), I.size() |
| | ) |
| | assert N == U.size( |
| | 0 |
| | ), "densepose outputs S {} and U {}" " should have equal first dim size".format( |
| | S.size(), U.size() |
| | ) |
| | assert N == V.size( |
| | 0 |
| | ), "densepose outputs S {} and V {}" " should have equal first dim size".format( |
| | S.size(), V.size() |
| | ) |
| | assert N == len( |
| | bboxes_xywh |
| | ), "number of bounding boxes {}" " should be equal to first dim size of outputs {}".format( |
| | len(bboxes_xywh), N |
| | ) |
| | for n in range(N): |
| | Sn = S[n].argmax(dim=0) |
| | In = I[n].argmax(dim=0) * (Sn > 0).long() |
| | segmentation = In.cpu().numpy().astype(np.uint8) |
| | mask = np.zeros(segmentation.shape, dtype=np.uint8) |
| | mask[segmentation > 0] = 1 |
| | bbox_xywh = bboxes_xywh[n] |
| |
|
| | if self.to_visualize == "I": |
| | vis = segmentation |
| | elif self.to_visualize in "UV": |
| | U_or_Vn = {"U": U, "V": V}[self.to_visualize][n].cpu().numpy().astype(np.float32) |
| | vis = np.zeros(segmentation.shape, dtype=np.float32) |
| | for partId in range(U_or_Vn.shape[0]): |
| | vis[segmentation == partId] = ( |
| | U_or_Vn[partId][segmentation == partId].clip(0, 1) * 255 |
| | ) |
| |
|
| | |
| | image_bgr = self.mask_visualizer.visualize(image_bgr, mask, vis, bbox_xywh) |
| |
|
| | return image_bgr |
| |
|
| |
|
| | class DensePoseOutputsUVisualizer(DensePoseOutputsVisualizer): |
| | def __init__(self, inplace=True, cmap=cv2.COLORMAP_PARULA, alpha=0.7, **kwargs): |
| | super().__init__(inplace=inplace, cmap=cmap, alpha=alpha, to_visualize="U", **kwargs) |
| |
|
| |
|
| | class DensePoseOutputsVVisualizer(DensePoseOutputsVisualizer): |
| | def __init__(self, inplace=True, cmap=cv2.COLORMAP_PARULA, alpha=0.7, **kwargs): |
| | super().__init__(inplace=inplace, cmap=cmap, alpha=alpha, to_visualize="V", **kwargs) |
| |
|
| |
|
| | class DensePoseOutputsFineSegmentationVisualizer(DensePoseOutputsVisualizer): |
| | def __init__(self, inplace=True, cmap=cv2.COLORMAP_PARULA, alpha=0.7, **kwargs): |
| | super().__init__(inplace=inplace, cmap=cmap, alpha=alpha, to_visualize="I", **kwargs) |
| |
|
