structures.py

# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
import base64
import numpy as np
from io import BytesIO
import torch
from PIL import Image
from torch.nn import functional as F


class DensePoseTransformData(object):

    # Horizontal symmetry label transforms used for horizontal flip
    MASK_LABEL_SYMMETRIES = [0, 1, 3, 2, 5, 4, 7, 6, 9, 8, 11, 10, 13, 12, 14]
    # fmt: off
    POINT_LABEL_SYMMETRIES = [ 0, 1, 2, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15, 18, 17, 20, 19, 22, 21, 24, 23]  # noqa
    # fmt: on

    def __init__(self, uv_symmetries):
        self.mask_label_symmetries = DensePoseTransformData.MASK_LABEL_SYMMETRIES
        self.point_label_symmetries = DensePoseTransformData.POINT_LABEL_SYMMETRIES
        self.uv_symmetries = uv_symmetries

    @staticmethod
    def load(fpath):
        import scipy.io

        uv_symmetry_map = scipy.io.loadmat(fpath)
        uv_symmetry_map_torch = {}
        for key in ["U_transforms", "V_transforms"]:
            uv_symmetry_map_torch[key] = []
            map_src = uv_symmetry_map[key]
            map_dst = uv_symmetry_map_torch[key]
            for i in range(map_src.shape[1]):
                map_dst.append(torch.from_numpy(map_src[0, i]).to(dtype=torch.float))
            uv_symmetry_map_torch[key] = torch.stack(map_dst, dim=0).to(
                device=torch.cuda.current_device()
            )
        transform_data = DensePoseTransformData(uv_symmetry_map_torch)
        return transform_data


