File size: 3,950 Bytes
73c83cf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
# Copyright (c) Facebook, Inc. and its affiliates.
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 # noqa
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
)
# pyre-fixme[61]: `vis` may not be initialized here.
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)
|