| """ |
| Objaverse Point Cloud Dataset Loader |
| |
| Load pre-processed .pt files. |
| Each .pt file is assumed to be a dictionary of PyTorch tensors. |
| """ |
|
|
| import os |
| import torch |
| import numpy as np |
| import open3d as o3d |
| from PIL import Image |
| import json |
| from collections.abc import Sequence |
| import random |
|
|
| from .builder import DATASETS |
| from .defaults import DefaultDataset, DefaultImagePointDataset |
|
|
|
|
| @DATASETS.register_module() |
| class Cap3DDataset(DefaultDataset): |
| """ |
| Cap3D Dataset for loading .pt files. |
| """ |
|
|
| def __init__( |
| self, |
| data_num=None, |
| **kwargs, |
| ): |
| super().__init__(**kwargs) |
| if data_num: |
| self.data_list = self.data_list[:data_num] |
|
|
| @staticmethod |
| def get_normals(center, coords): |
| Cs = np.repeat(center.reshape((1, -1)), coords.shape[0], axis=0) |
| view_dirs = coords - Cs |
| view_dirs = view_dirs / np.linalg.norm(view_dirs, axis=-1, keepdims=True) |
| pcd = o3d.geometry.PointCloud() |
| pcd.points = o3d.utility.Vector3dVector(coords) |
| pcd.estimate_normals() |
| normals = np.asarray(pcd.normals) |
| dot_product = np.sum(normals * view_dirs, axis=-1) |
| flip_mask = dot_product < 0 |
| normals[flip_mask] = -normals[flip_mask] |
|
|
| |
| normals = normals / np.linalg.norm(normals, axis=-1, keepdims=True) |
| return normals |
|
|
| def get_data(self, idx): |
| """ |
| The core data loading function. |
| |
| Args: |
| idx (int): The index of the data sample. |
| |
| Returns: |
| dict: A dictionary containing the loaded and processed point cloud data. |
| """ |
| data_path = self.data_list[idx % len(self.data_list)] |
| name = self.get_data_name(idx) |
| split = self.get_split_name(idx) |
|
|
| try: |
| data = torch.load(data_path, map_location="cpu") |
| except Exception as e: |
| raise IOError(f"Failed to load a .pt file at {data_path}. Error: {e}") |
|
|
| data = data.numpy().transpose() |
| color = data[:, 3:] |
| coord = data[:, :3] |
| normal = self.get_normals(np.mean(coord, axis=0), coord) |
|
|
| data_dict = {} |
|
|
| data_dict["coord"] = coord |
| data_dict["color"] = color |
| data_dict["normal"] = normal |
|
|
| data_dict["name"] = name |
| data_dict["split"] = split |
|
|
| return data_dict |
|
|
|
|
| @DATASETS.register_module() |
| class Cap3DImagePointDataset(DefaultImagePointDataset): |
| def __init__( |
| self, |
| data_num=-1, |
| **kwargs, |
| ): |
| self.data_num = data_num |
| super().__init__(**kwargs) |
|
|
| 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 |
|
|
| if self.data_num < 0 or self.data_num >= len(data_list): |
| selected_data_list = data_list |
| selected_split_list = split_list |
| return selected_data_list, selected_split_list |
|
|
| all_keys = list(data_list.keys()) |
| selected_keys = random.sample(all_keys, self.data_num) |
| selected_data_list = {key: data_list[key] for key in selected_keys} |
| valid_keys = set(selected_keys) |
|
|
| selected_split_list = {} |
| for split_name, key_list in split_list.items(): |
| filtered_keys = [key for key in key_list if key in valid_keys] |
| if filtered_keys: |
| selected_split_list[split_name] = filtered_keys |
| return selected_data_list, selected_split_list |
|
|
| @staticmethod |
| def get_normals(center, coords): |
| Cs = np.repeat(center.reshape((1, -1)), coords.shape[0], axis=0) |
| view_dirs = coords - Cs |
| view_dirs = view_dirs / np.linalg.norm(view_dirs, axis=-1, keepdims=True) |
| pcd = o3d.geometry.PointCloud() |
| pcd.points = o3d.utility.Vector3dVector(coords) |
| pcd.estimate_normals() |
| normals = np.asarray(pcd.normals) |
| dot_product = np.sum(normals * view_dirs, axis=-1) |
| flip_mask = dot_product < 0 |
| normals[flip_mask] = -normals[flip_mask] |
| normals = normals / np.linalg.norm(normals, axis=-1, keepdims=True) |
| return normals |
|
|
| def get_pc_data(self, data_path): |
| try: |
| data = torch.load(data_path, map_location="cpu") |
| except Exception as e: |
| raise IOError(f"Failed to load a .pt file at {data_path}. Error: {e}") |
|
|
| data = data.numpy().transpose() |
| color = data[:, 3:] |
| coord = data[:, :3] |
| normal = self.get_normals(np.mean(coord, axis=0), coord) |
|
|
| data_dict = {} |
|
|
| data_dict["coord"] = coord |
| data_dict["color"] = color |
| data_dict["normal"] = normal |
|
|
| return data_dict |
|
|
| 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"] |
| pc_dict = self.get_pc_data(pointclouds_path) |
| data_dict.update(pc_dict) |
| imgs_path = data_path["images"] |
| imgs = [Image.open(asset).convert("RGB") 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 |
|
|
| 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 |
|
|