class DensePoseDataRelative(object):
    """
    Dense pose relative annotations that can be applied to any bounding box:
        x - normalized X coordinates [0, 255] of annotated points
        y - normalized Y coordinates [0, 255] of annotated points
        i - body part labels 0,...,24 for annotated points
        u - body part U coordinates [0, 1] for annotated points
        v - body part V coordinates [0, 1] for annotated points
        segm - 256x256 segmentation mask with values 0,...,14
    To obtain absolute x and y data wrt some bounding box one needs to first
    divide the data by 256, multiply by the respective bounding box size
    and add bounding box offset:
        x_img = x0 + x_norm * w / 256.0
        y_img = y0 + y_norm * h / 256.0
    Segmentation masks are typically sampled to get image-based masks.
    """

    # Key for normalized X coordinates in annotation dict
    X_KEY = "dp_x"
    # Key for normalized Y coordinates in annotation dict
    Y_KEY = "dp_y"
    # Key for U part coordinates in annotation dict
    U_KEY = "dp_U"
    # Key for V part coordinates in annotation dict
    V_KEY = "dp_V"
    # Key for I point labels in annotation dict
    I_KEY = "dp_I"
    # Key for segmentation mask in annotation dict
    S_KEY = "dp_masks"
    # Number of body parts in segmentation masks
    N_BODY_PARTS = 14
    # Number of parts in point labels
    N_PART_LABELS = 24
    MASK_SIZE = 256

    def __init__(self, annotation, cleanup=False):
        is_valid, reason_not_valid = DensePoseDataRelative.validate_annotation(annotation)
        assert is_valid, "Invalid DensePose annotations: {}".format(reason_not_valid)
        self.x = torch.as_tensor(annotation[DensePoseDataRelative.X_KEY])
        self.y = torch.as_tensor(annotation[DensePoseDataRelative.Y_KEY])
        self.i = torch.as_tensor(annotation[DensePoseDataRelative.I_KEY])
        self.u = torch.as_tensor(annotation[DensePoseDataRelative.U_KEY])
        self.v = torch.as_tensor(annotation[DensePoseDataRelative.V_KEY])
        self.segm = DensePoseDataRelative.extract_segmentation_mask(annotation)
        self.device = torch.device("cpu")
        if cleanup:
            DensePoseDataRelative.cleanup_annotation(annotation)

    def to(self, device):
        if self.device == device:
            return self
        new_data = DensePoseDataRelative.__new__(DensePoseDataRelative)
        new_data.x = self.x
        new_data.x = self.x.to(device)
        new_data.y = self.y.to(device)
        new_data.i = self.i.to(device)
        new_data.u = self.u.to(device)
        new_data.v = self.v.to(device)
        new_data.segm = self.segm.to(device)
        new_data.device = device
        return new_data

    @staticmethod
    def extract_segmentation_mask(annotation):
        import pycocotools.mask as mask_utils

        poly_specs = annotation[DensePoseDataRelative.S_KEY]
        segm = torch.zeros((DensePoseDataRelative.MASK_SIZE,) * 2, dtype=torch.float32)
        for i in range(DensePoseDataRelative.N_BODY_PARTS):
            poly_i = poly_specs[i]
            if poly_i:
                mask_i = mask_utils.decode(poly_i)
                segm[mask_i > 0] = i + 1
        return segm

    @staticmethod
    def validate_annotation(annotation):
        for key in [
            DensePoseDataRelative.X_KEY,
            DensePoseDataRelative.Y_KEY,
            DensePoseDataRelative.I_KEY,
            DensePoseDataRelative.U_KEY,
            DensePoseDataRelative.V_KEY,
            DensePoseDataRelative.S_KEY,
        ]:
            if key not in annotation:
                return False, "no {key} data in the annotation".format(key=key)
        return True, None

    @staticmethod
    def cleanup_annotation(annotation):
        for key in [
            DensePoseDataRelative.X_KEY,
            DensePoseDataRelative.Y_KEY,
            DensePoseDataRelative.I_KEY,
            DensePoseDataRelative.U_KEY,
            DensePoseDataRelative.V_KEY,
            DensePoseDataRelative.S_KEY,
        ]:
            if key in annotation:
                del annotation[key]

    def apply_transform(self, transforms, densepose_transform_data):
        self._transform_pts(transforms, densepose_transform_data)
        self._transform_segm(transforms, densepose_transform_data)

    def _transform_pts(self, transforms, dp_transform_data):
        import detectron2.data.transforms as T

        # NOTE: This assumes that HorizFlipTransform is the only one that does flip
        do_hflip = sum(isinstance(t, T.HFlipTransform) for t in transforms.transforms) % 2 == 1
        if do_hflip:
            self.x = self.segm.size(1) - self.x
            self._flip_iuv_semantics(dp_transform_data)

    def _flip_iuv_semantics(self, dp_transform_data: DensePoseTransformData) -> None:
        i_old = self.i.clone()
        uv_symmetries = dp_transform_data.uv_symmetries
        pt_label_symmetries = dp_transform_data.point_label_symmetries
        for i in range(self.N_PART_LABELS):
            if i + 1 in i_old:
                annot_indices_i = i_old == i + 1
                if pt_label_symmetries[i + 1] != i + 1:
                    self.i[annot_indices_i] = pt_label_symmetries[i + 1]
                u_loc = (self.u[annot_indices_i] * 255).long()
                v_loc = (self.v[annot_indices_i] * 255).long()
                self.u[annot_indices_i] = uv_symmetries["U_transforms"][i][v_loc, u_loc].to(
                    device=self.u.device
                )
                self.v[annot_indices_i] = uv_symmetries["V_transforms"][i][v_loc, u_loc].to(
                    device=self.v.device
                )

    def _transform_segm(self, transforms, dp_transform_data):
        import detectron2.data.transforms as T

        # NOTE: This assumes that HorizFlipTransform is the only one that does flip
        do_hflip = sum(isinstance(t, T.HFlipTransform) for t in transforms.transforms) % 2 == 1
        if do_hflip:
            self.segm = torch.flip(self.segm, [1])
            self._flip_segm_semantics(dp_transform_data)

    def _flip_segm_semantics(self, dp_transform_data):
        old_segm = self.segm.clone()
        mask_label_symmetries = dp_transform_data.mask_label_symmetries
        for i in range(self.N_BODY_PARTS):
            if mask_label_symmetries[i + 1] != i + 1:
                self.segm[old_segm == i + 1] = mask_label_symmetries[i + 1]


