| """ |
| PartNetE Datasets |
| |
| Author: Yujia Zhang (yujia.zhang.cs@gmail.com) |
| Please cite our work if the code is helpful to you. |
| """ |
|
|
| import os |
| import glob |
| from collections.abc import Sequence |
| import numpy as np |
| import json |
|
|
| from .defaults import DefaultDataset |
| from .builder import DATASETS |
|
|
| from pointcept.utils.cache import shared_dict |
|
|
|
|
| @DATASETS.register_module() |
| class PartNetEDataset(DefaultDataset): |
| def __init__( |
| self, class_names, num_parts, data_root, meta_path, split, *args, **kwargs |
| ): |
| self.categories = class_names |
| self.num_parts = num_parts |
| self.num_part_offset = np.concatenate(([0], np.cumsum(self.num_parts))) |
| with open(meta_path, "r", encoding="utf-8") as file: |
| meta_data = json.load(file) |
| self.category2part = {} |
| self.parts = [] |
| for class_id, class_name in enumerate(self.categories): |
| self.category2part[class_name] = self.num_part_offset[class_id] + list( |
| range(self.num_parts[class_id]) |
| ) |
| category_part_name = [ |
| class_name + "_" + part_name |
| for part_name in ["other"] + meta_data[class_name] |
| ] |
| self.parts.extend(category_part_name) |
| super().__init__(data_root=data_root, split=split, *args, **kwargs) |
|
|
| def get_data_list(self): |
| if isinstance(self.split, str): |
| data_list = glob.glob(os.path.join(self.data_root, self.split, "*/*")) |
| elif isinstance(self.split, Sequence): |
| data_list = [] |
| for split in self.split: |
| data_list += glob.glob(os.path.join(self.data_root, split, "*/*")) |
| else: |
| raise NotImplementedError |
| return data_list |
|
|
| def get_data_name(self, idx): |
| file_path = self.data_list[idx % len(self.data_list)] |
| dir_path, id_name = os.path.split(file_path) |
| object_name = os.path.basename(dir_path) |
| data_name = f"{object_name}_{id_name}" |
| return data_name |
|
|
| 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 |
|
|
| object_name = name.split("_")[0] |
| cls_token = self.categories.index(object_name) |
| data_dict["cls_token"] = cls_token |
|
|
| 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) |
| + self.num_part_offset[cls_token] |
| + 1 |
| ) |
| 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 |
|
|
