densepose/modeling/test_time_augmentation.py

# Copyright (c) Facebook, Inc. and its affiliates.
import copy
import numpy as np
import torch
from fvcore.transforms import HFlipTransform, TransformList
from torch.nn import functional as F

from detectron2.data.transforms import RandomRotation, RotationTransform, apply_transform_gens
from detectron2.modeling.postprocessing import detector_postprocess
from detectron2.modeling.test_time_augmentation import DatasetMapperTTA, GeneralizedRCNNWithTTA

from ..converters import HFlipConverter


class DensePoseDatasetMapperTTA(DatasetMapperTTA):
    def __init__(self, cfg):
        super().__init__(cfg=cfg)
        self.angles = cfg.TEST.AUG.ROTATION_ANGLES

    def __call__(self, dataset_dict):
        ret = super().__call__(dataset_dict=dataset_dict)
        numpy_image = dataset_dict["image"].permute(1, 2, 0).numpy()
        for angle in self.angles:
            rotate = RandomRotation(angle=angle, expand=True)
            new_numpy_image, tfms = apply_transform_gens([rotate], np.copy(numpy_image))
            torch_image = torch.from_numpy(np.ascontiguousarray(new_numpy_image.transpose(2, 0, 1)))
            dic = copy.deepcopy(dataset_dict)
            # In DatasetMapperTTA, there is a pre_tfm transform (resize or no-op) that is
            # added at the beginning of each TransformList. That's '.transforms[0]'.
            dic["transforms"] = TransformList(
                [ret[-1]["transforms"].transforms[0]] + tfms.transforms
            )
            dic["image"] = torch_image
            ret.append(dic)
        return ret


class DensePoseGeneralizedRCNNWithTTA(GeneralizedRCNNWithTTA):
    def __init__(self, cfg, model, transform_data, tta_mapper=None, batch_size=1):
        """
        Args:
            cfg (CfgNode):
            model (GeneralizedRCNN): a GeneralizedRCNN to apply TTA on.
            transform_data (DensePoseTransformData): contains symmetry label
                transforms used for horizontal flip
            tta_mapper (callable): takes a dataset dict and returns a list of
                augmented versions of the dataset dict. Defaults to
                `DatasetMapperTTA(cfg)`.
            batch_size (int): batch the augmented images into this batch size for inference.
        """
        self._transform_data = transform_data.to(model.device)
        super().__init__(cfg=cfg, model=model, tta_mapper=tta_mapper, batch_size=batch_size)

    # the implementation follows closely the one from detectron2/modeling
    def _inference_one_image(self, input):
        """
        Args:
            input (dict): one dataset dict with "image" field being a CHW tensor

        Returns:
            dict: one output dict
        """
        orig_shape = (input["height"], input["width"])
        # For some reason, resize with uint8 slightly increases box AP but decreases densepose AP
        input["image"] = input["image"].to(torch.uint8)
        augmented_inputs, tfms = self._get_augmented_inputs(input)
        # Detect boxes from all augmented versions
        with self._turn_off_roi_heads(["mask_on", "keypoint_on", "densepose_on"]):
            # temporarily disable roi heads
            all_boxes, all_scores, all_classes = self._get_augmented_boxes(augmented_inputs, tfms)
        merged_instances = self._merge_detections(all_boxes, all_scores, all_classes, orig_shape)

        if self.cfg.MODEL.MASK_ON or self.cfg.MODEL.DENSEPOSE_ON:
            # Use the detected boxes to obtain new fields
            augmented_instances = self._rescale_detected_boxes(
                augmented_inputs, merged_instances, tfms
            )
            # run forward on the detected boxes
            outputs = self._batch_inference(augmented_inputs, augmented_instances)
            # Delete now useless variables to avoid being out of memory
            del augmented_inputs, augmented_instances
            # average the predictions
            if self.cfg.MODEL.MASK_ON:
                merged_instances.pred_masks = self._reduce_pred_masks(outputs, tfms)
            if self.cfg.MODEL.DENSEPOSE_ON:
                merged_instances.pred_densepose = self._reduce_pred_densepose(outputs, tfms)
            # postprocess
            merged_instances = detector_postprocess(merged_instances, *orig_shape)
            return {"instances": merged_instances}
        else:
            return {"instances": merged_instances}

    def _get_augmented_boxes(self, augmented_inputs, tfms):
        # Heavily based on detectron2/modeling/test_time_augmentation.py
        # Only difference is that RotationTransform is excluded from bbox computation
        # 1: forward with all augmented images
        outputs = self._batch_inference(augmented_inputs)
        # 2: union the results
        all_boxes = []
        all_scores = []
        all_classes = []
        for output, tfm in zip(outputs, tfms):
            # Need to inverse the transforms on boxes, to obtain results on original image
            if not any(isinstance(t, RotationTransform) for t in tfm.transforms):
                # Some transforms can't compute bbox correctly
                pred_boxes = output.pred_boxes.tensor
                original_pred_boxes = tfm.inverse().apply_box(pred_boxes.cpu().numpy())
                all_boxes.append(torch.from_numpy(original_pred_boxes).to(pred_boxes.device))
                all_scores.extend(output.scores)
                all_classes.extend(output.pred_classes)
        all_boxes = torch.cat(all_boxes, dim=0)
        return all_boxes, all_scores, all_classes

    def _reduce_pred_densepose(self, outputs, tfms):
        # Should apply inverse transforms on densepose preds.
        # We assume only rotation, resize & flip are used. pred_masks is a scale-invariant
        # representation, so we handle the other ones specially
        for idx, (output, tfm) in enumerate(zip(outputs, tfms)):
            for t in tfm.transforms:
                for attr in ["coarse_segm", "fine_segm", "u", "v"]:
                    setattr(
                        output.pred_densepose,
                        attr,
                        _inverse_rotation(
                            getattr(output.pred_densepose, attr), output.pred_boxes.tensor, t
                        ),
                    )
            if any(isinstance(t, HFlipTransform) for t in tfm.transforms):
                output.pred_densepose = HFlipConverter.convert(
                    output.pred_densepose, self._transform_data
                )
            self._incremental_avg_dp(outputs[0].pred_densepose, output.pred_densepose, idx)
        return outputs[0].pred_densepose

    # incrementally computed average: u_(n + 1) = u_n + (x_(n+1) - u_n) / (n + 1).
    def _incremental_avg_dp(self, avg, new_el, idx):
        for attr in ["coarse_segm", "fine_segm", "u", "v"]:
            setattr(avg, attr, (getattr(avg, attr) * idx + getattr(new_el, attr)) / (idx + 1))
            if idx:
                # Deletion of the > 0 index intermediary values to prevent GPU OOM
                setattr(new_el, attr, None)
        return avg


