Spaces:
Running
on
Zero
Running
on
Zero
| # Copyright (c) Facebook, Inc. and its affiliates. | |
| import numpy as np | |
| from typing import List, Optional, Tuple | |
| import torch | |
| from detectron2.data.detection_utils import read_image | |
| from ..structures import DensePoseChartResult | |
| from .base import Boxes, Image | |
| from .densepose_results import DensePoseResultsVisualizer | |
| def get_texture_atlas(path: Optional[str]) -> Optional[np.ndarray]: | |
| if path is None: | |
| return None | |
| # Reading images like that downsamples 16-bit images to 8-bit | |
| # If 16-bit images are needed, we can replace that by cv2.imread with the | |
| # cv2.IMREAD_UNCHANGED flag (with cv2 we also need it to keep alpha channels) | |
| # The rest of the pipeline would need to be adapted to 16-bit images too | |
| bgr_image = read_image(path) | |
| rgb_image = np.copy(bgr_image) # Convert BGR -> RGB | |
| rgb_image[:, :, :3] = rgb_image[:, :, 2::-1] # Works with alpha channel | |
| return rgb_image | |
| class DensePoseResultsVisualizerWithTexture(DensePoseResultsVisualizer): | |
| """ | |
| texture_atlas: An image, size 6N * 4N, with N * N squares for each of the 24 body parts. | |
| It must follow the grid found at https://github.com/facebookresearch/DensePose/blob/master/DensePoseData/demo_data/texture_atlas_200.png # noqa | |
| For each body part, U is proportional to the x coordinate, and (1 - V) to y | |
| """ | |
| def __init__(self, texture_atlas, **kwargs): | |
| self.texture_atlas = texture_atlas | |
| self.body_part_size = texture_atlas.shape[0] // 6 | |
| assert self.body_part_size == texture_atlas.shape[1] // 4 | |
| def visualize( | |
| self, | |
| image_bgr: Image, | |
| results_and_boxes_xywh: Tuple[Optional[List[DensePoseChartResult]], Optional[Boxes]], | |
| ) -> Image: | |
| densepose_result, boxes_xywh = results_and_boxes_xywh | |
| if densepose_result is None or boxes_xywh is None: | |
| return image_bgr | |
| boxes_xywh = boxes_xywh.int().cpu().numpy() | |
| texture_image, alpha = self.get_texture() | |
| for i, result in enumerate(densepose_result): | |
| iuv_array = torch.cat((result.labels[None], result.uv.clamp(0, 1))) | |
| x, y, w, h = boxes_xywh[i] | |
| bbox_image = image_bgr[y : y + h, x : x + w] | |
| image_bgr[y : y + h, x : x + w] = self.generate_image_with_texture( | |
| texture_image, alpha, bbox_image, iuv_array.cpu().numpy() | |
| ) | |
| return image_bgr | |
| def get_texture(self): | |
| N = self.body_part_size | |
| texture_image = np.zeros([24, N, N, self.texture_atlas.shape[-1]]) | |
| for i in range(4): | |
| for j in range(6): | |
| texture_image[(6 * i + j), :, :, :] = self.texture_atlas[ | |
| N * j : N * (j + 1), N * i : N * (i + 1), : | |
| ] | |
| if texture_image.shape[-1] == 4: # Image with alpha channel | |
| alpha = texture_image[:, :, :, -1] / 255.0 | |
| texture_image = texture_image[:, :, :, :3] | |
| else: | |
| alpha = texture_image.sum(axis=-1) > 0 | |
| return texture_image, alpha | |
| def generate_image_with_texture(self, texture_image, alpha, bbox_image_bgr, iuv_array): | |
| I, U, V = iuv_array | |
| generated_image_bgr = bbox_image_bgr.copy() | |
| for PartInd in range(1, 25): | |
| x, y = np.where(I == PartInd) | |
| x_index = (U[x, y] * (self.body_part_size - 1)).astype(int) | |
| y_index = ((1 - V[x, y]) * (self.body_part_size - 1)).astype(int) | |
| part_alpha = np.expand_dims(alpha[PartInd - 1, y_index, x_index], -1) | |
| generated_image_bgr[I == PartInd] = ( | |
| generated_image_bgr[I == PartInd] * (1 - part_alpha) | |
| + texture_image[PartInd - 1, y_index, x_index] * part_alpha | |
| ) | |
| return generated_image_bgr.astype(np.uint8) | |