from copy import deepcopy from pathlib import Path from typing import Any, Dict, List import numpy as np import torch import torch.utils.data as torchdata import torchvision.transforms as tvf from PIL import Image from pathlib import Path from ...models.utils import deg2rad, rotmat2d from ...utils.io import read_image from ...utils.wrappers import Camera from ..image import pad_image, rectify_image, resize_image from ..utils import decompose_rotmat from ..schema import MIADataConfiguration class MapLocDataset(torchdata.Dataset): def __init__( self, stage: str, cfg: MIADataConfiguration, names: List[str], data: Dict[str, Any], image_dirs: Dict[str, Path], seg_mask_dirs: Dict[str, Path], flood_masks_dirs: Dict[str, Path], image_ext: str = "", ): self.stage = stage self.cfg = deepcopy(cfg) self.data = data self.image_dirs = image_dirs self.seg_mask_dirs = seg_mask_dirs self.flood_masks_dirs = flood_masks_dirs self.names = names self.image_ext = image_ext tfs = [] self.tfs = tvf.Compose(tfs) self.augmentations = self.get_augmentations() def __len__(self): return len(self.names) def __getitem__(self, idx): if self.stage == "train" and self.cfg.random: seed = None else: seed = [self.cfg.seed, idx] (seed,) = np.random.SeedSequence(seed).generate_state(1) scene, seq, name = self.names[idx] view = self.get_view( idx, scene, seq, name, seed ) return view def get_augmentations(self): if self.stage != "train" or not self.cfg.augmentations.enabled: print(f"No Augmentation!", "\n" * 10) self.cfg.augmentations.random_flip = 0.0 return tvf.Compose([]) print(f"Augmentation!", "\n" * 10) augmentations = [ tvf.ColorJitter( brightness=self.cfg.augmentations.brightness, contrast=self.cfg.augmentations.contrast, saturation=self.cfg.augmentations.saturation, hue=self.cfg.augmentations.hue, ) ] if self.cfg.augmentations.random_resized_crop: augmentations.append( tvf.RandomResizedCrop(scale=(0.8, 1.0)) ) # RandomResizedCrop if self.cfg.augmentations.gaussian_noise.enabled: augmentations.append( tvf.GaussianNoise( mean=self.cfg.augmentations.gaussian_noise.mean, std=self.cfg.augmentations.gaussian_noise.std, ) ) # Gaussian noise if self.cfg.augmentations.brightness_contrast.enabled: augmentations.append( tvf.ColorJitter( brightness=self.cfg.augmentations.brightness_contrast.brightness_factor, contrast=self.cfg.augmentations.brightness_contrast.contrast_factor, saturation=0, # Keep saturation at 0 for brightness and contrast adjustment hue=0, ) ) # Brightness and contrast adjustment return tvf.Compose(augmentations) def random_flip(self, image, cam, valid, seg_mask, flood_mask, conf_mask): if torch.rand(1) < self.cfg.augmentations.random_flip: image = torch.flip(image, [-1]) cam = cam.flip() valid = torch.flip(valid, [-1]) seg_mask = torch.flip(seg_mask, [1]) flood_mask = torch.flip(flood_mask, [-1]) conf_mask = torch.flip(conf_mask, [-1]) return image, cam, valid, seg_mask, flood_mask, conf_mask def get_view(self, idx, scene, seq, name, seed): data = { "index": idx, "name": name, "scene": scene, "sequence": seq, } cam_dict = self.data["cameras"][scene][seq][self.data["camera_id"][idx]] cam = Camera.from_dict(cam_dict).float() if "roll_pitch_yaw" in self.data: roll, pitch, yaw = self.data["roll_pitch_yaw"][idx].numpy() else: roll, pitch, yaw = decompose_rotmat( self.data["R_c2w"][idx].numpy()) image = read_image(self.image_dirs[scene] / (name + self.image_ext)) image = Image.fromarray(image) image = self.augmentations(image) image = np.array(image) if "plane_params" in self.data: # transform the plane parameters from world to camera frames plane_w = self.data["plane_params"][idx] data["ground_plane"] = torch.cat( [rotmat2d(deg2rad(torch.tensor(yaw))) @ plane_w[:2], plane_w[2:]] ) image, valid, cam, roll, pitch = self.process_image( image, cam, roll, pitch, seed ) if "chunk_index" in self.data: # TODO: (cherie) do we need this? data["chunk_id"] = (scene, seq, self.data["chunk_index"][idx]) # Semantic map extraction seg_mask_path = self.seg_mask_dirs[scene] / \ (name.split("_")[0] + ".npy") seg_masks_ours = np.load(seg_mask_path) mask_center = ( seg_masks_ours.shape[0] // 2, seg_masks_ours.shape[1] // 2) seg_masks_ours = seg_masks_ours[mask_center[0] - 100:mask_center[0], mask_center[1] - 50: mask_center[1] + 50] if self.cfg.num_classes == 6: seg_masks_ours = seg_masks_ours[..., [0, 1, 2, 4, 6, 7]] flood_mask_path = self.flood_masks_dirs[scene] / \ (name.split("_")[0] + ".npy") flood_mask = np.load(flood_mask_path) flood_mask = flood_mask[mask_center[0]-100:mask_center[0], mask_center[1] - 50: mask_center[1] + 50] confidence_map = flood_mask.copy() confidence_map = (confidence_map - confidence_map.min()) / \ (confidence_map.max() - confidence_map.min() + 1e-6) seg_masks_ours = torch.from_numpy(seg_masks_ours).float() flood_mask = torch.from_numpy(flood_mask).float() confidence_map = torch.from_numpy(confidence_map).float() # Map Augmentations with torch.random.fork_rng(devices=[]): torch.manual_seed(seed) image, cam, valid, seg_masks_ours, flood_mask, confidence_map = self.random_flip( image, cam, valid, seg_masks_ours, flood_mask, confidence_map) return { **data, "image": image, "valid": valid, "camera": cam, "seg_masks": seg_masks_ours, "flood_masks": flood_mask, "roll_pitch_yaw": torch.tensor((roll, pitch, yaw)).float(), "confidence_map": confidence_map # "pixels_per_meter": torch.tensor(canvas.ppm).float(), } def process_image(self, image, cam, roll, pitch, seed): image = ( torch.from_numpy(np.ascontiguousarray(image)) .permute(2, 0, 1) .float() .div_(255) ) if not self.cfg.gravity_align: # Turn off gravity alignment roll = 0.0 pitch = 0.0 image, valid = rectify_image(image, cam, roll, pitch) else: image, valid = rectify_image( image, cam, roll, pitch if self.cfg.rectify_pitch else None ) roll = 0.0 if self.cfg.rectify_pitch: pitch = 0.0 if self.cfg.target_focal_length is not None: # Resize to a canonical focal length factor = self.cfg.target_focal_length / cam.f.numpy() size = (np.array(image.shape[-2:][::-1]) * factor).astype(int) image, _, cam, valid = resize_image( image, size, camera=cam, valid=valid) size_out = self.cfg.resize_image if size_out is None: # Round the edges up such that they are multiple of a factor stride = self.cfg.pad_to_multiple size_out = (np.ceil((size / stride)) * stride).astype(int) # Crop or pad such that both edges are of the given size image, valid, cam = pad_image( image, size_out, cam, valid, crop_and_center=True ) elif self.cfg.resize_image is not None: image, _, cam, valid = resize_image( image, self.cfg.resize_image, fn=max, camera=cam, valid=valid ) if self.cfg.pad_to_square: # Pad such that both edges are of the given size image, valid, cam = pad_image( image, self.cfg.resize_image, cam, valid) if self.cfg.reduce_fov is not None: h, w = image.shape[-2:] f = float(cam.f[0]) fov = np.arctan(w / f / 2) w_new = round(2 * f * np.tan(self.cfg.reduce_fov * fov)) image, valid, cam = pad_image( image, (w_new, h), cam, valid, crop_and_center=True ) with torch.random.fork_rng(devices=[]): torch.manual_seed(seed) image = self.tfs(image) return image, valid, cam, roll, pitch