def normalized_coords_transform(x0, y0, w, h):
    """
    Coordinates transform that maps top left corner to (-1, -1) and bottom
    right corner to (1, 1). Used for torch.grid_sample to initialize the
    grid
    """

    def f(p):
        return (2 * (p[0] - x0) / w - 1, 2 * (p[1] - y0) / h - 1)

    return f


class DensePoseOutput(object):
    def __init__(self, S, I, U, V, confidences):
        """
        Args:
            S (`torch.Tensor`): coarse segmentation tensor of size (N, A, H, W)
            I (`torch.Tensor`): fine segmentation tensor of size (N, C, H, W)
            U (`torch.Tensor`): U coordinates for each fine segmentation label of size (N, C, H, W)
            V (`torch.Tensor`): V coordinates for each fine segmentation label of size (N, C, H, W)
            confidences (dict of str -> `torch.Tensor`) estimated confidence model parameters
        """
        self.S = S
        self.I = I  # noqa: E741
        self.U = U
        self.V = V
        self.confidences = confidences
        self._check_output_dims(S, I, U, V)

    def _check_output_dims(self, S, I, U, V):
        assert (
            len(S.size()) == 4
        ), "Segmentation output should have 4 " "dimensions (NCHW), but has size {}".format(
            S.size()
        )
        assert (
            len(I.size()) == 4
        ), "Segmentation output should have 4 " "dimensions (NCHW), but has size {}".format(
            S.size()
        )
        assert (
            len(U.size()) == 4
        ), "Segmentation output should have 4 " "dimensions (NCHW), but has size {}".format(
            S.size()
        )
        assert (
            len(V.size()) == 4
        ), "Segmentation output should have 4 " "dimensions (NCHW), but has size {}".format(
            S.size()
        )
        assert len(S) == len(I), (
            "Number of output segmentation planes {} "
            "should be equal to the number of output part index "
            "planes {}".format(len(S), len(I))
        )
        assert S.size()[2:] == I.size()[2:], (
            "Output segmentation plane size {} "
            "should be equal to the output part index "
            "plane size {}".format(S.size()[2:], I.size()[2:])
        )
        assert I.size() == U.size(), (
            "Part index output shape {} "
            "should be the same as U coordinates output shape {}".format(I.size(), U.size())
        )
        assert I.size() == V.size(), (
            "Part index output shape {} "
            "should be the same as V coordinates output shape {}".format(I.size(), V.size())
        )

    def resize(self, image_size_hw):
        # do nothing - outputs are invariant to resize
        pass

    def _crop(self, S, I, U, V, bbox_old_xywh, bbox_new_xywh):
        """
        Resample S, I, U, V from bbox_old to the cropped bbox_new
        """
        x0old, y0old, wold, hold = bbox_old_xywh
        x0new, y0new, wnew, hnew = bbox_new_xywh
        tr_coords = normalized_coords_transform(x0old, y0old, wold, hold)
        topleft = (x0new, y0new)
        bottomright = (x0new + wnew, y0new + hnew)
        topleft_norm = tr_coords(topleft)
        bottomright_norm = tr_coords(bottomright)
        hsize = S.size(1)
        wsize = S.size(2)
        grid = torch.meshgrid(
            torch.arange(
                topleft_norm[1],
                bottomright_norm[1],
                (bottomright_norm[1] - topleft_norm[1]) / hsize,
            )[:hsize],
            torch.arange(
                topleft_norm[0],
                bottomright_norm[0],
                (bottomright_norm[0] - topleft_norm[0]) / wsize,
            )[:wsize],
        )
        grid = torch.stack(grid, dim=2).to(S.device)
        assert (
            grid.size(0) == hsize
        ), "Resampled grid expected " "height={}, actual height={}".format(hsize, grid.size(0))
        assert grid.size(1) == wsize, "Resampled grid expected " "width={}, actual width={}".format(
            wsize, grid.size(1)
        )
        S_new = F.grid_sample(
            S.unsqueeze(0),
            torch.unsqueeze(grid, 0),
            mode="bilinear",
            padding_mode="border",
            align_corners=True,
        ).squeeze(0)
        I_new = F.grid_sample(
            I.unsqueeze(0),
            torch.unsqueeze(grid, 0),
            mode="bilinear",
            padding_mode="border",
            align_corners=True,
        ).squeeze(0)
        U_new = F.grid_sample(
            U.unsqueeze(0),
            torch.unsqueeze(grid, 0),
            mode="bilinear",
            padding_mode="border",
            align_corners=True,
        ).squeeze(0)
        V_new = F.grid_sample(
            V.unsqueeze(0),
            torch.unsqueeze(grid, 0),
            mode="bilinear",
            padding_mode="border",
            align_corners=True,
        ).squeeze(0)
        return S_new, I_new, U_new, V_new

    def crop(self, indices_cropped, bboxes_old, bboxes_new):
        """
        Crop outputs for selected bounding boxes to the new bounding boxes.
        """
        # VK: cropping is ignored for now
        # for i, ic in enumerate(indices_cropped):
        #    self.S[ic], self.I[ic], self.U[ic], self.V[ic] = \
        #        self._crop(self.S[ic], self.I[ic], self.U[ic], self.V[ic],
        #        bboxes_old[i], bboxes_new[i])
        pass

    def hflip(self, transform_data: DensePoseTransformData) -> None:
        """
        Change S, I, U and V to take into account a Horizontal flip.
        """
        if self.I.shape[0] > 0:
            for el in "SIUV":
                self.__dict__[el] = torch.flip(self.__dict__[el], [3])
            self._flip_iuv_semantics_tensor(transform_data)
            self._flip_segm_semantics_tensor(transform_data)

    def _flip_iuv_semantics_tensor(self, dp_transform_data: DensePoseTransformData) -> None:
        point_label_symmetries = dp_transform_data.point_label_symmetries
        uv_symmetries = dp_transform_data.uv_symmetries

        N, C, H, W = self.U.shape
        u_loc = (self.U[:, 1:, :, :].clamp(0, 1) * 255).long()
        v_loc = (self.V[:, 1:, :, :].clamp(0, 1) * 255).long()
        Iindex = torch.arange(C - 1, device=self.U.device)[None, :, None, None].expand(
            N, C - 1, H, W
        )
        self.U[:, 1:, :, :] = uv_symmetries["U_transforms"][Iindex, v_loc, u_loc].to(
            device=self.U.device
        )
        self.V[:, 1:, :, :] = uv_symmetries["V_transforms"][Iindex, v_loc, u_loc].to(
            device=self.V.device
        )

        for el in "IUV":
            self.__dict__[el] = self.__dict__[el][:, point_label_symmetries, :, :]

    def _flip_segm_semantics_tensor(self, dp_transform_data):
        if self.S.shape[1] == DensePoseDataRelative.N_BODY_PARTS + 1:
            self.S = self.S[:, dp_transform_data.mask_label_symmetries, :, :]

    def to_result(self, boxes_xywh):
        """
        Convert DensePose outputs to results format. Results are more compact,
        but cannot be resampled any more
        """
        result = DensePoseResult(boxes_xywh, self.S, self.I, self.U, self.V)
        return result

    def __getitem__(self, item):
        if isinstance(item, int):
            S_selected = self.S[item].unsqueeze(0)
            I_selected = self.I[item].unsqueeze(0)
            U_selected = self.U[item].unsqueeze(0)
            V_selected = self.V[item].unsqueeze(0)
            conf_selected = {}
            for key in self.confidences:
                conf_selected[key] = self.confidences[key][item].unsqueeze(0)
        else:
            S_selected = self.S[item]
            I_selected = self.I[item]
            U_selected = self.U[item]
            V_selected = self.V[item]
            conf_selected = {}
            for key in self.confidences:
                conf_selected[key] = self.confidences[key][item]
        return DensePoseOutput(S_selected, I_selected, U_selected, V_selected, conf_selected)

    def __str__(self):
        s = "DensePoseOutput S {}, I {}, U {}, V {}".format(
            list(self.S.size()), list(self.I.size()), list(self.U.size()), list(self.V.size())
        )
        s_conf = "confidences: [{}]".format(
            ", ".join([f"{key} {list(self.confidences[key].size())}" for key in self.confidences])
        )
        return ", ".join([s, s_conf])

    def __len__(self):
        return self.S.size(0)


