| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| import numpy as np |
| import mmcv |
| from mmdet.datasets.builder import PIPELINES |
| import torch |
| from PIL import Image |
|
|
|
|
| @PIPELINES.register_module() |
| class PadMultiViewImage(): |
| """Pad the multi-view image. |
| There are two padding modes: (1) pad to a fixed size and (2) pad to the |
| minimum size that is divisible by some number. |
| Added keys are "pad_shape", "pad_fixed_size", "pad_size_divisor", |
| Args: |
| size (tuple, optional): Fixed padding size. |
| size_divisor (int, optional): The divisor of padded size. |
| pad_val (float, optional): Padding value, 0 by default. |
| """ |
| def __init__(self, size=None, size_divisor=None, pad_val=0): |
| self.size = size |
| self.size_divisor = size_divisor |
| self.pad_val = pad_val |
| assert size is not None or size_divisor is not None |
| assert size_divisor is None or size is None |
| |
| def _pad_img(self, results): |
| """Pad images according to ``self.size``.""" |
| if self.size is not None: |
| padded_img = [mmcv.impad(img, |
| shape = self.size, pad_val=self.pad_val) for img in results['img']] |
| elif self.size_divisor is not None: |
| padded_img = [mmcv.impad_to_multiple(img, |
| self.size_divisor, pad_val=self.pad_val) for img in results['img']] |
| results['img_shape'] = [img.shape for img in results['img']] |
| results['img'] = padded_img |
| results['pad_shape'] = [img.shape for img in padded_img] |
| results['pad_fix_size'] = self.size |
| results['pad_size_divisor'] = self.size_divisor |
| |
| def __call__(self, results): |
| """Call function to pad images, masks, semantic segmentation maps. |
| Args: |
| results (dict): Result dict from loading pipeline. |
| Returns: |
| dict: Updated result dict. |
| """ |
| self._pad_img(results) |
| return results |
|
|
|
|
| def __repr__(self): |
| repr_str = self.__class__.__name__ |
| repr_str += f'(size={self.size}, ' |
| repr_str += f'size_divisor={self.size_divisor}, ' |
| repr_str += f'pad_val={self.pad_val})' |
| return repr_str |
|
|
|
|
| @PIPELINES.register_module() |
| class NormalizeMultiviewImage(object): |
| """Normalize the image. |
| Added key is "img_norm_cfg". |
| Args: |
| mean (sequence): Mean values of 3 channels. |
| std (sequence): Std values of 3 channels. |
| to_rgb (bool): Whether to convert the image from BGR to RGB, |
| default is true. |
| """ |
|
|
| def __init__(self, mean, std, to_rgb=True): |
| self.mean = np.array(mean, dtype=np.float32) |
| self.std = np.array(std, dtype=np.float32) |
| self.to_rgb = to_rgb |
|
|
| def __call__(self, results): |
| """Call function to normalize images. |
| Args: |
| results (dict): Result dict from loading pipeline. |
| Returns: |
| dict: Normalized results, 'img_norm_cfg' key is added into |
| result dict. |
| """ |
| results['img'] = [mmcv.imnormalize( |
| img, self.mean, self.std, self.to_rgb) for img in results['img']] |
| results['img_norm_cfg'] = dict( |
| mean=self.mean, std=self.std, to_rgb=self.to_rgb) |
| return results |
|
|
| def __repr__(self): |
| repr_str = self.__class__.__name__ |
| repr_str += f'(mean={self.mean}, std={self.std}, to_rgb={self.to_rgb})' |
| return repr_str |
|
|
|
|
| @PIPELINES.register_module() |
| class ResizeCropFlipRotImage(): |
| def __init__(self, data_aug_conf=None, with_2d=True, filter_invisible=True, training=True): |
| self.data_aug_conf = data_aug_conf |
| self.training = training |
| self.min_size = 2.0 |
| self.with_2d = with_2d |
| self.filter_invisible = filter_invisible |
|
|
| def __call__(self, results): |
|
|
| imgs = results['img'] |
| N = len(imgs) |
| new_imgs = [] |
| new_gt_bboxes = [] |
| new_centers2d = [] |
| new_gt_labels = [] |
| new_depths = [] |
| assert self.data_aug_conf['rot_lim'] == (0.0, 0.0), "Rotation is not currently supported" |
|
|
| resize, resize_dims, crop, flip, rotate = self._sample_augmentation() |
|
|
|
|
| for i in range(N): |
| img = Image.fromarray(np.uint8(imgs[i])) |
| img, ida_mat = self._img_transform( |
| img, |
| resize=resize, |
| resize_dims=resize_dims, |
| crop=crop, |
| flip=flip, |
| rotate=rotate, |
| ) |
| if self.training and self.with_2d: |
| gt_bboxes = results['gt_bboxes'][i] |
| centers2d = results['centers2d'][i] |
| gt_labels = results['gt_labels'][i] |
| depths = results['depths'][i] |
| if len(gt_bboxes) != 0: |
| gt_bboxes, centers2d, gt_labels, depths = self._bboxes_transform( |
| gt_bboxes, |
| centers2d, |
| gt_labels, |
| depths, |
| resize=resize, |
| crop=crop, |
| flip=flip, |
| ) |
| if len(gt_bboxes) != 0 and self.filter_invisible: |
| gt_bboxes, centers2d, gt_labels, depths = self._filter_invisible(gt_bboxes, centers2d, gt_labels, depths) |
|
|
| new_gt_bboxes.append(gt_bboxes) |
| new_centers2d.append(centers2d) |
| new_gt_labels.append(gt_labels) |
| new_depths.append(depths) |
|
|
| new_imgs.append(np.array(img).astype(np.float32)) |
| results['intrinsics'][i][:3, :3] = ida_mat @ results['intrinsics'][i][:3, :3] |
| results['gt_bboxes'] = new_gt_bboxes |
| results['centers2d'] = new_centers2d |
| results['gt_labels'] = new_gt_labels |
| results['depths'] = new_depths |
| results['img'] = new_imgs |
| results['lidar2img'] = [results['intrinsics'][i] @ results['extrinsics'][i] for i in range(len(results['extrinsics']))] |
|
|
| return results |
|
|
| def _bboxes_transform(self, bboxes, centers2d, gt_labels, depths,resize, crop, flip): |
| assert len(bboxes) == len(centers2d) == len(gt_labels) == len(depths) |
| fH, fW = self.data_aug_conf["final_dim"] |
| bboxes = bboxes * resize |
| bboxes[:, 0] = bboxes[:, 0] - crop[0] |
| bboxes[:, 1] = bboxes[:, 1] - crop[1] |
| bboxes[:, 2] = bboxes[:, 2] - crop[0] |
| bboxes[:, 3] = bboxes[:, 3] - crop[1] |
| bboxes[:, 0] = np.clip(bboxes[:, 0], 0, fW) |
| bboxes[:, 2] = np.clip(bboxes[:, 2], 0, fW) |
| bboxes[:, 1] = np.clip(bboxes[:, 1], 0, fH) |
| bboxes[:, 3] = np.clip(bboxes[:, 3], 0, fH) |
| keep = ((bboxes[:, 2] - bboxes[:, 0]) >= self.min_size) & ((bboxes[:, 3] - bboxes[:, 1]) >= self.min_size) |
|
|
|
|
| if flip: |
| x0 = bboxes[:, 0].copy() |
| x1 = bboxes[:, 2].copy() |
| bboxes[:, 2] = fW - x0 |
| bboxes[:, 0] = fW - x1 |
| bboxes = bboxes[keep] |
|
|
| centers2d = centers2d * resize |
| centers2d[:, 0] = centers2d[:, 0] - crop[0] |
| centers2d[:, 1] = centers2d[:, 1] - crop[1] |
| centers2d[:, 0] = np.clip(centers2d[:, 0], 0, fW) |
| centers2d[:, 1] = np.clip(centers2d[:, 1], 0, fH) |
| if flip: |
| centers2d[:, 0] = fW - centers2d[:, 0] |
|
|
| centers2d = centers2d[keep] |
| gt_labels = gt_labels[keep] |
| depths = depths[keep] |
|
|
| return bboxes, centers2d, gt_labels, depths |
|
|
|
|
| def _filter_invisible(self, bboxes, centers2d, gt_labels, depths): |
| |
| assert len(bboxes) == len(centers2d) == len(gt_labels) == len(depths) |
| fH, fW = self.data_aug_conf["final_dim"] |
| indices_maps = np.zeros((fH,fW)) |
| tmp_bboxes = np.zeros_like(bboxes) |
| tmp_bboxes[:, :2] = np.ceil(bboxes[:, :2]) |
| tmp_bboxes[:, 2:] = np.floor(bboxes[:, 2:]) |
| tmp_bboxes = tmp_bboxes.astype(np.int64) |
| sort_idx = np.argsort(-depths, axis=0, kind='stable') |
| tmp_bboxes = tmp_bboxes[sort_idx] |
| bboxes = bboxes[sort_idx] |
| depths = depths[sort_idx] |
| centers2d = centers2d[sort_idx] |
| gt_labels = gt_labels[sort_idx] |
| for i in range(bboxes.shape[0]): |
| u1, v1, u2, v2 = tmp_bboxes[i] |
| indices_maps[v1:v2, u1:u2] = i |
| indices_res = np.unique(indices_maps).astype(np.int64) |
| bboxes = bboxes[indices_res] |
| depths = depths[indices_res] |
| centers2d = centers2d[indices_res] |
| gt_labels = gt_labels[indices_res] |
|
|
| return bboxes, centers2d, gt_labels, depths |
|
|
|
|
|
|
| def _get_rot(self, h): |
| return torch.Tensor( |
| [ |
| [np.cos(h), np.sin(h)], |
| [-np.sin(h), np.cos(h)], |
| ] |
| ) |
|
|
| def _img_transform(self, img, resize, resize_dims, crop, flip, rotate): |
| ida_rot = torch.eye(2) |
| ida_tran = torch.zeros(2) |
| |
| img = img.resize(resize_dims) |
| img = img.crop(crop) |
| if flip: |
| img = img.transpose(method=Image.FLIP_LEFT_RIGHT) |
| img = img.rotate(rotate) |
|
|
| |
| ida_rot *= resize |
| ida_tran -= torch.Tensor(crop[:2]) |
| if flip: |
| A = torch.Tensor([[-1, 0], [0, 1]]) |
| b = torch.Tensor([crop[2] - crop[0], 0]) |
| ida_rot = A.matmul(ida_rot) |
| ida_tran = A.matmul(ida_tran) + b |
| A = self._get_rot(rotate / 180 * np.pi) |
| b = torch.Tensor([crop[2] - crop[0], crop[3] - crop[1]]) / 2 |
| b = A.matmul(-b) + b |
| ida_rot = A.matmul(ida_rot) |
| ida_tran = A.matmul(ida_tran) + b |
| ida_mat = torch.eye(3) |
| ida_mat[:2, :2] = ida_rot |
| ida_mat[:2, 2] = ida_tran |
| return img, ida_mat |
|
|
| def _sample_augmentation(self): |
| H, W = self.data_aug_conf["H"], self.data_aug_conf["W"] |
| fH, fW = self.data_aug_conf["final_dim"] |
| if self.training: |
| resize = np.random.uniform(*self.data_aug_conf["resize_lim"]) |
| resize_dims = (int(W * resize), int(H * resize)) |
| newW, newH = resize_dims |
| crop_h = int((1 - np.random.uniform(*self.data_aug_conf["bot_pct_lim"])) * newH) - fH |
| crop_w = int(np.random.uniform(0, max(0, newW - fW))) |
| crop = (crop_w, crop_h, crop_w + fW, crop_h + fH) |
| flip = False |
| if self.data_aug_conf["rand_flip"] and np.random.choice([0, 1]): |
| flip = True |
| rotate = np.random.uniform(*self.data_aug_conf["rot_lim"]) |
| else: |
| resize = max(fH / H, fW / W) |
| resize_dims = (int(W * resize), int(H * resize)) |
| newW, newH = resize_dims |
| crop_h = int((1 - np.mean(self.data_aug_conf["bot_pct_lim"])) * newH) - fH |
| crop_w = int(max(0, newW - fW) / 2) |
| crop = (crop_w, crop_h, crop_w + fW, crop_h + fH) |
| flip = False |
| rotate = 0 |
| return resize, resize_dims, crop, flip, rotate |
|
|
| @PIPELINES.register_module() |
| class GlobalRotScaleTransImage(): |
| def __init__( |
| self, |
| rot_range=[-0.3925, 0.3925], |
| scale_ratio_range=[0.95, 1.05], |
| translation_std=[0, 0, 0], |
| reverse_angle=False, |
| training=True, |
| ): |
|
|
| self.rot_range = rot_range |
| self.scale_ratio_range = scale_ratio_range |
| self.translation_std = translation_std |
|
|
| self.reverse_angle = reverse_angle |
| self.training = training |
|
|
| def __call__(self, results): |
| |
| translation_std = np.array(self.translation_std, dtype=np.float32) |
|
|
| rot_angle = np.random.uniform(*self.rot_range) |
| scale_ratio = np.random.uniform(*self.scale_ratio_range) |
| trans = np.random.normal(scale=translation_std, size=3).T |
|
|
| self._rotate_bev_along_z(results, rot_angle) |
| if self.reverse_angle: |
| rot_angle = rot_angle * -1 |
| results["gt_bboxes_3d"].rotate( |
| np.array(rot_angle) |
| ) |
|
|
| |
| self._scale_xyz(results, scale_ratio) |
| results["gt_bboxes_3d"].scale(scale_ratio) |
|
|
| |
| self._trans_xyz(results, trans) |
| results["gt_bboxes_3d"].translate(trans) |
|
|
| return results |
|
|
| def _trans_xyz(self, results, trans): |
| trans_mat = torch.eye(4, 4) |
| trans_mat[:3, -1] = torch.from_numpy(trans).reshape(1, 3) |
| trans_mat_inv = torch.inverse(trans_mat) |
| num_view = len(results["lidar2img"]) |
| results['ego_pose'] = (torch.tensor(results["ego_pose"]).float() @ trans_mat_inv).numpy() |
| results['ego_pose_inv'] = (trans_mat.float() @ torch.tensor(results["ego_pose_inv"])).numpy() |
|
|
| for view in range(num_view): |
| results["lidar2img"][view] = (torch.tensor(results["lidar2img"][view]).float() @ trans_mat_inv).numpy() |
|
|
|
|
| def _rotate_bev_along_z(self, results, angle): |
| rot_cos = torch.cos(torch.tensor(angle)) |
| rot_sin = torch.sin(torch.tensor(angle)) |
|
|
| rot_mat = torch.tensor([[rot_cos, rot_sin, 0, 0], [-rot_sin, rot_cos, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]) |
| rot_mat_inv = torch.inverse(rot_mat) |
|
|
| results['ego_pose'] = (torch.tensor(results["ego_pose"]).float() @ rot_mat_inv).numpy() |
| results['ego_pose_inv'] = (rot_mat.float() @ torch.tensor(results["ego_pose_inv"])).numpy() |
| num_view = len(results["lidar2img"]) |
| for view in range(num_view): |
| results["lidar2img"][view] = (torch.tensor(results["lidar2img"][view]).float() @ rot_mat_inv).numpy() |
|
|
| def _scale_xyz(self, results, scale_ratio): |
| scale_mat = torch.tensor( |
| [ |
| [scale_ratio, 0, 0, 0], |
| [0, scale_ratio, 0, 0], |
| [0, 0, scale_ratio, 0], |
| [0, 0, 0, 1], |
| ] |
| ) |
|
|
| scale_mat_inv = torch.inverse(scale_mat) |
|
|
| results['ego_pose'] = (torch.tensor(results["ego_pose"]).float() @ scale_mat_inv).numpy() |
| results['ego_pose_inv'] = (scale_mat @ torch.tensor(results["ego_pose_inv"]).float()).numpy() |
|
|
| num_view = len(results["lidar2img"]) |
| for view in range(num_view): |
| results["lidar2img"][view] = (torch.tensor(results["lidar2img"][view]).float() @ scale_mat_inv).numpy() |
|
|
