File size: 2,711 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 | """
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
|