def _inverse_rotation(densepose_attrs, boxes, transform):
    # resample outputs to image size and rotate back the densepose preds
    # on the rotated images to the space of the original image
    if len(boxes) == 0 or not isinstance(transform, RotationTransform):
        return densepose_attrs
    boxes = boxes.int().cpu().numpy()
    wh_boxes = boxes[:, 2:] - boxes[:, :2]  # bboxes in the rotated space
    inv_boxes = rotate_box_inverse(transform, boxes).astype(int)  # bboxes in original image
    wh_diff = (inv_boxes[:, 2:] - inv_boxes[:, :2] - wh_boxes) // 2  # diff between new/old bboxes
    rotation_matrix = torch.tensor([transform.rm_image]).to(device=densepose_attrs.device).float()
    rotation_matrix[:, :, -1] = 0
    # To apply grid_sample for rotation, we need to have enough space to fit the original and
    # rotated bboxes. l_bds and r_bds are the left/right bounds that will be used to
    # crop the difference once the rotation is done
    l_bds = np.maximum(0, -wh_diff)
    for i in range(len(densepose_attrs)):
        if min(wh_boxes[i]) <= 0:
            continue
        densepose_attr = densepose_attrs[[i]].clone()
        # 1. Interpolate densepose attribute to size of the rotated bbox
        densepose_attr = F.interpolate(densepose_attr, wh_boxes[i].tolist()[::-1], mode="bilinear")
        # 2. Pad the interpolated attribute so it has room for the original + rotated bbox
        densepose_attr = F.pad(densepose_attr, tuple(np.repeat(np.maximum(0, wh_diff[i]), 2)))
        # 3. Compute rotation grid and transform
        grid = F.affine_grid(rotation_matrix, size=densepose_attr.shape)
        densepose_attr = F.grid_sample(densepose_attr, grid)
        # 4. Compute right bounds and crop the densepose_attr to the size of the original bbox
        r_bds = densepose_attr.shape[2:][::-1] - l_bds[i]
        densepose_attr = densepose_attr[:, :, l_bds[i][1] : r_bds[1], l_bds[i][0] : r_bds[0]]
        if min(densepose_attr.shape) > 0:
            # Interpolate back to the original size of the densepose attribute
            densepose_attr = F.interpolate(
                densepose_attr, densepose_attrs.shape[-2:], mode="bilinear"
            )
            # Adding a very small probability to the background class to fill padded zones
            densepose_attr[:, 0] += 1e-10
            densepose_attrs[i] = densepose_attr
    return densepose_attrs


def rotate_box_inverse(rot_tfm, rotated_box):
    """
    rotated_box is a N * 4 array of [x0, y0, x1, y1] boxes
    When a bbox is rotated, it gets bigger, because we need to surround the tilted bbox
    So when a bbox is rotated then inverse-rotated, it is much bigger than the original
    This function aims to invert the rotation on the box, but also resize it to its original size
    """
    # 1. Compute the inverse rotation of the rotated bboxes (bigger than it )
    invrot_box = rot_tfm.inverse().apply_box(rotated_box)
    h, w = rotated_box[:, 3] - rotated_box[:, 1], rotated_box[:, 2] - rotated_box[:, 0]
    ih, iw = invrot_box[:, 3] - invrot_box[:, 1], invrot_box[:, 2] - invrot_box[:, 0]
    assert 2 * rot_tfm.abs_sin**2 != 1, "45 degrees angle can't be inverted"
    # 2. Inverse the corresponding computation in the rotation transform
    # to get the original height/width of the rotated boxes
    orig_h = (h * rot_tfm.abs_cos - w * rot_tfm.abs_sin) / (1 - 2 * rot_tfm.abs_sin**2)
    orig_w = (w * rot_tfm.abs_cos - h * rot_tfm.abs_sin) / (1 - 2 * rot_tfm.abs_sin**2)
    # 3. Resize the inverse-rotated bboxes to their original size
    invrot_box[:, 0] += (iw - orig_w) / 2
    invrot_box[:, 1] += (ih - orig_h) / 2
    invrot_box[:, 2] -= (iw - orig_w) / 2
    invrot_box[:, 3] -= (ih - orig_h) / 2

    return invrot_box