def viz_obb(pc,
label,
mask,
angle_classes,
angle_residuals,
size_classes,
size_residuals,
name=''):
""" Visualize oriented bounding box ground truth
pc: (N,3)
label: (K,3) K == MAX_NUM_OBJ
mask: (K,)
angle_classes: (K,)
angle_residuals: (K,)
size_classes: (K,)
size_residuals: (K,3)
"""
oriented_boxes = []
K = label.shape[0]
for i in range(K):
if mask[i] == 0: continue
obb = np.zeros(7)
obb[0:3] = label[i, 0:3]
heading_angle = 0 # hard code to 0
box_size = DC.mean_size_arr[size_classes[i], :] + size_residuals[i, :]
obb[3:6] = box_size
obb[6] = -1 * heading_angle
print(obb)
oriented_boxes.append(obb)
pc_util.write_oriented_bbox(oriented_boxes, 'gt_obbs{}.ply'.format(name))
pc_util.write_ply(label[mask == 1, :], 'gt_centroids{}.ply'.format(name))
def viz_obb(pc, label, mask, angle_classes, angle_residuals, size_classes,
size_residuals):
""" Visualize oriented bounding box ground truth
pc: (N,3)
label: (K,3) K == MAX_NUM_OBJ
mask: (K,)
angle_classes: (K,)
angle_residuals: (K,)
size_classes: (K,)
size_residuals: (K,3)
"""
oriented_boxes = []
K = label.shape[0]
for i in range(K):
if mask[i] == 0: continue
obb = np.zeros(7)
obb[0:3] = label[i, 0:3]
heading_angle = DC.class2angle(angle_classes[i], angle_residuals[i])
box_size = DC.class2size(size_classes[i], size_residuals[i])
obb[3:6] = box_size
obb[6] = -1 * heading_angle
print(obb)
oriented_boxes.append(obb)
pc_util.write_oriented_bbox(oriented_boxes, 'gt_obbs.ply')
pc_util.write_ply(label[mask == 1, :], 'gt_centroids.ply')
def viz_votes(pc, point_votes, point_votes_mask, name=''):
""" Visualize point votes and point votes mask labels
pc: (N,3 or 6), point_votes: (N,9), point_votes_mask: (N,)
"""
inds = (point_votes_mask == 1)
pc_obj = pc[inds, 0:3]
pc_obj_voted1 = pc_obj + point_votes[inds, 0:3]
pc_util.write_ply(pc_obj, 'pc_obj{}.ply'.format(name))
pc_util.write_ply(pc_obj_voted1, 'pc_obj_voted1{}.ply'.format(name))
def virtual_scane_one_model(model_dir, worker_id):
print('Scanning ' + model_dir)
######### prepare some folder for tmp work #############
tmp_model_name = 'tmp' + str(worker_id) + '.ply'
TMP_DATA_PATH = './tmp' + str(worker_id)
TMP_PLY_POINTCLOUD_PATH = './tmp' + str(worker_id) + '.ply_output'
if not os.path.exists(TMP_DATA_PATH):
os.makedirs(TMP_DATA_PATH)
clean_dir(TMP_PLY_POINTCLOUD_PATH)
# generate camera parameters
cam_view_points, cam_target_points = generate_camera_view_target_points()
model_filename = os.path.join(model_dir, 'models/model_normalized.obj')
if not os.path.exists(model_filename):
print('File not found: %s' % (model_filename))
return
model_basename = os.path.basename(model_dir)
prev_clean_output_filename = os.path.join(PREV_OUTPUT_DATA_PATH,
model_basename + '_clean.ply')
if os.path.exists(prev_clean_output_filename):
print('Previously scanned, skip.', prev_clean_output_filename)
return
ply_tmp_name = os.path.join(TMP_DATA_PATH, tmp_model_name)
mesh_util.convert_obj2ply(model_filename,
ply_tmp_name,
recenter=True,
center_mode='box_center')
cmd_str = EXE_VIRTUAL_SCANNER + ' ' + ply_tmp_name + ' ' + CMD_POSTFIX.format(
','.join(str(e) for e in cam_view_points), ','.join(
str(e) for e in cam_target_points))
os.system(cmd_str)
# collect all scanned point clouds
all_xyz = []
pcd_files = glob.glob(TMP_PLY_POINTCLOUD_PATH + '/*.pcd')
for pf in pcd_files:
xyz = pc_util.read_pcd(pf)
all_xyz.extend(xyz)
all_points = np.array(all_xyz)
print('Collecte #points:', all_points.shape)
if all_points.shape[0] < 2048:
print('Failed to scan sufficient points! Move to next model.')
return
all_points = pc_util.remove_duplicated_points(all_points)
print('Total points after merge: %d' % (all_points.shape[0]))
clean_output_filename = os.path.join(OUTPUT_DATA_PATH,
model_basename + '_clean.ply')
pc_util.write_ply(all_points, clean_output_filename)
print('Save point cloud to ' + clean_output_filename)
return
def data_viz(data_dir, dump_dir=os.path.join(BASE_DIR, 'data_viz_dump')):
''' Examine and visualize ycbgrasp dataset. '''
ycb = ycb_object(data_dir)
idxs = np.array(range(0, len(ycb)))
if not os.path.exists(dump_dir):
os.mkdir(dump_dir)
for idx in range(len(ycb)):
if idx % 10:
continue
data_idx = idxs[idx]
print('data index: ', data_idx)
pc = ycb.get_pointcloud(data_idx)
pc = pc[:, 0:3]
pc = pc_util.random_sampling(pc, args.num_point)
pc_util.write_ply(pc, os.path.join(dump_dir, str(idx) + '_pc.ply'))
print('Complete!')
def crop_patch(self, save_root_path, use_dijkstra=True, scale_ratio=1, gpu_id=0):
if save_root_path[-1]=='/':
save_root_path = save_root_path[:-1]
if not os.path.exists(save_root_path):
os.makedirs(save_root_path)
print('Creating TF session...')
config = tf.ConfigProto()
config.gpu_options.visible_device_list=str(gpu_id)
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
self.sess = tf.Session(config=config)
seeds = self.get_idx(self.patch_num)
print(seeds)
self.sess.close()
i = -1
for seed in (seeds):
i = i+1
# patch_size = self.patch_size*np.random.randint(1,scale_ratio+1)
assert scale_ratio>=1.0
patch_size = int(self.patch_size*np.random.uniform(1.0, scale_ratio))
try:
if use_dijkstra:
idx = self.geodesic_knn(seed, patch_size)
else:
idx = self.bfs_knn(seed, patch_size)
except Exception as e:
exc_type, exc_obj, exc_tb = sys.exc_info()
fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1]
print(exc_type, fname, exc_tb.tb_lineno)
idxx = np.random.permutation(patch_size)[:self.patch_size]
idxx.sort()
idx = idx[idxx]
point = self.data[idx]
#print("patch:%d point:%d" % (i, patch_size))
output_filename = '%s/%s_%d.ply' % (save_root_path, self.name, i)
pc_util.write_ply(point, output_filename)
def viz_votes(pc, point_votes, point_votes_mask):
""" Visualize point votes and point votes mask labels
pc: (N,3 or 6), point_votes: (N,9), point_votes_mask: (N,)
"""
inds = (point_votes_mask == 1)
pc_obj = pc[inds, 0:3]
pc_obj_voted1 = pc_obj + point_votes[inds, 0:3]
pc_obj_voted2 = pc_obj + point_votes[inds, 3:6]
pc_obj_voted3 = pc_obj + point_votes[inds, 6:9]
pc_util.write_ply(pc_obj, 'pc_obj.ply')
pc_util.write_ply(pc_obj_voted1, 'pc_obj_voted1.ply')
pc_util.write_ply(pc_obj_voted2, 'pc_obj_voted2.ply')
pc_util.write_ply(pc_obj_voted3, 'pc_obj_voted3.ply')
def data_viz(data_dir, dump_dir=os.path.join(BASE_DIR, 'data_viz_dump')):
''' Examine and visualize ycbgrasp dataset. '''
ycb = ycb_object(data_dir)
idxs = np.array(range(0, len(ycb)))
#np.random.seed(0)
#np.random.shuffle(idxs)
if not os.path.exists(dump_dir):
os.mkdir(dump_dir)
for idx in range(len(ycb)):
if idx % 500:
continue
data_idx = idxs[idx]
print('data index: ', data_idx)
pc = ycb.get_pointcloud(data_idx)
objects = ycb.get_label_objects(data_idx)
oriented_boxes = []
for obj in objects:
object_pc, inds = ycbgrasp_utils.get_object_points(
pc, obj.classname)
pc_util.write_ply(
object_pc,
os.path.join(dump_dir,
str(idx) + '_' + obj.classname + '_pc.ply'))
if len(object_pc) > 300:
obb = np.zeros((7))
obb[0:3] = obj.centroid
obb[3:6] = np.array([obj.l, obj.w, obj.h]) * 2
obb[6] = obj.heading_angle
oriented_boxes.append(obb)
if len(oriented_boxes) > 0:
oriented_boxes = np.vstack(tuple(oriented_boxes))
pc_util.write_oriented_bbox(
oriented_boxes, os.path.join(dump_dir,
str(idx) + '_obbs.ply'))
pc = pc[:, 0:3]
pc_util.write_ply(pc, os.path.join(dump_dir, str(idx) + '_pc.ply'))
print('Complete!')
obb = np.zeros(7)
obb[0:3] = label[i, 0:3]
heading_angle = 0 # hard code to 0
box_size = DC.mean_size_arr[size_classes[i], :] + size_residuals[i, :]
obb[3:6] = box_size
obb[6] = -1 * heading_angle
print(obb)
oriented_boxes.append(obb)
pc_util.write_oriented_bbox(oriented_boxes, 'gt_obbs{}.ply'.format(name))
pc_util.write_ply(label[mask == 1, :], 'gt_centroids{}.ply'.format(name))
if __name__ == '__main__':
dset = ScannetDetectionDataset(use_height=True, num_points=40000)
for i_example in range(4):
example = dset.__getitem__(1)
pc_util.write_ply(example['point_clouds'],
'pc_{}.ply'.format(i_example))
viz_votes(example['point_clouds'],
example['vote_label'],
example['vote_label_mask'],
name=i_example)
viz_obb(pc=example['point_clouds'],
label=example['center_label'],
mask=example['box_label_mask'],
angle_classes=None,
angle_residuals=None,
size_classes=example['size_class_label'],
size_residuals=example['size_residual_label'],
name=i_example)
def data_viz(data_dir, dump_dir=os.path.join(BASE_DIR, 'data_viz_dump')):
''' Examine and visualize SUN RGB-D data. '''
sunrgbd = sunrgbd_object(data_dir)
idxs = np.array(range(1, len(sunrgbd) + 1))
np.random.seed(0)
np.random.shuffle(idxs)
for idx in range(len(sunrgbd)):
data_idx = idxs[idx]
print('-' * 10, 'data index: ', data_idx)
pc = sunrgbd.get_depth(data_idx)
print('Point cloud shape:', pc.shape)
# Project points to image
calib = sunrgbd.get_calibration(data_idx)
uv, d = calib.project_upright_depth_to_image(pc[:, 0:3])
print('Point UV:', uv)
print('Point depth:', d)
import matplotlib.pyplot as plt
cmap = plt.cm.get_cmap('hsv', 256)
cmap = np.array([cmap(i) for i in range(256)])[:, :3] * 255
img = sunrgbd.get_image(data_idx)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
for i in range(uv.shape[0]):
depth = d[i]
color = cmap[int(120.0 / depth), :]
cv2.circle(img, (int(np.round(uv[i, 0])), int(np.round(uv[i, 1]))),
2,
color=tuple(color),
thickness=-1)
if not os.path.exists(dump_dir):
os.mkdir(dump_dir)
Image.fromarray(img).save(os.path.join(dump_dir, 'img_depth.jpg'))
# Load box labels
objects = sunrgbd.get_label_objects(data_idx)
print('Objects:', objects)
# Draw 2D boxes on image
img = sunrgbd.get_image(data_idx)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
for i, obj in enumerate(objects):
cv2.rectangle(img, (int(obj.xmin), int(obj.ymin)),
(int(obj.xmax), int(obj.ymax)), (0, 255, 0), 2)
cv2.putText(img, '%d %s' % (i, obj.classname),
(max(int(obj.xmin), 15), max(int(obj.ymin), 15)),
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 0, 0), 2)
Image.fromarray(img).save(os.path.join(dump_dir, 'img_box2d.jpg'))
# Dump OBJ files for the colored point cloud
for num_point in [10000, 20000, 40000, 80000]:
sampled_pcrgb = pc_util.random_sampling(pc, num_point)
pc_util.write_ply_rgb(
sampled_pcrgb[:, 0:3],
(sampled_pcrgb[:, 3:] * 256).astype(np.int8),
os.path.join(dump_dir, 'pcrgb_%dk.obj' % (num_point // 1000)))
# Dump OBJ files for 3D bounding boxes
# l,w,h correspond to dx,dy,dz
# heading angle is from +X rotating towards -Y
# (+X is degree, -Y is 90 degrees)
oriented_boxes = []
for obj in objects:
obb = np.zeros((7))
obb[0:3] = obj.centroid
# Some conversion to map with default setting of w,l,h
# and angle in box dumping
obb[3:6] = np.array([obj.l, obj.w, obj.h]) * 2
obb[6] = -1 * obj.heading_angle
print('Object cls, heading, l, w, h:',\
obj.classname, obj.heading_angle, obj.l, obj.w, obj.h)
oriented_boxes.append(obb)
if len(oriented_boxes) > 0:
oriented_boxes = np.vstack(tuple(oriented_boxes))
pc_util.write_oriented_bbox(oriented_boxes,
os.path.join(dump_dir, 'obbs.ply'))
else:
print('-' * 30)
continue
# Draw 3D boxes on depth points
box3d = []
ori3d = []
for obj in objects:
corners_3d_image, corners_3d = sunrgbd_utils.compute_box_3d(
obj, calib)
ori_3d_image, ori_3d = sunrgbd_utils.compute_orientation_3d(
obj, calib)
print('Corners 3D: ', corners_3d)
box3d.append(corners_3d)
ori3d.append(ori_3d)
pc_box3d = np.concatenate(box3d, 0)
pc_ori3d = np.concatenate(ori3d, 0)
print(pc_box3d.shape)
print(pc_ori3d.shape)
pc_util.write_ply(pc_box3d, os.path.join(dump_dir,
'box3d_corners.ply'))
pc_util.write_ply(pc_ori3d, os.path.join(dump_dir, 'box3d_ori.ply'))
print('-' * 30)
print('Point clouds and bounding boxes saved to PLY files under %s' %
(dump_dir))
print('Type anything to continue to the next sample...')
input()
def scannet_projection():
mask = torch.cuda.LongTensor(1 * column_height)
depth_images = torch.cuda.FloatTensor(1, proj_image_dims[1], proj_image_dims[0])
color_images = torch.cuda.FloatTensor(1, 3, input_image_dims[1], input_image_dims[0])
camera_poses = torch.cuda.FloatTensor(1, 4, 4)
label_images = torch.cuda.LongTensor(1, proj_image_dims[1], proj_image_dims[0])
# load_images
data_path = os.path.join(BASE_DIR, 'test_data')
load_frames_multi(data_path, np.array([20]), depth_images, color_images, camera_poses, color_mean, color_std)
# Load alignments
lines = open(BASE_DIR + '/scannet/scans/scene0001_00/scene0001_00.txt').readlines()
for line in lines:
if 'axisAlignment' in line:
axis_align_matrix = [float(x) \
for x in line.rstrip().strip('axisAlignment = ').split(' ')]
break
axis_align_matrix = np.array(axis_align_matrix).reshape((4, 4))
# TODO: for test purpose, we only need one loop now
for d, c in zip(depth_images, camera_poses):
# boundingmin, boundingmax, world_to_camera = projection.compute_projection(d, c, axis_align_matrix)
boundingmin, boundingmax, world_to_camera = projection.compute_projection(d, c)
boundingmin = boundingmin.cpu().numpy()
boundingmax = boundingmax.cpu().numpy()
break
mesh_vertices = read_mesh_vertices(filename=BASE_DIR + '/scannet/scans/scene0001_00/scene0001_00_vh_clean_2.ply')
# add addition ones to the end or each position in mesh_vertices
mesh_vertices = np.append(mesh_vertices, np.ones((mesh_vertices.shape[0], 1)), axis=1)
# TODO: this filter is problematic
# filter out the vertices that are not in the frustum (here treated as box-shape)
filter1 = mesh_vertices[:, 0]
filter1 = (filter1 >= boundingmin[0]) & (filter1 <= boundingmax[0])
filter2 = mesh_vertices[:, 1]
filter2 = (filter2 >= boundingmin[1]) & (filter2 <= boundingmax[1])
filter3 = mesh_vertices[:, 2]
filter3 = (filter3 >= boundingmin[2]) & (filter3 <= boundingmax[2])
filter_all = filter1 & filter2 & filter3
valid_vertices = mesh_vertices[filter_all]
# transform to current frame
world_to_camera = world_to_camera.cpu().numpy()
N = valid_vertices.shape[0]
valid_vertices_T = np.transpose(valid_vertices)
pcamera = np.matmul(world_to_camera, valid_vertices_T) # (4,4) x (4,N) => (4,N)
# Alternative way to compute p
# intrinsic_np = intrinsic.cpu().numpy()
# pimage = np.matmul(intrinsic_np, pcamera) # (4,4) x (4,N) => (4,N)
# ppixel = 1 / pimage[2] * np.array(pimage[0:2, :]) # shape: (2,N)
# p = np.concatenate((ppixel, pimage[2:, :]), axis=0)
# p = np.transpose(p)
p = np.transpose(pcamera) # shape: (N,4)
# project into image
p[:, 0] = (p[:, 0] * projection.intrinsic[0][0].cpu().numpy()) / p[:, 2] + projection.intrinsic[0][2].cpu().numpy()
p[:, 1] = (p[:, 1] * projection.intrinsic[1][1].cpu().numpy()) / p[:, 2] + projection.intrinsic[1][2].cpu().numpy()
pi = np.rint(p)
x = pi[:, 0]
x_filter = (x >= 0) & (x < proj_image_dims[0])
y = pi[:, 1]
y_filter = (y >= 0) & (y < proj_image_dims[1])
filterxy = x_filter & y_filter
pi = pi[filterxy]
p = p[filterxy]
reconstructed_depth_map = np.zeros((proj_image_dims[1], proj_image_dims[0]))
# for i1 in range(320):
# for i2 in range(240):
# x = pi[:, 0]
# x_filter = x == i1
# y = pi[:, 1]
# y_filter = y == i2
#
# correct_depth = depth_map_to_compare[i2][i1]
#
# dpth = p[:, 2]
# dpth_filter = (dpth <= correct_depth+0.5) & (dpth >= correct_depth-0.5)
#
# filterxyd = x_filter & y_filter & dpth_filter
# p_temp = p[filterxyd]
# if p_temp.shape[0] == 0:
# reconstructed_depth_map[i2][i1] = 0.0
# else:
# reconstructed_depth_map[i2][i1] = p_temp[0,2]
# print(i1,i2)
#
p_combined = np.concatenate((pi[:, 0:2], p[:, 2:3]), axis=1)
# find correspondence in a 320 x 240 image, and fill in the depth value:
# reconstructed_depth_map = np.zeros((500, 500))
# TODO: compare with depth map
for p in p_combined:
reconstructed_depth_map[int(p[1]), int(p[0])] = p[2]
num_non_zeros = np.count_nonzero(reconstructed_depth_map)
imageio.imwrite(BASE_DIR + '/reconstructed_depth_map.png', reconstructed_depth_map)
# output point cloud from depth_image
depth_map_to_compare = depth_images[0].cpu().numpy()
pointstowrite = np.ones((320 * 240, 4))
colors = np.ones((320 * 240, 4))
# TODO: convert this to matrix multiplication
print("back-projection, depth-map -> 3d points")
for i1 in tqdm(range(proj_image_dims[0])):
for i2 in range(proj_image_dims[1]):
pcamera = projection.depth_to_skeleton(i1, i2, depth_map_to_compare[i2, i1]).unsqueeze(1).cpu().numpy()
pcamera = np.append(pcamera, np.ones((1, 1)), axis=0)
camera2world = camera_poses[0].cpu().numpy()
world = np.matmul(camera2world, pcamera)
world = world.reshape((1, 4))
pointstowrite[i1 * i2, :] = world[0, :]
pointstowrite = pointstowrite[:, 0:3]
pc_util.write_ply(pointstowrite, BASE_DIR + '/scannet/testobject.ply')
para_config_gan['3D-EPN_train_point_cloud_dir'] = '/workspace/pointnet2/pc2pc/data/3D-EPN_dataset/shapenet_dim32_sdf_pc/02958343/point_cloud'
para_config_gan['3D-EPN_test_point_cloud_dir'] = '/workspace/pointnet2/pc2pc/data/3D-EPN_dataset/test-images_dim32_sdf_pc/02958343/point_cloud'
NOISY_TRAIN_DATASET = ShapeNet_3DEPN_PointsDataset(para_config_gan['3D-EPN_train_point_cloud_dir'], batch_size=12, npoint=2048, shuffle=True, split='trainval', preprocess=False)
'''
REALDATASET = RealWorldPointsDataset(
'/workspace/pointnet2/pc2pc/data/scannet_v2_chairs_alilgned_v2/point_cloud',
batch_size=6,
npoint=2048,
shuffle=False,
split='trainval',
random_seed=0)
data_batch = REALDATASET.next_batch()
pc_util.write_ply(data_batch[0], 'real_data.ply')
pc_util.write_ply(data_batch[1], 'real_data_1.ply')
pc_util.write_ply(data_batch[2], 'real_data_2.ply')
'''
TRAIN_DATASET = ShapeNetPartPointsDataset('/workspace/pointnet2/pc2pc/data/ShapeNet_v2_point_cloud/03001627/point_cloud_clean', batch_size=6, npoint=2048, shuffle=False, split='test', random_seed=0)
data_batch = TRAIN_DATASET.next_batch_noise_partial_by_percentage()
pc_util.write_ply(data_batch[0], 'data.ply')
'''
#data_batch = TRAIN_DATASET.aug_data_batch(data_batch, rot=False, trans=0.1, snap2ground=False)
#pc_util.write_ply(data_batch[0], 'data_aug.ply')
'''
noisy_partial_data_batch = TRAIN_DATASET.next_batch_noise_added_with_partial(partial_portion=1)
pc_util.write_ply(noisy_partial_data_batch[0], 'test_noisy_partial.ply')
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = os.path.dirname(BASE_DIR)
sys.path.append(os.path.join(ROOT_DIR, '../utils'))
import pc_util
dataset_dir = '/workspace/pointnet2/pc2pc/data/3D-EPN_dataset/test-images_dim32_sdf_pc'
output_dir = '/workspace/pointnet2/pc2pc/data/3D-EPN_dataset/test-images_dim32_sdf_pc_processed'
if not os.path.exists(output_dir):
os.makedirs(output_dir)
cls_ids = os.listdir(dataset_dir)
for cls_id in cls_ids:
cls_dir = os.path.join(dataset_dir, cls_id)
output_cls_dir = os.path.join(output_dir, cls_id, 'point_cloud')
if not os.path.exists(output_cls_dir):
os.makedirs(output_cls_dir)
point_cloud_names = os.listdir(cls_dir)
for pc_n in tqdm(point_cloud_names):
pc_filename = os.path.join(cls_dir, pc_n)
out_pc_filename = os.path.join(output_cls_dir, pc_n)
points = pc_util.read_ply_xyz(pc_filename)
rotated_points = pc_util.rotate_point_cloud_by_axis_angle(
points, [0, 1, 0], 90)
pc_util.write_ply(rotated_points, out_pc_filename)
ret_dict['point_boundary_mask_z'] = point_boundary_mask_z.astype(np.float32)
ret_dict['point_boundary_mask_xy'] = point_boundary_mask_xy.astype(np.float32)
ret_dict['point_boundary_offset_z'] = point_boundary_offset_z.astype(np.float32)
ret_dict['point_boundary_offset_xy'] = point_boundary_offset_xy.astype(np.float32)
ret_dict['point_boundary_sem_z'] = point_boundary_sem_z.astype(np.float32)
ret_dict['point_boundary_sem_xy'] = point_boundary_sem_xy.astype(np.float32)
ret_dict['point_line_mask'] = point_line_mask.astype(np.float32)
ret_dict['point_line_offset'] = point_line_offset.astype(np.float32)
ret_dict['point_line_sem'] = point_line_sem.astype(np.float32)
ret_dict['vote_label'] = point_votes.astype(np.float32)
ret_dict['vote_label_mask'] = point_votes_mask.astype(np.int64)
ret_dict['scan_idx'] = np.array(idx).astype(np.int64)
ret_dict['pcl_color'] = pcl_color
ret_dict['num_instance'] = num_instance
return ret_dict
if __name__=='__main__':
dset = ScannetDetectionDataset(split_set='train', use_height=True, num_points=50000, augment=False, use_angle=False)
for i_example in range(len(dset.scan_names)):
example = dset.__getitem__(i_example)
pc_util.write_ply(example['point_clouds'], 'pc_{}.ply'.format(i_example))
pc_util.write_ply(example['point_clouds'][example['point_line_mask']==1,0:3], 'pc_obj_line{}.ply'.format(i_example))
pc_util.write_ply(example['point_clouds'][example['point_boundary_mask_z']==1,0:3], 'pc_obj_boundary_z{}.ply'.format(i_example))
pc_util.write_ply(example['point_clouds'][example['point_boundary_mask_xy']==1,0:3], 'pc_obj_boundary_xy{}.ply'.format(i_example))
print (i_example)
def dump_results(end_points, dump_dir, config, inference_switch=False):
if not os.path.exists(dump_dir):
os.system('mkdir %s' % (dump_dir))
# INPUT
point_clouds = end_points['point_clouds'].cpu().numpy()
batch_size = point_clouds.shape[0]
# NETWORK OUTPUTS
seed_xyz = end_points['seed_xyz'].detach().cpu().numpy() # (B,num_seed,3)
if 'vote_object_xyz' in end_points:
aggregated_vote_object_xyz = end_points[
'aggregated_vote_object_xyz'].detach().cpu().numpy()
vote_object_xyz = end_points['vote_object_xyz'].detach().cpu().numpy(
) # (B,num_seed,3)
aggregated_vote_object_xyz = end_points[
'aggregated_vote_object_xyz'].detach().cpu().numpy()
objectness_scores = end_points['objectness_scores'].detach().cpu().numpy(
) # (B,K,2)
pred_center = end_points['center'].detach().cpu().numpy() # (B,K,3)
pred_rot = end_points['rot_6d'] # (B,K,3)
predict_rot = pred_rot.reshape(pred_rot.shape[0] * pred_rot.shape[1],
pred_rot.shape[2])
predict_rmat = tools.compute_rotation_matrix_from_ortho6d(predict_rot)
predict_rmat = predict_rmat.reshape(pred_rot.shape[0], pred_rot.shape[1],
9)
# OTHERS
#pred_mask = end_points['pred_mask'] # B,num_proposal
idx_beg = 0
save_pose_file = os.path.join(dump_dir, 'pred_poses.txt')
f = open(save_pose_file, "w")
for i in range(batch_size):
pc = point_clouds[i, :, :]
objectness_prob = softmax(objectness_scores[i, :, :])[:, 1] # (K,)
# Dump various point clouds
pc_util.write_ply(
pc, os.path.join(dump_dir, '%06d_pc.ply' % (idx_beg + i)))
pc_util.write_ply(
seed_xyz[i, :, :],
os.path.join(dump_dir, '%06d_seed_pc.ply' % (idx_beg + i)))
if 'vote_object_xyz' in end_points:
pc_util.write_ply(
end_points['vote_object_xyz'][i, :, :],
os.path.join(dump_dir, '%06d_vgen_pc.ply' % (idx_beg + i)))
pc_util.write_ply(
aggregated_vote_object_xyz[i, :, :],
os.path.join(dump_dir,
'%06d_aggregated_vote_pc.ply' % (idx_beg + i)))
pc_util.write_ply(
aggregated_vote_object_xyz[i, :, :],
os.path.join(dump_dir,
'%06d_aggregated_vote_pc.ply' % (idx_beg + i)))
if np.sum(objectness_prob > DUMP_CONF_THRESH) > 0:
pc_util.write_ply(
pred_center[i, objectness_prob > DUMP_CONF_THRESH, 0:3],
os.path.join(dump_dir,
'%06d_confident_proposal_pc.ply' % (idx_beg + i)))
# Dump predicted poses
if np.sum(objectness_prob > DUMP_CONF_THRESH) > 0:
num_proposal = pred_center.shape[1]
for j in range(num_proposal):
loc = pred_center[i, j, 0:3]
for ite in loc:
str_num = '{:.6f}'.format(ite)
f.write(str_num)
f.write(' ')
f.write("\n")
for j in range(num_proposal):
rmat = predict_rmat[i, j, 0:9]
for ind, ite in enumerate(rmat):
str_num = '{:.6f}'.format(ite)
f.write(str_num)
if (ind + 1) % 3 == 0:
f.write('\n')
else:
f.write(' ')
f.close()
df_mesh.apply_translation(trans_v)
df_mesh.apply_scale(scale_factor)
# rotate to make the face -z, from +z to -z
rot_m = axangle2aff([0, 1, 0], np.pi)
df_mesh.apply_transform(rot_m)
return df_mesh
volumn_name_list = os.listdir(EPN_test_result_dir)
for vname in volumn_name_list:
if not vname.endswith('.npy'): continue
v_filename = os.path.join(EPN_test_result_dir, vname)
v = np.load(v_filename)
vertices, triangles = mcubes.marching_cubes(v, isoval)
# isosuface from distance field
mesh = trimesh.Trimesh(vertices=vertices, faces=triangles)
mesh = align_dfmesh_scanpc(mesh, v.shape[0])
# samples from isosurface
recon_pc, _ = trimesh.sample.sample_surface(mesh, 2048)
recon_pc = np.array(recon_pc)
output_pcloud_re_dir = os.path.join(output_dir, 'pcloud', 'reconstruction')
if not os.path.exists(output_pcloud_re_dir):
os.makedirs(output_pcloud_re_dir)
re_output_filename = os.path.join(output_pcloud_re_dir,
vname.split('.')[0] + '.ply')
pc_util.write_ply(recon_pc, re_output_filename)
def dump_results(end_points, dump_dir, config, inference_switch=False):
if not os.path.exists(dump_dir):
os.system('mkdir %s' % (dump_dir))
# INPUT
point_clouds = end_points['point_clouds'].cpu().numpy()
batch_size = point_clouds.shape[0]
# NETWORK OUTPUTS
seed_xyz = end_points['seed_xyz'].detach().cpu().numpy() # (B,num_seed,3)
if 'vote_xyz' in end_points:
aggregated_vote_xyz = end_points['aggregated_vote_xyz'].detach().cpu(
).numpy()
vote_xyz = end_points['vote_xyz'].detach().cpu().numpy(
) # (B,num_seed,3)
aggregated_vote_xyz = end_points['aggregated_vote_xyz'].detach().cpu(
).numpy()
objectness_scores = end_points['objectness_scores'].detach().cpu().numpy(
) # (B,K,2)
pred_center = end_points['center'].detach().cpu().numpy() # (B,K,3)
pred_angle_class = torch.argmax(end_points['angle_scores'],
-1) # B,num_proposal
pred_angle_residual = torch.gather(
end_points['angle_residuals'], 2,
pred_angle_class.unsqueeze(-1)) # B,num_proposal,1
pred_angle_class = pred_angle_class.detach().cpu().numpy(
) # B,num_proposal
pred_angle_residual = pred_angle_residual.squeeze(
2).detach().cpu().numpy() # B,num_proposal
pred_viewpoint_class = torch.argmax(end_points['viewpoint_scores'],
-1) # B,num_proposal
pred_sem_cls = torch.argmax(end_points['sem_cls_scores'],
-1) # B,num_proposal
pred_quality = end_points['quality'].detach().cpu().numpy(
) # B,num_proposal
pred_width = end_points['width'].detach().cpu().numpy() # B,num_proposal
# OTHERS
#pred_mask = end_points['pred_mask'] # B,num_proposal
idx_beg = 0
save_grasp_file = os.path.join(dump_dir, 'pred_grasps.txt')
f = open(save_grasp_file, "w")
f.write("object x y z viewpoint angle quality width\n")
for i in range(batch_size):
pc = point_clouds[i, :, :]
objectness_prob = softmax(objectness_scores[i, :, :])[:, 1] # (K,)
# Dump various point clouds
pc_util.write_ply(
pc, os.path.join(dump_dir, '%06d_pc.ply' % (idx_beg + i)))
pc_util.write_ply(
seed_xyz[i, :, :],
os.path.join(dump_dir, '%06d_seed_pc.ply' % (idx_beg + i)))
if 'vote_xyz' in end_points:
pc_util.write_ply(
end_points['vote_xyz'][i, :, :],
os.path.join(dump_dir, '%06d_vgen_pc.ply' % (idx_beg + i)))
pc_util.write_ply(
aggregated_vote_xyz[i, :, :],
os.path.join(dump_dir,
'%06d_aggregated_vote_pc.ply' % (idx_beg + i)))
pc_util.write_ply(
aggregated_vote_xyz[i, :, :],
os.path.join(dump_dir,
'%06d_aggregated_vote_pc.ply' % (idx_beg + i)))
if np.sum(objectness_prob > DUMP_CONF_THRESH) > 0:
pc_util.write_ply(
pred_center[i, objectness_prob > DUMP_CONF_THRESH, 0:3],
os.path.join(dump_dir,
'%06d_confident_proposal_pc.ply' % (idx_beg + i)))
# Dump predicted grasps
if np.sum(objectness_prob > DUMP_CONF_THRESH) > 0:
num_proposal = pred_center.shape[1]
for j in range(num_proposal):
grasp = config.param2grasp(pred_sem_cls[i,
j], pred_center[i, j,
0:3],
pred_viewpoint_class[i, j],
pred_angle_class[i, j],
pred_angle_residual[i, j],
pred_quality[i, j,
0], pred_width[i, j,
0])
f.write(grasp[0])
f.write(' ')
for ite in grasp[1:]:
str_num = '{:.6f}'.format(ite)
f.write(str_num)
f.write(' ')
f.write("\n")
f.close()
def project_point_clound_to_image(point_cloud, calib_Rtilt, calib_K):
calib = sunrgbd_utils.SUNRGBD_Calibration_from_array(calib_Rtilt, calib_K)
uv, d = calib.project_upright_depth_to_image(point_cloud[:, 0:3])
# pc2 = np.copy(point_cloud)
# pc2 = np.dot(np.transpose(calib_Rtilt), np.transpose(pc2))
# pc2 = np.transpose(pc2)
# pc2[:,[0,1,2]] = pc2[:,[0,2,1]] # cam X,Y,Z = depth X,-Z,Y
# pc2[:,1] *= -1
# uv = np.dot(pc2, np.transpose(calib_Rtilt)) # (n,3)
# uv[:,0] /= uv[:,2]
# uv[:,1] /= uv[:,2]
return uv
if __name__ == '__main__':
d = SunrgbdDetectionVotesDataset(use_height=True,
use_color=True,
use_v1=True,
augment=True)
sample = d[200]
print(sample['vote_label'].shape, sample['vote_label_mask'].shape)
pc_util.write_ply(sample['point_clouds'], 'pc.ply')
viz_votes(sample['point_clouds'], sample['vote_label'],
sample['vote_label_mask'])
viz_obb(sample['point_clouds'], sample['center_label'],
sample['box_label_mask'], sample['heading_class_label'],
sample['heading_residual_label'], sample['size_class_label'],
sample['size_residual_label'])
output_dir = '../data/ShapeNet_v1_point_cloud/03636649/point_cloud_clean_full_ds'
if not os.path.exists(output_dir):
os.makedirs(output_dir)
down_sample_rate = 0.125
ply_filename_list = [
os.path.join(point_cloud_dir, f) for f in os.listdir(point_cloud_dir)
]
ply_filename_list.sort()
for pf in tqdm(ply_filename_list):
points = pc_util.read_ply_xyz(pf)
choice = np.random.choice(points.shape[0],
int(points.shape[0] * down_sample_rate))
sampled_points = points[choice]
if sampled_points.shape[0] < 1000:
print('Skip, probably empty scan. %s' % (pf))
continue
# ensure that the bbox is centerred at the original
pts_min = np.amin(sampled_points, axis=0, keepdims=True)
pts_max = np.amax(sampled_points, axis=0, keepdims=True)
bbox_center = (pts_min + pts_max) / 2.0
sampled_points = sampled_points - bbox_center
output_filename = os.path.join(output_dir, os.path.basename(pf))
pc_util.write_ply(sampled_points, output_filename)
def dump_results(end_points, dump_dir, config, inference_switch=False):
''' Dump results.
Args:
end_points: dict
{..., pred_mask}
pred_mask is a binary mask array of size (batch_size, num_proposal) computed by running NMS and empty box removal
Returns:
None
'''
if not os.path.exists(dump_dir):
os.system('mkdir %s'%(dump_dir))
# INPUT
point_clouds = end_points['point_clouds'].cpu().numpy()
batch_size = point_clouds.shape[0]
# NETWORK OUTPUTS
seed_xyz = end_points['seed_xyz'].detach().cpu().numpy() # (B,num_seed,3)
if 'vote_xyz' in end_points:
aggregated_vote_xyz = end_points['aggregated_vote_xyz'].detach().cpu().numpy()
vote_xyz = end_points['vote_xyz'].detach().cpu().numpy() # (B,num_seed,3)
aggregated_vote_xyz = end_points['aggregated_vote_xyz'].detach().cpu().numpy()
objectness_scores = end_points['objectness_scores'].detach().cpu().numpy() # (B,K,2)
pred_center = end_points['center'].detach().cpu().numpy() # (B,K,3)
pred_heading_class = torch.argmax(end_points['heading_scores'], -1) # B,num_proposal
pred_heading_residual = torch.gather(end_points['heading_residuals'], 2, pred_heading_class.unsqueeze(-1)) # B,num_proposal,1
pred_heading_class = pred_heading_class.detach().cpu().numpy() # B,num_proposal
pred_heading_residual = pred_heading_residual.squeeze(2).detach().cpu().numpy() # B,num_proposal
pred_size_class = torch.argmax(end_points['size_scores'], -1) # B,num_proposal
pred_size_residual = torch.gather(end_points['size_residuals'], 2, pred_size_class.unsqueeze(-1).unsqueeze(-1).repeat(1,1,1,3)) # B,num_proposal,1,3
pred_size_residual = pred_size_residual.squeeze(2).detach().cpu().numpy() # B,num_proposal,3
# OTHERS
pred_mask = end_points['pred_mask'] # B,num_proposal
idx_beg = 0
for i in range(batch_size):
pc = point_clouds[i,:,:]
objectness_prob = softmax(objectness_scores[i,:,:])[:,1] # (K,)
# Dump various point clouds
pc_util.write_ply(pc, os.path.join(dump_dir, '%06d_pc.ply'%(idx_beg+i)))
pc_util.write_ply(seed_xyz[i,:,:], os.path.join(dump_dir, '%06d_seed_pc.ply'%(idx_beg+i)))
#pc_util.write_ply_rgb(pc[:,0:3], (pc[:,3:]*256).astype(np.int8),os.path.join(dump_dir, '%06d_rgb_pc.ply'%(idx_beg+i))) #jason color
if 'vote_xyz' in end_points:
pc_util.write_ply(end_points['vote_xyz'][i,:,:], os.path.join(dump_dir, '%06d_vgen_pc.ply'%(idx_beg+i)))
pc_util.write_ply(aggregated_vote_xyz[i,:,:], os.path.join(dump_dir, '%06d_aggregated_vote_pc.ply'%(idx_beg+i)))
pc_util.write_ply(aggregated_vote_xyz[i,:,:], os.path.join(dump_dir, '%06d_aggregated_vote_pc.ply'%(idx_beg+i)))
pc_util.write_ply(pred_center[i,:,0:3], os.path.join(dump_dir, '%06d_proposal_pc.ply'%(idx_beg+i)))
if np.sum(objectness_prob>DUMP_CONF_THRESH)>0:
pc_util.write_ply(pred_center[i,objectness_prob>DUMP_CONF_THRESH,0:3], os.path.join(dump_dir, '%06d_confident_proposal_pc.ply'%(idx_beg+i)))
# Dump predicted bounding boxes
if np.sum(objectness_prob>DUMP_CONF_THRESH)>0:
num_proposal = pred_center.shape[1]
obbs = []
for j in range(num_proposal):
obb = config.param2obb(pred_center[i,j,0:3], pred_heading_class[i,j], pred_heading_residual[i,j],
pred_size_class[i,j], pred_size_residual[i,j])
obbs.append(obb)
if len(obbs)>0:
obbs = np.vstack(tuple(obbs)) # (num_proposal, 7)
pc_util.write_oriented_bbox(obbs[objectness_prob>DUMP_CONF_THRESH,:], os.path.join(dump_dir, '%06d_pred_confident_bbox.ply'%(idx_beg+i)))
pc_util.write_oriented_bbox(obbs[np.logical_and(objectness_prob>DUMP_CONF_THRESH, pred_mask[i,:]==1),:], os.path.join(dump_dir, '%06d_pred_confident_nms_bbox.ply'%(idx_beg+i)))
pc_util.write_oriented_bbox(obbs[pred_mask[i,:]==1,:], os.path.join(dump_dir, '%06d_pred_nms_bbox.ply'%(idx_beg+i)))
pc_util.write_oriented_bbox(obbs, os.path.join(dump_dir, '%06d_pred_bbox.ply'%(idx_beg+i)))
# Return if it is at inference time. No dumping of groundtruths
if inference_switch:
return
# LABELS
gt_center = end_points['center_label'].cpu().numpy() # (B,MAX_NUM_OBJ,3)
gt_mask = end_points['box_label_mask'].cpu().numpy() # B,K2
gt_heading_class = end_points['heading_class_label'].cpu().numpy() # B,K2
gt_heading_residual = end_points['heading_residual_label'].cpu().numpy() # B,K2
gt_size_class = end_points['size_class_label'].cpu().numpy() # B,K2
gt_size_residual = end_points['size_residual_label'].cpu().numpy() # B,K2,3
objectness_label = end_points['objectness_label'].detach().cpu().numpy() # (B,K,)
objectness_mask = end_points['objectness_mask'].detach().cpu().numpy() # (B,K,)
for i in range(batch_size):
if np.sum(objectness_label[i,:])>0:
pc_util.write_ply(pred_center[i,objectness_label[i,:]>0,0:3], os.path.join(dump_dir, '%06d_gt_positive_proposal_pc.ply'%(idx_beg+i)))
if np.sum(objectness_mask[i,:])>0:
pc_util.write_ply(pred_center[i,objectness_mask[i,:]>0,0:3], os.path.join(dump_dir, '%06d_gt_mask_proposal_pc.ply'%(idx_beg+i)))
pc_util.write_ply(gt_center[i,:,0:3], os.path.join(dump_dir, '%06d_gt_centroid_pc.ply'%(idx_beg+i)))
pc_util.write_ply_color(pred_center[i,:,0:3], objectness_label[i,:], os.path.join(dump_dir, '%06d_proposal_pc_objectness_label.obj'%(idx_beg+i)))
# Dump GT bounding boxes
obbs = []
for j in range(gt_center.shape[1]):
if gt_mask[i,j] == 0: continue
obb = config.param2obb(gt_center[i,j,0:3], gt_heading_class[i,j], gt_heading_residual[i,j],
gt_size_class[i,j], gt_size_residual[i,j])
obbs.append(obb)
if len(obbs)>0:
obbs = np.vstack(tuple(obbs)) # (num_gt_objects, 7)
pc_util.write_oriented_bbox(obbs, os.path.join(dump_dir, '%06d_gt_bbox.ply'%(idx_beg+i)))
# OPTIONALL, also dump prediction and gt details
if 'batch_pred_map_cls' in end_points:
for ii in range(batch_size):
fout = open(os.path.join(dump_dir, '%06d_pred_map_cls.txt'%(ii)), 'w')
for t in end_points['batch_pred_map_cls'][ii]:
fout.write(str(t[0])+' ')
fout.write(",".join([str(x) for x in list(t[1].flatten())]))
fout.write(' '+str(t[2]))
fout.write('\n')
fout.close()
if 'batch_gt_map_cls' in end_points:
for ii in range(batch_size):
fout = open(os.path.join(dump_dir, '%06d_gt_map_cls.txt'%(ii)), 'w')
for t in end_points['batch_gt_map_cls'][ii]:
fout.write(str(t[0])+' ')
fout.write(",".join([str(x) for x in list(t[1].flatten())]))
fout.write('\n')
fout.close()
os.makedirs(out_dir)
if not os.path.exists(out_gt_dir):
os.mkdir(out_gt_dir)
if not os.path.exists(out_recon_dir):
os.mkdir(out_recon_dir)
if not os.path.exists(out_recon_mesh_dir):
os.mkdir(out_recon_mesh_dir)
df = read_df_from_txt(df_txt_fn)
df_mesh = get_isosurface(df, iso_val)
cls_id, mdl_id = get_clsId_modelId(df_txt_fn)
scanned_pc_filename = os.path.join(SHAPENET_POINTCLOUD_DIR, cls_id, 'point_cloud_clean', mdl_id+'_clean.ply')
if not os.path.exists(scanned_pc_filename):
print('No scanning available: %s'%(scanned_pc_filename))
continue
scan_pc = pc_util.read_ply_xyz(scanned_pc_filename)
if 'v1' in SHAPENET_POINTCLOUD_DIR:
# for v1 data, rotate it to align with v2
scan_pc = pc_util.rotate_point_cloud_by_axis_angle(scan_pc, [0,1,0], 90)
scan_pc =pc_util.sample_point_cloud(scan_pc, points_sample_nb)
pc_util.write_ply(scan_pc, os.path.join(out_gt_dir, mdl_id+'.ply'))
df_mesh = align_dfmesh_scanpc(df_mesh, dim, scan_pc)
df_mesh.export(os.path.join(out_recon_mesh_dir, mdl_id+'.ply'))
#recon_samples, _ = trimesh.sample.sample_surface_even(df_mesh, points_sample_nb)
recon_samples, _ = trimesh.sample.sample_surface(df_mesh, points_sample_nb)
pc_util.write_ply(np.array(recon_samples), os.path.join(out_recon_dir, mdl_id+'.ply'))