File size: 3,824 Bytes
7ab05dd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | """
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
|