# 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