| """ |
| HK Remote Data Dataset |
| |
| Author: Yujia Zhang (yujia.zhang.cs@gmail.com) |
| Please cite our work if the code is helpful to you. |
| """ |
|
|
| import os |
| import numpy as np |
| import open3d as o3d |
| import torch |
|
|
| from .defaults import DefaultDataset |
| from .builder import DATASETS |
| from pointcept.utils.cache import shared_dict |
|
|
|
|
| @DATASETS.register_module() |
| class PartNetDataDataset(DefaultDataset): |
| def __init__(self, if_img=True, crop_h=630, crop_w=1120, patch_size=14, **kwargs): |
| 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.if_img = if_img |
| super().__init__(**kwargs) |
|
|
| def get_data(self, idx): |
| data_path = self.data_list[idx % len(self.data_list)] |
| pc_data_path = os.path.join( |
| data_path, "point_sample", "sample-points-all-pts-nor-rgba-10000.ply" |
| ) |
| 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 = {} |
| pc = o3d.io.read_point_cloud(pc_data_path) |
| data_dict["coord"] = np.asarray(pc.points) |
| data_dict["color"] = np.asarray(pc.colors) |
| data_dict["normal"] = np.asarray(pc.normals) |
| data_dict["name"] = name |
| data_dict["split"] = split |
|
|
| if self.if_img: |
| 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([0], dtype=np.int32) |
| data_dict["correspondence"] = np.ones( |
| (data_dict["coord"].shape[0], 0, 2), dtype=np.float32 |
| ) * (-1) |
|
|
| 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 |
|
|
