""" Default Datasets Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com), Yujia Zhang (yujia.zhang.cs@gmail.com) Please cite our work if the code is helpful to you. """ import os import glob import json import numpy as np import torch from copy import deepcopy from torch.utils.data import Dataset from collections.abc import Sequence from torchvision.transforms import InterpolationMode from PIL import Image from torchvision.transforms import transforms as T from pointcept.utils.logger import get_root_logger from pointcept.utils.cache import shared_dict from .builder import DATASETS, build_dataset from .transform import Compose, TRANSFORMS INTERPOLATION_MODE = { "bilinear": InterpolationMode.BILINEAR, "bicubic": InterpolationMode.BICUBIC, "nearest": InterpolationMode.NEAREST, } @DATASETS.register_module() class DefaultDataset(Dataset): VALID_ASSETS = [ "coord", "color", "normal", "strength", "segment", "instance", "pose", ] def __init__( self, split="train", data_root="data/dataset", transform=None, test_mode=False, test_cfg=None, cache=False, ignore_index=-1, loop=1, ): super(DefaultDataset, self).__init__() self.data_root = data_root self.split = split self.transform = Compose(transform) self.cache = cache self.ignore_index = ignore_index self.loop = ( loop if not test_mode else 1 ) # force make loop = 1 while in test mode self.test_mode = test_mode self.test_cfg = test_cfg if test_mode else None if test_mode: self.test_voxelize = TRANSFORMS.build(self.test_cfg.voxelize) self.test_crop = ( TRANSFORMS.build(self.test_cfg.crop) if self.test_cfg.crop else None ) self.post_transform = Compose(self.test_cfg.post_transform) self.aug_transform = [Compose(aug) for aug in self.test_cfg.aug_transform] self.data_list = self.get_data_list() logger = get_root_logger() logger.info( "Totally {} x {} samples in {} {} set.".format( len(self.data_list), self.loop, os.path.basename(self.data_root), split ) ) def get_data_list(self): if isinstance(self.split, str): split_list = [self.split] elif isinstance(self.split, Sequence): split_list = self.split else: raise NotImplementedError data_list = [] for split in split_list: if os.path.isfile(os.path.join(self.data_root, split)): with open(os.path.join(self.data_root, split)) as f: data_list += [ os.path.join(self.data_root, data) for data in json.load(f) ] else: data_list += glob.glob(os.path.join(self.data_root, split, "*")) return data_list def get_data(self, idx): data_path = self.data_list[idx % len(self.data_list)] name = self.get_data_name(idx) split = self.get_split_name(idx) if self.cache: cache_name = f"pointcept-{name}" return shared_dict(cache_name) data_dict = {} assets = os.listdir(data_path) for asset in assets: if not asset.endswith(".npy"): continue if asset[:-4] not in self.VALID_ASSETS: continue data_dict[asset[:-4]] = np.load(os.path.join(data_path, asset)) data_dict["name"] = name data_dict["split"] = split if "coord" in data_dict.keys(): data_dict["coord"] = data_dict["coord"].astype(np.float32) if "color" in data_dict.keys(): data_dict["color"] = data_dict["color"].astype(np.float32) if "normal" in data_dict.keys(): data_dict["normal"] = data_dict["normal"].astype(np.float32) if "segment" in data_dict.keys(): data_dict["segment"] = data_dict["segment"].reshape([-1]).astype(np.int32) else: data_dict["segment"] = ( np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1 ) if "instance" in data_dict.keys(): data_dict["instance"] = data_dict["instance"].reshape([-1]).astype(np.int32) else: data_dict["instance"] = ( np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1 ) return data_dict def get_data_name(self, idx): return os.path.basename(self.data_list[idx % len(self.data_list)]) def get_split_name(self, idx): return os.path.basename( os.path.dirname(self.data_list[idx % len(self.data_list)]) ) def prepare_train_data(self, idx): # load data data_dict = self.get_data(idx) data_dict = self.transform(data_dict) return data_dict def prepare_test_data(self, idx): # load data data_dict = self.get_data(idx) data_dict = self.transform(data_dict) result_dict = dict(segment=data_dict.pop("segment"), name=data_dict.pop("name")) if "origin_segment" in data_dict: assert "inverse" in data_dict result_dict["origin_segment"] = data_dict.pop("origin_segment") result_dict["inverse"] = data_dict.pop("inverse") data_dict_list = [] for aug in self.aug_transform: data_dict_list.append(aug(deepcopy(data_dict))) fragment_list = [] for data in data_dict_list: if self.test_voxelize is not None: data_part_list = self.test_voxelize(data) else: data["index"] = np.arange(data["coord"].shape[0]) data_part_list = [data] for data_part in data_part_list: if self.test_crop is not None: data_part = self.test_crop(data_part) else: data_part = [data_part] fragment_list += data_part for i in range(len(fragment_list)): fragment_list[i] = self.post_transform(fragment_list[i]) result_dict["fragment_list"] = fragment_list return result_dict def __getitem__(self, idx): if self.test_mode: return self.prepare_test_data(idx) else: return self.prepare_train_data(idx) def __len__(self): return len(self.data_list) * self.loop @DATASETS.register_module() class DefaultImagePointDataset(Dataset): PC_VALID_ASSETS = [ "coord", "color", "normal", "superpoint", "pose", "segment", "instance", "strength", ] def __init__( self, split="train", data_root="data/dataset", transform=None, test_mode=False, test_cfg=None, cache=False, ignore_index=-1, loop=1, crop_h=630, crop_w=1120, patch_size=14, interpolation="bilinear", if_img=True, ): super(DefaultImagePointDataset, self).__init__() self.if_img = if_img self.data_root = data_root self.split = split self.transform = Compose(transform) self.cache = cache self.ignore_index = ignore_index self.loop = ( loop if not test_mode else 1 ) # force make loop = 1 while in test mode self.test_mode = test_mode self.test_cfg = test_cfg if test_mode else None if test_mode: self.test_voxelize = TRANSFORMS.build(self.test_cfg.voxelize) self.test_crop = ( TRANSFORMS.build(self.test_cfg.crop) if self.test_cfg.crop else None ) self.post_transform = Compose(self.test_cfg.post_transform) self.aug_transform = [Compose(aug) for aug in self.test_cfg.aug_transform] self.data_list, self.split_list = self.get_data_list() if isinstance(self.data_list, dict): self.data_name = list(self.data_list.keys()) elif isinstance(self.data_list[0], dict): self.data_name = list([data["token"] for data in self.data_list]) else: self.data_name = self.data_list logger = get_root_logger() logger.info( "Totally {} x {} samples in {} {} set.".format( len(self.data_name), self.loop, os.path.basename(self.data_root), split ) ) self.crop_h = crop_h self.crop_w = crop_w self.patch_size = patch_size self.patch_h = crop_h // patch_size self.patch_w = crop_w // patch_size self.transform_img = T.Compose( [ T.Resize( (self.patch_h * self.patch_size, self.patch_w * self.patch_size), interpolation=INTERPOLATION_MODE[interpolation], ), T.ToTensor(), ] ) def resize_correspondence_info( self, correspondence, size, size0, crop_size, _alignment ): h, w = size h0, w0 = size0 left, top, right, bottom = crop_size crop_h = bottom - top crop_w = right - left mask_crop = ( (correspondence[:, 1] >= top) & (correspondence[:, 1] < bottom) & (correspondence[:, 0] >= left) & (correspondence[:, 0] < right) ) correspondence = correspondence[mask_crop] correspondence[:, 1] -= top correspondence[:, 0] -= left correspondence[:, 1] = (correspondence[:, 1] * h / crop_h / _alignment).astype( np.float32 ) correspondence[:, 0] = (correspondence[:, 0] * w / crop_w / _alignment).astype( np.float32 ) correspondence = correspondence[:, [1, 0, 2]] correspondence = np.unique(correspondence, axis=0) return correspondence def get_data_list(self): split_list = {} if isinstance(self.split, str): data_path = os.path.join(self.data_root, "splits", f"{self.split}.json") with open(data_path, "r", encoding="utf-8") as file: data_list = json.load(file) split_list[self.split] = list(data_list.keys()) elif isinstance(self.split, Sequence): data_list = {} for split in self.split: data_path = os.path.join(self.data_root, "splits", f"{split}.json") with open(data_path, "r", encoding="utf-8") as file: data_split_dict = json.load(file) data_list.update(data_split_dict) split_list[split] = list(data_split_dict.keys()) else: raise NotImplementedError return data_list, split_list def get_data_name(self, idx): return self.data_name[idx % len(self.data_name)] def get_split_name(self, idx): for split, names in self.split_list.items(): if self.data_name[idx % len(self.data_name)] in names: return split return None def get_data(self, idx): data_dict = {} name = self.get_data_name(idx) split = self.get_split_name(idx) data_dict["name"] = name data_dict["split"] = split data_path = self.data_list[name] pointclouds_path = data_path["pointclouds"] assets = os.listdir(pointclouds_path) for asset in assets: if not asset.endswith(".npy"): continue if asset[:-4] not in self.PC_VALID_ASSETS: continue data_dict[asset[:-4]] = np.load(os.path.join(pointclouds_path, asset)) if self.if_img: imgs_path = data_path["images"] imgs = [Image.open(asset) for asset in imgs_path] if len(imgs) > 0: img_width, img_height = imgs[0].size div_w = img_width // self.patch_w div_h = img_height // self.patch_h div_min = max(min(div_w, div_h), 1) crop_img_width = div_min * self.patch_w crop_img_height = div_min * self.patch_h left = int((img_width - crop_img_width) / 2) top = int((img_height - crop_img_height) / 2) right = int((img_width + crop_img_width) / 2) bottom = int((img_height + crop_img_height) / 2) imgs = [img.crop((left, top, right, bottom)) for img in imgs] imgs = [self.transform_img(img) for img in imgs] imgs_list = torch.stack(imgs) data_dict["images"] = imgs_list.float() else: data_dict["images"] = torch.empty( ( 0, 3, self.patch_h * self.patch_size, self.patch_w * self.patch_size, ) ) data_dict["img_num"] = np.array( [data_dict["images"].shape[0]], dtype=np.int32 ) correspondences_path = data_path["correspondences"] correspondence_infos = np.ones( (data_dict["coord"].shape[0], len(correspondences_path), 2), dtype=np.float32, ) * (-1) for asset_id, asset in enumerate(correspondences_path): correspondence_info = np.load(asset).astype(np.float32) if np.array_equal(correspondence_info, -np.ones((1, 3))): continue correspondence_info = self.resize_correspondence_info( correspondence_info, (self.patch_h * self.patch_size, self.patch_w * self.patch_size), (img_height, img_width), (left, top, right, bottom), self.patch_size, ) correspondence_infos[ correspondence_info[:, -1].astype(np.int32), asset_id, : ] = correspondence_info[:, :-1] data_dict["correspondence"] = correspondence_infos # .reshape(-1, 2) if "coord" in data_dict.keys(): data_dict["coord"] = data_dict["coord"].astype(np.float32) if "color" in data_dict.keys(): data_dict["color"] = data_dict["color"].astype(np.float32) if "normal" in data_dict.keys(): data_dict["normal"] = data_dict["normal"].astype(np.float32) if "segment" in data_dict.keys(): data_dict["segment"] = data_dict["segment"].reshape([-1]).astype(np.int32) else: data_dict["segment"] = ( np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1 ) if "instance" in data_dict.keys(): data_dict["instance"] = data_dict["instance"].reshape([-1]).astype(np.int32) else: data_dict["instance"] = ( np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1 ) return data_dict def prepare_train_data(self, idx): # load data data_dict = self.get_data(idx) data_dict = self.transform(data_dict) return data_dict def prepare_test_data(self, idx): # load data data_dict = self.get_data(idx) data_dict = self.transform(data_dict) result_dict = dict(segment=data_dict.pop("segment"), name=data_dict.pop("name")) if "origin_segment" in data_dict: assert "inverse" in data_dict result_dict["origin_segment"] = data_dict.pop("origin_segment") result_dict["inverse"] = data_dict.pop("inverse") data_dict_list = [] for aug in self.aug_transform: data_dict_list.append(aug(deepcopy(data_dict))) fragment_list = [] for data in data_dict_list: if self.test_voxelize is not None: data_part_list = self.test_voxelize(data) else: data["index"] = np.arange(data["coord"].shape[0]) data_part_list = [data] for data_part in data_part_list: if self.test_crop is not None: data_part = self.test_crop(data_part) else: data_part = [data_part] fragment_list += data_part for i in range(len(fragment_list)): fragment_list[i] = self.post_transform(fragment_list[i]) result_dict["fragment_list"] = fragment_list return result_dict def __getitem__(self, idx): if self.test_mode: return self.prepare_test_data(idx) else: return self.prepare_train_data(idx) def __len__(self): return len(self.data_list) * self.loop @DATASETS.register_module() class DefaultMultiViewImagePointDataset(DefaultImagePointDataset): def __init__(self, depth_scale=1000.0, if_img=True, **kwargs): self.depth_scale = depth_scale self.if_img = if_img super().__init__(if_img=if_img, **kwargs) def get_data_list(self): split_list = {} if isinstance(self.split, str): splits_to_load = [self.split] elif isinstance(self.split, Sequence): splits_to_load = self.split else: raise NotImplementedError data_list = {} for split in splits_to_load: data_path = os.path.join(self.data_root, "splits", f"{split}.json") if not os.path.exists(data_path): raise FileNotFoundError(f"Split file not found: {data_path}") with open(data_path, "r", encoding="utf-8") as file: data_split_dict = json.load(file) data_list.update(data_split_dict) split_list[split] = list(data_split_dict.keys()) return data_list, split_list @staticmethod def get_normals(center, coords): normals = coords - center[None, :] normals = normals / np.linalg.norm(normals, axis=-1, keepdims=True) return normals def get_data(self, idx): name = self.get_data_name(idx) split = self.get_split_name(idx) data_info = self.data_list[name] all_points_world, all_colors, camera_params = [], [], [] correspondences = [] coord_num = 0 for i in range(len(data_info["images"])): rgb_path = data_info["images"][i] depth_path = data_info["depths"][i] pose_path = data_info["Ts"][i] intr_path = data_info["Ks"][i] try: with Image.open(rgb_path) as img: color_img = np.array(img) with Image.open(depth_path) as img: depth_img = np.array(img) cam_pose_world = np.load(pose_path) cam_intr = np.load(intr_path) camera_params.append({"pose": cam_pose_world, "intr": cam_intr}) except FileNotFoundError: continue depth_m = depth_img.astype(np.float32) / self.depth_scale h, w = depth_m.shape fx, fy, cx, cy = ( cam_intr[0, 0], cam_intr[1, 1], cam_intr[0, 2], cam_intr[1, 2], ) v, u = np.indices((h, w)) valid_mask = (depth_m > 0) & (depth_m < 0.8) z_cam = depth_m[valid_mask] x_cam, y_cam = (u[valid_mask] - cx) * z_cam / fx, ( v[valid_mask] - cy ) * z_cam / fy points_cam = np.vstack((x_cam, y_cam, z_cam)).T points_world = ( cam_pose_world @ np.hstack([points_cam, np.ones((points_cam.shape[0], 1))]).T ).T[:, :3] all_points_world.append(points_world) all_colors.append(color_img[valid_mask]) pixel = np.vstack([u[valid_mask], v[valid_mask]]).T pixel = np.concatenate([pixel, np.arange(pixel.shape[0])[:, None]], axis=1) pixel[:, 2] += coord_num coord_num += points_world.shape[0] correspondences.append(pixel) coord = ( np.concatenate(all_points_world, axis=0) if all_points_world else np.zeros((0, 3)) ) coord[:, 2] = -coord[:, 2] color = np.concatenate(all_colors, axis=0) if all_colors else np.zeros((0, 3)) normal = self.get_normals(np.array([0, 0, 0]), coord) if self.if_img: imgs_path = data_info["images"] imgs = [Image.open(asset) for asset in imgs_path] img_width, img_height = imgs[0].size div_w = img_width // self.patch_w div_h = img_height // self.patch_h div_min = max(min(div_w, div_h), 1) crop_img_width = div_min * self.patch_w crop_img_height = div_min * self.patch_h left = int((img_width - crop_img_width) / 2) top = int((img_height - crop_img_height) / 2) right = int((img_width + crop_img_width) / 2) bottom = int((img_height + crop_img_height) / 2) imgs = [img.crop((left, top, right, bottom)) for img in imgs] imgs = [self.transform_img(img) for img in imgs] if len(imgs) > 0: imgs_list = torch.stack(imgs) imgs_tensor = imgs_list.float() else: imgs_tensor = torch.empty( ( 0, 3, self.patch_h * self.patch_size, self.patch_w * self.patch_size, ) ) img_num = np.array([imgs_tensor.shape[0]], dtype=np.int32) num_points, num_imgs = coord.shape[0], len(imgs) correspondence_infos = ( np.ones((num_points, num_imgs, 2), dtype=np.float32) * -1 ) for asset_id, correspondence_info in enumerate(correspondences): correspondence_info = self.resize_correspondence_info( correspondence_info, (self.patch_h * self.patch_size, self.patch_w * self.patch_size), (img_height, img_width), (left, top, right, bottom), self.patch_size, ) correspondence_infos[ correspondence_info[:, -1].astype(np.int32), asset_id, : ] = correspondence_info[:, :-1] data_dict = { "name": name, "split": split, "coord": coord.astype(np.float32), "color": color.astype(np.float32), "normal": normal.astype(np.float32), "images": imgs_tensor, "img_num": img_num, "correspondence": correspondence_infos, } else: data_dict = { "name": name, "split": split, "coord": coord.astype(np.float32), "color": color.astype(np.float32), "normal": normal.astype(np.float32), } return data_dict @DATASETS.register_module() class ConcatDataset(Dataset): def __init__(self, datasets, loop=1): super(ConcatDataset, self).__init__() self.datasets = [build_dataset(dataset) for dataset in datasets] self.loop = loop self.data_list = self.get_data_list() logger = get_root_logger() logger.info( "Totally {} x {} samples in the concat set.".format( len(self.data_list), self.loop ) ) def get_data_list(self): data_list = [] for i in range(len(self.datasets)): data_list.extend( zip( np.ones(len(self.datasets[i]), dtype=int) * i, np.arange(len(self.datasets[i])), ) ) return data_list def get_data(self, idx): dataset_idx, data_idx = self.data_list[idx % len(self.data_list)] return self.datasets[dataset_idx][data_idx] def get_data_name(self, idx): dataset_idx, data_idx = self.data_list[idx % len(self.data_list)] return self.datasets[dataset_idx].get_data_name(data_idx) def __getitem__(self, idx): return self.get_data(idx) def __len__(self): return len(self.data_list) * self.loop