|
|
import copy |
|
|
import json |
|
|
import os |
|
|
|
|
|
import numpy as np |
|
|
from scipy.linalg import polar |
|
|
from scipy.spatial.transform import Rotation |
|
|
import torch |
|
|
from torch.utils.data import Dataset |
|
|
|
|
|
from .utils import exists |
|
|
from .utils.logger import print_log |
|
|
|
|
|
|
|
|
def create_dataset(cfg_dataset): |
|
|
kwargs = cfg_dataset |
|
|
name = kwargs.pop('name') |
|
|
dataset = get_dataset(name)(**kwargs) |
|
|
print_log(f"Dataset '{name}' init: kwargs={kwargs}, len={len(dataset)}") |
|
|
return dataset |
|
|
|
|
|
def get_dataset(name): |
|
|
return { |
|
|
'base': PrimitiveDataset, |
|
|
}[name] |
|
|
|
|
|
|
|
|
SHAPE_CODE = { |
|
|
'CubeBevel': 0, |
|
|
'SphereSharp': 1, |
|
|
'CylinderSharp': 2, |
|
|
} |
|
|
|
|
|
|
|
|
class PrimitiveDataset(Dataset): |
|
|
def __init__(self, |
|
|
pc_dir, |
|
|
bs_dir, |
|
|
max_length=144, |
|
|
range_scale=[0, 1], |
|
|
range_rotation=[-180, 180], |
|
|
range_translation=[-1, 1], |
|
|
rotation_type='euler', |
|
|
pc_format='pc', |
|
|
): |
|
|
self.data_filename = os.listdir(pc_dir) |
|
|
|
|
|
self.pc_dir = pc_dir |
|
|
self.max_length = max_length |
|
|
self.range_scale = range_scale |
|
|
self.range_rotation = range_rotation |
|
|
self.range_translation = range_translation |
|
|
self.rotation_type = rotation_type |
|
|
self.pc_format = pc_format |
|
|
|
|
|
with open(os.path.join(bs_dir, 'basic_shapes.json'), 'r', encoding='utf-8') as f: |
|
|
basic_shapes = json.load(f) |
|
|
|
|
|
self.typeid_map = { |
|
|
1101002001034001: 'CubeBevel', |
|
|
1101002001034010: 'SphereSharp', |
|
|
1101002001034002: 'CylinderSharp', |
|
|
} |
|
|
|
|
|
def __len__(self): |
|
|
return len(self.data_filename) |
|
|
|
|
|
def __getitem__(self, idx): |
|
|
pc_file = os.path.join(self.pc_dir, self.data_filename[idx]) |
|
|
pc = o3d.io.read_point_cloud(pc_file) |
|
|
|
|
|
model_data = {} |
|
|
|
|
|
points = torch.from_numpy(np.asarray(pc.points)).float() |
|
|
colors = torch.from_numpy(np.asarray(pc.colors)).float() |
|
|
normals = torch.from_numpy(np.asarray(pc.normals)).float() |
|
|
if self.pc_format == 'pc': |
|
|
model_data['pc'] = torch.concatenate([points, colors], dim=-1).T |
|
|
elif self.pc_format == 'pn': |
|
|
model_data['pc'] = torch.concatenate([points, normals], dim=-1) |
|
|
elif self.pc_format == 'pcn': |
|
|
model_data['pc'] = torch.concatenate([points, colors, normals], dim=-1) |
|
|
else: |
|
|
raise ValueError(f'invalid pc_format: {self.pc_format}') |
|
|
|
|
|
return model_data |
|
|
|
|
|
|
|
|
def get_typeid_shapename_mapping(shapenames, config_data): |
|
|
typeid_map = {} |
|
|
for info in config_data.values(): |
|
|
for shapename in shapenames: |
|
|
if shapename[3:-4] in info['bpPath']: |
|
|
typeid_map[info['typeId']] = shapename.split('_')[3] |
|
|
break |
|
|
return typeid_map |
|
|
|
|
|
|
|
|
def check_valid_range(data, value_range): |
|
|
lo, hi = value_range |
|
|
assert hi > lo |
|
|
return np.logical_and(data >= lo, hi <= hi).all() |
|
|
|
|
|
|
|
|
def quat_to_euler(quat, degree=True): |
|
|
return Rotation.from_quat(quat).as_euler('XYZ', degrees=degree) |
|
|
|
|
|
|
|
|
def quat_to_rotvec(quat, degree=True): |
|
|
return Rotation.from_quat(quat).as_rotvec(degrees=degree) |
|
|
|
|
|
|
|
|
def rotate_axis(euler): |
|
|
trans = np.eye(4, 4) |
|
|
trans[:3, :3] = Rotation.from_euler('xyz', euler).as_matrix() |
|
|
return trans |
|
|
|
|
|
|
|
|
def SRT_quat_to_matrix(scale, quat, translation): |
|
|
rotation_matrix = Rotation.from_quat(quat).as_matrix() |
|
|
transform_matrix = np.eye(4) |
|
|
transform_matrix[:3, :3] = rotation_matrix * scale |
|
|
transform_matrix[:3, 3] = translation |
|
|
return transform_matrix |
|
|
|
|
|
|
|
|
def matrix_to_SRT_quat2(transform_matrix): |
|
|
transform_matrix = np.array(transform_matrix) |
|
|
translation = transform_matrix[:3, 3] |
|
|
rotation_matrix, scale_matrix = polar(transform_matrix[:3,:3]) |
|
|
quat = Rotation.from_matrix(rotation_matrix).as_quat() |
|
|
scale = np.diag(scale_matrix) |
|
|
return scale, quat, translation |
|
|
|
|
|
|
|
|
def apply_transform_to_block(block, trans_aug): |
|
|
precision_loss = False |
|
|
trans = SRT_quat_to_matrix( |
|
|
block['data']['scale'], |
|
|
block['data']['rotation'], |
|
|
block['data']['location'] |
|
|
) |
|
|
|
|
|
trans = trans_aug @ trans |
|
|
scale, quat, translation = matrix_to_SRT_quat2(trans) |
|
|
|
|
|
trans_rec = SRT_quat_to_matrix(scale, quat, translation) |
|
|
if not np.allclose(trans, trans_rec, atol=1e-1): |
|
|
precision_loss = True |
|
|
return precision_loss, {} |
|
|
|
|
|
new_block = copy.deepcopy(block) |
|
|
new_block['data']['scale'] = scale.tolist() |
|
|
new_block['data']['rotation'] = quat.tolist() |
|
|
new_block['data']['location'] = translation.tolist() |
|
|
return precision_loss, new_block |
|
|
|
