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