class DensePoseResult(object):
    def __init__(self, boxes_xywh, S, I, U, V):
        self.results = []
        self.boxes_xywh = boxes_xywh.cpu().tolist()
        assert len(boxes_xywh.size()) == 2
        assert boxes_xywh.size(1) == 4
        for i, box_xywh in enumerate(boxes_xywh):
            result_i = self._output_to_result(box_xywh, S[[i]], I[[i]], U[[i]], V[[i]])
            result_numpy_i = result_i.cpu().numpy()
            result_encoded_i = DensePoseResult.encode_png_data(result_numpy_i)
            result_encoded_with_shape_i = (result_numpy_i.shape, result_encoded_i)
            self.results.append(result_encoded_with_shape_i)

    def __str__(self):
        s = "DensePoseResult: N={} [{}]".format(
            len(self.results), ", ".join([str(list(r[0])) for r in self.results])
        )
        return s

    def _output_to_result(self, box_xywh, S, I, U, V):
        x, y, w, h = box_xywh
        w = max(int(w), 1)
        h = max(int(h), 1)
        result = torch.zeros([3, h, w], dtype=torch.uint8, device=U.device)
        assert (
            len(S.size()) == 4
        ), "AnnIndex tensor size should have {} " "dimensions but has {}".format(4, len(S.size()))
        s_bbox = F.interpolate(S, (h, w), mode="bilinear", align_corners=False).argmax(dim=1)
        assert (
            len(I.size()) == 4
        ), "IndexUV tensor size should have {} " "dimensions but has {}".format(4, len(S.size()))
        i_bbox = (
            F.interpolate(I, (h, w), mode="bilinear", align_corners=False).argmax(dim=1)
            * (s_bbox > 0).long()
        ).squeeze(0)
        assert len(U.size()) == 4, "U tensor size should have {} " "dimensions but has {}".format(
            4, len(U.size())
        )
        u_bbox = F.interpolate(U, (h, w), mode="bilinear", align_corners=False)
        assert len(V.size()) == 4, "V tensor size should have {} " "dimensions but has {}".format(
            4, len(V.size())
        )
        v_bbox = F.interpolate(V, (h, w), mode="bilinear", align_corners=False)
        result[0] = i_bbox
        for part_id in range(1, u_bbox.size(1)):
            result[1][i_bbox == part_id] = (
                (u_bbox[0, part_id][i_bbox == part_id] * 255).clamp(0, 255).to(torch.uint8)
            )
            result[2][i_bbox == part_id] = (
                (v_bbox[0, part_id][i_bbox == part_id] * 255).clamp(0, 255).to(torch.uint8)
            )
        assert (
            result.size(1) == h
        ), "Results height {} should be equal" "to bounding box height {}".format(result.size(1), h)
        assert (
            result.size(2) == w
        ), "Results width {} should be equal" "to bounding box width {}".format(result.size(2), w)
        return result

    @staticmethod
    def encode_png_data(arr):
        """
        Encode array data as a PNG image using the highest compression rate
        @param arr [in] Data stored in an array of size (3, M, N) of type uint8
        @return Base64-encoded string containing PNG-compressed data
        """
        assert len(arr.shape) == 3, "Expected a 3D array as an input," " got a {0}D array".format(
            len(arr.shape)
        )
        assert arr.shape[0] == 3, "Expected first array dimension of size 3," " got {0}".format(
            arr.shape[0]
        )
        assert arr.dtype == np.uint8, "Expected an array of type np.uint8, " " got {0}".format(
            arr.dtype
        )
        data = np.moveaxis(arr, 0, -1)
        im = Image.fromarray(data)
        fstream = BytesIO()
        im.save(fstream, format="png", optimize=True)
        s = base64.encodebytes(fstream.getvalue()).decode()
        return s

    @staticmethod
    def decode_png_data(shape, s):
        """
        Decode array data from a string that contains PNG-compressed data
        @param Base64-encoded string containing PNG-compressed data
        @return Data stored in an array of size (3, M, N) of type uint8
        """
        fstream = BytesIO(base64.decodebytes(s.encode()))
        im = Image.open(fstream)
        data = np.moveaxis(np.array(im.getdata(), dtype=np.uint8), -1, 0)
        return data.reshape(shape)

    def __len__(self):
        return len(self.results)

    def __getitem__(self, item):
        result_encoded = self.results[item]
        bbox_xywh = self.boxes_xywh[item]
        return result_encoded, bbox_xywh


class DensePoseList(object):

    _TORCH_DEVICE_CPU = torch.device("cpu")

    def __init__(self, densepose_datas, boxes_xyxy_abs, image_size_hw, device=_TORCH_DEVICE_CPU):
        assert len(densepose_datas) == len(
            boxes_xyxy_abs
        ), "Attempt to initialize DensePoseList with {} DensePose datas " "and {} boxes".format(
            len(densepose_datas), len(boxes_xyxy_abs)
        )
        self.densepose_datas = []
        for densepose_data in densepose_datas:
            assert isinstance(densepose_data, DensePoseDataRelative) or densepose_data is None, (
                "Attempt to initialize DensePoseList with DensePose datas "
                "of type {}, expected DensePoseDataRelative".format(type(densepose_data))
            )
            densepose_data_ondevice = (
                densepose_data.to(device) if densepose_data is not None else None
            )
            self.densepose_datas.append(densepose_data_ondevice)
        self.boxes_xyxy_abs = boxes_xyxy_abs.to(device)
        self.image_size_hw = image_size_hw
        self.device = device

    def to(self, device):
        if self.device == device:
            return self
        return DensePoseList(self.densepose_datas, self.boxes_xyxy_abs, self.image_size_hw, device)

    def __iter__(self):
        return iter(self.densepose_datas)

    def __len__(self):
        return len(self.densepose_datas)

    def __repr__(self):
        s = self.__class__.__name__ + "("
        s += "num_instances={}, ".format(len(self.densepose_datas))
        s += "image_width={}, ".format(self.image_size_hw[1])
        s += "image_height={})".format(self.image_size_hw[0])
        return s

    def __getitem__(self, item):
        if isinstance(item, int):
            densepose_data_rel = self.densepose_datas[item]
            return densepose_data_rel
        elif isinstance(item, slice):
            densepose_datas_rel = self.densepose_datas[item]
            boxes_xyxy_abs = self.boxes_xyxy_abs[item]
            return DensePoseList(
                densepose_datas_rel, boxes_xyxy_abs, self.image_size_hw, self.device
            )
        elif isinstance(item, torch.Tensor) and (item.dtype == torch.bool):
            densepose_datas_rel = [self.densepose_datas[i] for i, x in enumerate(item) if x > 0]
            boxes_xyxy_abs = self.boxes_xyxy_abs[item]
            return DensePoseList(
                densepose_datas_rel, boxes_xyxy_abs, self.image_size_hw, self.device
            )
        else:
            densepose_datas_rel = [self.densepose_datas[i] for i in item]
            boxes_xyxy_abs = self.boxes_xyxy_abs[item]
            return DensePoseList(
                densepose_datas_rel, boxes_xyxy_abs, self.image_size_hw, self.device
            )