├── empty_cloud.pcd ├── README.md ├── ycb └── objects.txt ├── ycb_filter_point_cloud.py ├── ycb_transform_point_cloud.py ├── ycb_downloader.py └── ycb_generate_point_cloud.py /empty_cloud.pcd: -------------------------------------------------------------------------------- 1 | # .PCD v.7 - Point Cloud Data file format 2 | VERSION .7 3 | FIELDS x y z rgb 4 | SIZE 4 4 4 4 5 | TYPE F F F F 6 | COUNT 1 1 1 1 7 | WIDTH 0 8 | HEIGHT 1 9 | VIEWPOINT 0 0 0 1 0 0 0 10 | POINTS 0 11 | DATA binary 12 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ycb-scripts 2 | Scripts to generate and process point clouds in [YCB dataset](http://www.ycbbenchmarks.com/). 3 | 4 | ## Usage 5 | 6 | ### Download Dataset 7 | ```bash 8 | python ycb_downloader.py 9 | ``` 10 | 11 | ### Generate Clouds From RGB-D Images 12 | 13 | ```bash 14 | python ycb_generate_point_cloud.py 15 | ``` 16 | 17 | ### Transform Clouds to Unified Camera Coordinate System 18 | ```bash 19 | python ycb_transform_point_cloud.py 20 | ``` 21 | 22 | ### Denoise Transformed Clouds 23 | ```bash 24 | python ycb_filter_point_cloud.py 25 | ``` 26 | 27 | After processing, all point clouds are denoised and transformed to the coordinate system of camera "NP5". It is helpful for object grasp pose generation. 28 | -------------------------------------------------------------------------------- /ycb/objects.txt: -------------------------------------------------------------------------------- 1 | 001_chips_can 2 | 002_master_chef_can 3 | 003_cracker_box 4 | 004_sugar_box 5 | 005_tomato_soup_can 6 | 006_mustard_bottle 7 | 007_tuna_fish_can 8 | 008_pudding_box 9 | 009_gelatin_box 10 | 010_potted_meat_can 11 | 011_banana 12 | 012_strawberry 13 | 013_apple 14 | 014_lemon 15 | 015_peach 16 | 016_pear 17 | 017_orange 18 | 018_plum 19 | 019_pitcher_base 20 | 021_bleach_cleanser 21 | 022_windex_bottle 22 | 024_bowl 23 | 026_sponge 24 | 027_skillet 25 | 029_plate 26 | 030_fork 27 | 031_spoon 28 | 032_knife 29 | 033_spatula 30 | 035_power_drill 31 | 036_wood_block 32 | 037_scissors 33 | 038_padlock 34 | 040_large_marker 35 | 041_small_marker 36 | 042_adjustable_wrench 37 | 043_phillips_screwdriver 38 | 044_flat_screwdriver 39 | 048_hammer 40 | 049_small_clamp 41 | 050_medium_clamp 42 | 051_large_clamp 43 | 052_extra_large_clamp 44 | 053_mini_soccer_ball 45 | 054_softball 46 | 055_baseball 47 | 056_tennis_ball 48 | 057_racquetball 49 | 058_golf_ball 50 | 059_chain 51 | 061_foam_brick 52 | 062_dice 53 | 063-a_marbles 54 | 065-a_cups 55 | 065-b_cups 56 | 065-c_cups 57 | 065-d_cups 58 | 065-e_cups 59 | 065-f_cups 60 | 065-g_cups 61 | 065-h_cups 62 | 065-i_cups 63 | 065-j_cups 64 | 070-a_colored_wood_blocks 65 | 071_nine_hole_peg_test 66 | 072-a_toy_airplane 67 | 072-b_toy_airplane 68 | 072-c_toy_airplane 69 | 072-d_toy_airplane 70 | 072-e_toy_airplane 71 | 072-f_toy_airplane 72 | 072-h_toy_airplane 73 | 072-i_toy_airplane 74 | 072-j_toy_airplane 75 | 072-k_toy_airplane 76 | 073-a_lego_duplo 77 | 073-b_lego_duplo 78 | 073-c_lego_duplo 79 | 073-d_lego_duplo 80 | 073-e_lego_duplo 81 | 073-f_lego_duplo 82 | 073-g_lego_duplo 83 | 073-h_lego_duplo 84 | 073-i_lego_duplo 85 | 073-j_lego_duplo 86 | 073-k_lego_duplo 87 | 073-l_lego_duplo 88 | 073-m_lego_duplo 89 | 076_timer 90 | 077_rubiks_cube 91 | -------------------------------------------------------------------------------- /ycb_filter_point_cloud.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import tqdm 4 | import open3d as o3 5 | 6 | # Parameters 7 | ycb_data_folder = "./ycb//" # Folder that contains the ycb data. 8 | viewpoint_cameras = ["NP1", "NP2", "NP3", "NP4", "NP5"] # Camera which the viewpoint will be generated. 9 | viewpoint_angles = [str(i) for i in range(0, 360, 3)] # Relative angle of the object w.r.t the camera (angle of the turntable). 10 | 11 | # Objects 12 | f = open(os.path.join(ycb_data_folder, 'objects.txt')) 13 | object_list = f.readlines() 14 | f.close() 15 | print('Get {0} objects in total.'.format(len(object_list))) 16 | for i, object_name in enumerate(object_list): 17 | object_list[i] = object_name.rstrip() 18 | 19 | # object_list = ['063-a_marbles'] 20 | 21 | voxel_size = 0.002 22 | dist_thresh = 0.005 23 | 24 | def compute_distance(A,B): 25 | A = A[:,:3].copy() 26 | B = B[:,:3].copy().T 27 | A2 = np.sum(np.square(A), axis=1, keepdims=True) 28 | B2 = np.sum(np.square(B), axis=0, keepdims=True) 29 | AB = np.matmul(A, B) 30 | dists = A2 + B2 - 2 * AB 31 | dists = np.sqrt(dists + 1e-6) 32 | return dists 33 | 34 | 35 | if __name__ == "__main__": 36 | reference_camera = "NP5" 37 | 38 | for target_object in object_list: 39 | print(target_object) 40 | output_dir = os.path.join(ycb_data_folder, target_object, "clouds_filtered") 41 | if not os.path.exists(output_dir): 42 | os.makedirs(output_dir) 43 | 44 | model_file = os.path.join(ycb_data_folder, target_object, 'clouds', 'merged_cloud.ply') 45 | model = o3.read_point_cloud(model_file) 46 | sampled_model = o3.voxel_down_sample(model, voxel_size=voxel_size) 47 | sampled_model_points = np.asarray(sampled_model.points) 48 | 49 | for viewpoint_camera in viewpoint_cameras: 50 | pbar = tqdm.tqdm(viewpoint_angles) 51 | for viewpoint_angle in pbar: 52 | pbar.set_description('Camera {0}, Angle {1}'.format(viewpoint_camera, viewpoint_angle)) 53 | 54 | cloud_file = os.path.join(ycb_data_folder, target_object, 'clouds_aligned', viewpoint_camera+'_'+viewpoint_angle+'.pcd') 55 | output_file = os.path.join(output_dir, viewpoint_camera+'_'+viewpoint_angle+'.pcd') 56 | # os.system('cp {0} {1}'.format(cloud_file, output_dir)) 57 | os.system('cp empty_cloud.pcd {}'.format(output_file)) 58 | 59 | cloud = o3.read_point_cloud(cloud_file) 60 | cloud_points = np.asarray(cloud.points) 61 | cloud_colors = np.asarray(cloud.colors) 62 | 63 | dists = compute_distance(cloud_points, sampled_model_points) 64 | mask = (dists.min(axis=1) < dist_thresh) 65 | 66 | filtered_cloud = o3.PointCloud() 67 | filtered_cloud.points = o3.Vector3dVector(cloud_points[mask]) 68 | filtered_cloud.colors = o3.Vector3dVector(cloud_colors[mask]) 69 | 70 | o3.write_point_cloud(output_file, filtered_cloud) 71 | -------------------------------------------------------------------------------- /ycb_transform_point_cloud.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import h5py as h5 4 | import math 5 | import tqdm 6 | import open3d as o3 7 | 8 | # Parameters 9 | ycb_data_folder = "./ycb//" # Folder that contains the ycb data. 10 | viewpoint_cameras = ["NP1", "NP2", "NP3", "NP4", "NP5"] # Camera which the viewpoint will be generated. 11 | viewpoint_angles = [str(i) for i in range(0, 360, 3)] # Relative angle of the object w.r.t the camera (angle of the turntable). 12 | 13 | # Objects 14 | f = open(os.path.join(ycb_data_folder, 'objects.txt')) 15 | object_list = f.readlines() 16 | f.close() 17 | print('Get {0} objects in total.'.format(len(object_list))) 18 | for i, object_name in enumerate(object_list): 19 | object_list[i] = object_name.rstrip() 20 | 21 | # object_list = ['072-a_toy_airplane'] 22 | 23 | def getTransform(calibration, viewpoint_camera, referenceCamera, transformation): 24 | CamFromRefKey = "H_{0}_from_{1}".format(viewpoint_camera, referenceCamera) 25 | CamFromRef = calibration[CamFromRefKey][:] 26 | 27 | TableFromRefKey = "H_table_from_reference_camera" 28 | TableFromRef = transformation[TableFromRefKey][:] 29 | return np.dot(TableFromRef, np.linalg.inv(CamFromRef)) 30 | 31 | 32 | if __name__ == "__main__": 33 | reference_camera = "NP5" 34 | 35 | for target_object in object_list: 36 | print(target_object) 37 | output_dir = os.path.join(ycb_data_folder, target_object, "clouds_aligned") 38 | if not os.path.exists(output_dir): 39 | os.makedirs(output_dir) 40 | for viewpoint_camera in viewpoint_cameras: 41 | pbar = tqdm.tqdm(viewpoint_angles) 42 | for viewpoint_angle in pbar: 43 | pbar.set_description('Camera {0}, Angle {1}'.format(viewpoint_camera, viewpoint_angle)) 44 | 45 | cloud_file = os.path.join(ycb_data_folder, target_object, 'clouds', viewpoint_camera+'_'+viewpoint_angle+'.pcd') 46 | output_file = os.path.join(output_dir, viewpoint_camera+'_'+viewpoint_angle+'.pcd') 47 | # os.system('cp {0} {1}'.format(cloud_file, output_dir)) 48 | os.system('cp empty_cloud.pcd {}'.format(output_file)) 49 | 50 | calibrationFilename = os.path.join(ycb_data_folder+target_object, "calibration.h5") 51 | transformationFilename = os.path.join(ycb_data_folder, target_object, 'poses', '{}_{}_pose.h5'.format(reference_camera, viewpoint_angle)) 52 | 53 | calibration = h5.File(calibrationFilename) 54 | transformation = h5.File(transformationFilename) 55 | H_table_from_cam = getTransform(calibration, viewpoint_camera, reference_camera, transformation) 56 | 57 | pointCloud = o3.read_point_cloud(cloud_file) 58 | points = np.asarray(pointCloud.points) 59 | ones = np.ones(len(points))[:,np.newaxis] 60 | 61 | points_ = np.concatenate([points,ones], axis=1) 62 | aligned_points_ = np.matmul(H_table_from_cam, points_.T) 63 | aligned_points = aligned_points_.T[:,:3] 64 | 65 | aligned_cloud = o3.PointCloud() 66 | aligned_cloud.points = o3.Vector3dVector(aligned_points) 67 | aligned_cloud.colors = pointCloud.colors 68 | 69 | o3.write_point_cloud(output_file, aligned_cloud) 70 | -------------------------------------------------------------------------------- /ycb_downloader.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import json 4 | import urllib 5 | import urllib2 6 | 7 | output_directory = "./ycb" 8 | 9 | # You can either set this to "all" or a list of the objects that you'd like to 10 | # download. 11 | # objects_to_download = "all" 12 | #objects_to_download = ["002_master_chef_can", "003_cracker_box", "004_sugar_box", "005_tomato_soup_can"] 13 | # objects_to_download = ["006_mustard_bottle", "007_tuna_fish_can", "008_pudding_box", "009_gelatin_box", "010_potted_meat_can", "011_banana", "012_strawberry", "013_apple", "014_lemon", "015_peach", "016_pear", "017_orange", "018_plum", "019_pitcher_base", "021_bleach_cleanser", "022_windex_bottle", "024_bowl", "026_sponge", "027_skillet", "029_plate", "030_fork", "031_spoon", "032_knife", "033_spatula", "035_power_drill", "036_wood_block", "037_scissors", "038_padlock", "039_key", "040_large_marker", "041_small_marker", "042_adjustable_wrench", "043_phillips_screwdriver", "044_flat_screwdriver", "048_hammer", "049_small_clamp", "050_medium_clamp", "051_large_clamp", "052_extra_large_clamp", "053_mini_soccer_ball", "054_softball", "055_baseball", "056_tennis_ball", "057_racquetball", "058_golf_ball", "059_chain", "061_foam_brick", "062_dice", "063-a_marbles", "065-a_cups", "065-b_cups", "065-c_cups", "065-d_cups", "065-e_cups", "065-f_cups", "065-g_cups", "065-h_cups", "065-i_cups", "065-j_cups", "070-a_colored_wood_blocks", "071_nine_hole_peg_test", "072-a_toy_airplane", "072-b_toy_airplane", "072-c_toy_airplane", "072-d_toy_airplane", "072-e_toy_airplane", "072-f_toy_airplane", "072-g_toy_airplane", "072-h_toy_airplane", "072-i_toy_airplane", "072-j_toy_airplane", "072-k_toy_airplane", "073-a_lego_duplo", "073-b_lego_duplo", "073-c_lego_duplo", "073-d_lego_duplo", "073-e_lego_duplo", "073-f_lego_duplo", "073-g_lego_duplo", "073-h_lego_duplo", "073-i_lego_duplo", "073-j_lego_duplo", "073-k_lego_duplo", "073-l_lego_duplo", "073-m_lego_duplo", "076_timer", "077_rubiks_cube"] 14 | # objects_to_download = ["025_mug"] 15 | 16 | f = open(os.path.join(output_directory, 'objects.txt')) 17 | objects_to_download = f.readlines() 18 | f.close() 19 | 20 | # You can edit this list to only download certain kinds of files. 21 | # 'berkeley_rgbd' contains all of the depth maps and images from the Carmines. 22 | # 'berkeley_rgb_highres' contains all of the high-res images from the Canon cameras. 23 | # 'berkeley_processed' contains all of the segmented point clouds and textured meshes. 24 | # 'google_16k' contains google meshes with 16k vertices. 25 | # 'google_64k' contains google meshes with 64k vertices. 26 | # 'google_512k' contains google meshes with 512k vertices. 27 | # See the website for more details. 28 | #files_to_download = ["berkeley_rgbd", "berkeley_rgb_highres", "berkeley_processed", "google_16k", "google_64k", "google_512k"] 29 | files_to_download = ["berkeley_rgbd", "berkeley_rgb_highres" "berkeley_processed"] 30 | 31 | # Extract all files from the downloaded .tgz, and remove .tgz files. 32 | # If false, will just download all .tgz files to output_directory 33 | extract = True 34 | 35 | base_url = "http://ycb-benchmarks.s3-website-us-east-1.amazonaws.com/data/" 36 | objects_url = base_url + "objects.json" 37 | 38 | if not os.path.exists(output_directory): 39 | os.makedirs(output_directory) 40 | 41 | def fetch_objects(url): 42 | response = urllib2.urlopen(url) 43 | html = response.read() 44 | objects = json.loads(html) 45 | return objects["objects"] 46 | 47 | def download_file(url, filename): 48 | u = urllib2.urlopen(url) 49 | f = open(filename, 'wb') 50 | meta = u.info() 51 | file_size = int(meta.getheaders("Content-Length")[0]) 52 | print("Downloading: %s (%s MB)" % (filename, file_size/1000000.0)) 53 | 54 | file_size_dl = 0 55 | block_sz = 65536 56 | while True: 57 | buffer = u.read(block_sz) 58 | if not buffer: 59 | break 60 | 61 | file_size_dl += len(buffer) 62 | f.write(buffer) 63 | status = r"%10d [%3.2f%%]" % (file_size_dl/1000000.0, file_size_dl * 100. / file_size) 64 | status = status + chr(8)*(len(status)+1) 65 | print(status,) 66 | f.close() 67 | 68 | def tgz_url(object, type): 69 | if type in ["berkeley_rgbd", "berkeley_rgb_highres"]: 70 | return base_url + "berkeley/{object}/{object}_{type}.tgz".format(object=object,type=type) 71 | elif type in ["berkeley_processed"]: 72 | return base_url + "berkeley/{object}/{object}_berkeley_meshes.tgz".format(object=object,type=type) 73 | else: 74 | return base_url + "google/{object}_{type}.tgz".format(object=object,type=type) 75 | 76 | def extract_tgz(filename, dir): 77 | tar_command = "tar -xzf {filename} -C {dir}".format(filename=filename,dir=dir) 78 | os.system(tar_command) 79 | os.remove(filename) 80 | 81 | def check_url(url): 82 | try: 83 | request = urllib2.Request(url) 84 | request.get_method = lambda : 'HEAD' 85 | response = urllib2.urlopen(request) 86 | return True 87 | except Exception as e: 88 | return False 89 | 90 | 91 | if __name__ == "__main__": 92 | 93 | objects = objects_to_download 94 | # objects = fetch_objects(objects_url) 95 | 96 | for object in objects: 97 | if objects_to_download == "all" or object in objects_to_download: 98 | for file_type in files_to_download: 99 | url = tgz_url(object, file_type) 100 | if not check_url(url): 101 | continue 102 | filename = "{path}/{object}_{file_type}.tgz".format(path=output_directory, 103 | object=object, 104 | file_type=file_type) 105 | download_file(url, filename) 106 | if extract: 107 | extract_tgz(filename, output_directory) 108 | -------------------------------------------------------------------------------- /ycb_generate_point_cloud.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import h5py as h5 4 | from imageio import imread 5 | import math 6 | import struct 7 | import tqdm 8 | import open3d as o3 9 | 10 | # Parameters 11 | ycb_data_folder = "./ycb//" # Folder that contains the ycb data. 12 | viewpoint_cameras = ["NP1", "NP2", "NP3", "NP4", "NP5"] # Camera which the viewpoint will be generated. 13 | viewpoint_angles = [str(i) for i in range(0, 360, 3)] # Relative angle of the object w.r.t the camera (angle of the turntable). 14 | 15 | # Objects 16 | f = open(os.path.join(ycb_data_folder, 'objects.txt')) 17 | object_list = f.readlines() 18 | f.close() 19 | print('Get {0} objects in total.'.format(len(object_list))) 20 | for i, object_name in enumerate(object_list): 21 | object_list[i] = object_name.rstrip() 22 | 23 | # object_list = ['001_chips_can'] 24 | # object_list = object_list[1:] 25 | # object_list = object_list[3:20] 26 | # object_list = object_list[20:38] 27 | # object_list = object_list[38:56] 28 | # object_list = object_list[56:74] 29 | # object_list = object_list[74:] 30 | 31 | 32 | def im2col(im, psize): 33 | n_channels = 1 if len(im.shape) == 2 else im.shape[0] 34 | (n_channels, rows, cols) = (1,) * (3 - len(im.shape)) + im.shape 35 | 36 | im_pad = np.zeros((n_channels, 37 | int(math.ceil(1.0 * rows / psize) * psize), 38 | int(math.ceil(1.0 * cols / psize) * psize))) 39 | im_pad[:, 0:rows, 0:cols] = im 40 | 41 | final = np.zeros((im_pad.shape[1], im_pad.shape[2], n_channels, 42 | psize, psize)) 43 | for c in range(n_channels): 44 | for x in range(psize): 45 | for y in range(psize): 46 | im_shift = np.vstack( 47 | (im_pad[c, x:], im_pad[c, :x])) 48 | im_shift = np.column_stack( 49 | (im_shift[:, y:], im_shift[:, :y])) 50 | final[x::psize, y::psize, c] = np.swapaxes( 51 | im_shift.reshape(int(im_pad.shape[1] / psize), psize, 52 | int(im_pad.shape[2] / psize), psize), 1, 2) 53 | 54 | return np.squeeze(final[0:rows - psize + 1, 0:cols - psize + 1]) 55 | 56 | def filterDiscontinuities(depthMap): 57 | filt_size = 7 58 | thresh = 1000 59 | 60 | # Ensure that filter sizes are okay 61 | assert filt_size % 2 == 1, "Can only use odd filter sizes." 62 | 63 | # Compute discontinuities 64 | offset = int((filt_size - 1) / 2) 65 | patches = 1.0 * im2col(depthMap, filt_size) 66 | mids = patches[:, :, offset, offset] 67 | mins = np.min(patches, axis=(2, 3)) 68 | maxes = np.max(patches, axis=(2, 3)) 69 | 70 | discont = np.maximum(np.abs(mins - mids), 71 | np.abs(maxes - mids)) 72 | mark = discont > thresh 73 | 74 | # Account for offsets 75 | final_mark = np.zeros((480, 640), dtype=np.uint16) 76 | final_mark[offset:offset + mark.shape[0], 77 | offset:offset + mark.shape[1]] = mark 78 | 79 | return depthMap * (1 - final_mark) 80 | 81 | def registerDepthMap(unregisteredDepthMap, 82 | rgbImage, 83 | depthK, 84 | rgbK, 85 | H_RGBFromDepth): 86 | 87 | 88 | unregisteredHeight = unregisteredDepthMap.shape[0] 89 | unregisteredWidth = unregisteredDepthMap.shape[1] 90 | 91 | registeredHeight = rgbImage.shape[0] 92 | registeredWidth = rgbImage.shape[1] 93 | 94 | registeredDepthMap = np.zeros((registeredHeight, registeredWidth)) 95 | 96 | xyzDepth = np.empty((4,1)) 97 | xyzRGB = np.empty((4,1)) 98 | 99 | # Ensure that the last value is 1 (homogeneous coordinates) 100 | xyzDepth[3] = 1 101 | 102 | invDepthFx = 1.0 / depthK[0,0] 103 | invDepthFy = 1.0 / depthK[1,1] 104 | depthCx = depthK[0,2] 105 | depthCy = depthK[1,2] 106 | 107 | rgbFx = rgbK[0,0] 108 | rgbFy = rgbK[1,1] 109 | rgbCx = rgbK[0,2] 110 | rgbCy = rgbK[1,2] 111 | 112 | undistorted = np.empty(2) 113 | for v in range(unregisteredHeight): 114 | for u in range(unregisteredWidth): 115 | 116 | depth = unregisteredDepthMap[v,u] 117 | if depth == 0: 118 | continue 119 | 120 | xyzDepth[0] = ((u - depthCx) * depth) * invDepthFx 121 | xyzDepth[1] = ((v - depthCy) * depth) * invDepthFy 122 | xyzDepth[2] = depth 123 | 124 | xyzRGB[0] = (H_RGBFromDepth[0,0] * xyzDepth[0] + 125 | H_RGBFromDepth[0,1] * xyzDepth[1] + 126 | H_RGBFromDepth[0,2] * xyzDepth[2] + 127 | H_RGBFromDepth[0,3]) 128 | xyzRGB[1] = (H_RGBFromDepth[1,0] * xyzDepth[0] + 129 | H_RGBFromDepth[1,1] * xyzDepth[1] + 130 | H_RGBFromDepth[1,2] * xyzDepth[2] + 131 | H_RGBFromDepth[1,3]) 132 | xyzRGB[2] = (H_RGBFromDepth[2,0] * xyzDepth[0] + 133 | H_RGBFromDepth[2,1] * xyzDepth[1] + 134 | H_RGBFromDepth[2,2] * xyzDepth[2] + 135 | H_RGBFromDepth[2,3]) 136 | 137 | invRGB_Z = 1.0 / xyzRGB[2] 138 | undistorted[0] = (rgbFx * xyzRGB[0]) * invRGB_Z + rgbCx 139 | undistorted[1] = (rgbFy * xyzRGB[1]) * invRGB_Z + rgbCy 140 | 141 | uRGB = int(undistorted[0] + 0.5) 142 | vRGB = int(undistorted[1] + 0.5) 143 | 144 | if (uRGB < 0 or uRGB >= registeredWidth) or (vRGB < 0 or vRGB >= registeredHeight): 145 | continue 146 | 147 | registeredDepth = xyzRGB[2] 148 | if registeredDepth > registeredDepthMap[vRGB,uRGB]: 149 | registeredDepthMap[vRGB,uRGB] = registeredDepth 150 | 151 | return registeredDepthMap 152 | 153 | def registeredDepthMapToPointCloud(depthMap, rgbImage, rgbK, organized=True, mask=None): 154 | rgbCx = rgbK[0,2] 155 | rgbCy = rgbK[1,2] 156 | invRGBFx = 1.0/rgbK[0,0] 157 | invRGBFy = 1.0/rgbK[1,1] 158 | 159 | height = depthMap.shape[0] 160 | width = depthMap.shape[1] 161 | 162 | if organized: 163 | cloud = np.empty((height, width, 6), dtype=np.float) 164 | else: 165 | cloud = np.empty((1, height*width, 6), dtype=np.float) 166 | 167 | goodPointsCount = 0 168 | for v in range(height): 169 | for u in range(width): 170 | 171 | depth = depthMap[v,u] 172 | 173 | if organized: 174 | row = v 175 | col = u 176 | else: 177 | row = 0 178 | col = goodPointsCount 179 | 180 | if depth <= 0 or (mask is not None and mask[v,u] > 0): 181 | if organized: 182 | if depth <= 0: 183 | cloud[row,col,0] = float('nan') 184 | cloud[row,col,1] = float('nan') 185 | cloud[row,col,2] = float('nan') 186 | cloud[row,col,3] = 0 187 | cloud[row,col,4] = 0 188 | cloud[row,col,5] = 0 189 | continue 190 | 191 | cloud[row,col,0] = (u - rgbCx) * depth * invRGBFx 192 | cloud[row,col,1] = (v - rgbCy) * depth * invRGBFy 193 | cloud[row,col,2] = depth 194 | cloud[row,col,3] = rgbImage[v,u,0] 195 | cloud[row,col,4] = rgbImage[v,u,1] 196 | cloud[row,col,5] = rgbImage[v,u,2] 197 | if not organized: 198 | goodPointsCount += 1 199 | 200 | if not organized: 201 | cloud = cloud[:,:goodPointsCount,:] 202 | 203 | return cloud 204 | 205 | 206 | def writePLY(filename, cloud, faces=[]): 207 | if len(cloud.shape) != 3: 208 | print("Expected pointCloud to have 3 dimensions. Got %d instead" % len(cloud.shape)) 209 | return 210 | 211 | color = True if cloud.shape[2] == 6 else False 212 | num_points = cloud.shape[0]*cloud.shape[1] 213 | 214 | header_lines = [ 215 | 'ply', 216 | 'format ascii 1.0', 217 | 'element vertex %d' % num_points, 218 | 'property float x', 219 | 'property float y', 220 | 'property float z', 221 | ] 222 | if color: 223 | header_lines.extend([ 224 | 'property uchar diffuse_red', 225 | 'property uchar diffuse_green', 226 | 'property uchar diffuse_blue', 227 | ]) 228 | if faces != None: 229 | header_lines.extend([ 230 | 'element face %d' % len(faces), 231 | 'property list uchar int vertex_indices' 232 | ]) 233 | 234 | header_lines.extend([ 235 | 'end_header', 236 | ]) 237 | 238 | f = open(filename, 'w+') 239 | f.write('\n'.join(header_lines)) 240 | f.write('\n') 241 | 242 | lines = [] 243 | for i in range(cloud.shape[0]): 244 | for j in range(cloud.shape[1]): 245 | if color: 246 | lines.append('%s %s %s %d %d %d' % tuple(cloud[i, j, :].tolist())) 247 | else: 248 | lines.append('%s %s %s' % tuple(cloud[i, j, :].tolist())) 249 | 250 | for face in faces: 251 | lines.append(('%d' + ' %d'*len(face)) % tuple([len(face)] + list(face))) 252 | 253 | f.write('\n'.join(lines) + '\n') 254 | f.close() 255 | 256 | def writePCD(pointCloud, filename, ascii=False): 257 | if len(pointCloud.shape) != 3: 258 | print("Expected pointCloud to have 3 dimensions. Got %d instead" % len(pointCloud.shape)) 259 | return 260 | with open(filename, 'w') as f: 261 | height = pointCloud.shape[0] 262 | width = pointCloud.shape[1] 263 | f.write("# .PCD v.7 - Point Cloud Data file format\n") 264 | f.write("VERSION .7\n") 265 | if pointCloud.shape[2] == 3: 266 | f.write("FIELDS x y z\n") 267 | f.write("SIZE 4 4 4\n") 268 | f.write("TYPE F F F\n") 269 | f.write("COUNT 1 1 1\n") 270 | else: 271 | f.write("FIELDS x y z rgb\n") 272 | f.write("SIZE 4 4 4 4\n") 273 | f.write("TYPE F F F F\n") 274 | f.write("COUNT 1 1 1 1\n") 275 | f.write("WIDTH %d\n" % width) 276 | f.write("HEIGHT %d\n" % height) 277 | f.write("VIEWPOINT 0 0 0 1 0 0 0\n") 278 | f.write("POINTS %d\n" % (height * width)) 279 | if ascii: 280 | f.write("DATA ascii\n") 281 | for row in range(height): 282 | for col in range(width): 283 | if pointCloud.shape[2] == 3: 284 | f.write("%f %f %f\n" % tuple(pointCloud[row, col, :])) 285 | else: 286 | f.write("%f %f %f" % tuple(pointCloud[row, col, :3])) 287 | r = int(pointCloud[row, col, 3]) 288 | g = int(pointCloud[row, col, 4]) 289 | b = int(pointCloud[row, col, 5]) 290 | rgb_int = (r << 16) | (g << 8) | b 291 | packed = struct.pack('i', rgb_int) 292 | rgb = struct.unpack('f', packed)[0] 293 | f.write(" %.12e\n" % rgb) 294 | else: 295 | f.write("DATA binary\n") 296 | if pointCloud.shape[2] == 6: 297 | # These are written as bgr because rgb is interpreted as a single 298 | # little-endian float. 299 | dt = np.dtype([('x', np.float32), 300 | ('y', np.float32), 301 | ('z', np.float32), 302 | ('b', np.uint8), 303 | ('g', np.uint8), 304 | ('r', np.uint8), 305 | ('I', np.uint8)]) 306 | pointCloud_tmp = np.zeros((height*width, 1), dtype=dt) 307 | for i, k in enumerate(['x', 'y', 'z', 'r', 'g', 'b']): 308 | pointCloud_tmp[k] = pointCloud[:, :, i].reshape((height*width, 1)) 309 | pointCloud_tmp.tofile(f) 310 | else: 311 | dt = np.dtype([('x', np.float32), 312 | ('y', np.float32), 313 | ('z', np.float32), 314 | ('I', np.uint8)]) 315 | pointCloud_tmp = np.zeros((height*width, 1), dtype=dt) 316 | for i, k in enumerate(['x', 'y', 'z']): 317 | pointCloud_tmp[k] = pointCloud[:, :, i].reshape((height*width, 1)) 318 | pointCloud_tmp.tofile(f) 319 | 320 | def getRGBFromDepthTransform(calibration, camera, referenceCamera): 321 | irKey = "H_{0}_ir_from_{1}".format(camera, referenceCamera) 322 | rgbKey = "H_{0}_from_{1}".format(camera, referenceCamera) 323 | 324 | rgbFromRef = calibration[rgbKey][:] 325 | irFromRef = calibration[irKey][:] 326 | 327 | return np.dot(rgbFromRef, np.linalg.inv(irFromRef)) 328 | 329 | 330 | if __name__ == "__main__": 331 | reference_camera = "NP5" 332 | 333 | for target_object in object_list: 334 | print(target_object) 335 | if not os.path.exists(ycb_data_folder+"/"+target_object+"/clouds"): 336 | os.makedirs(ycb_data_folder+"/"+target_object+"/clouds") 337 | for viewpoint_camera in viewpoint_cameras: 338 | pbar = tqdm.tqdm(viewpoint_angles) 339 | for viewpoint_angle in pbar: 340 | pbar.set_description('Camera {0}, Angle {1}'.format(viewpoint_camera, viewpoint_angle)) 341 | basename = "{0}_{1}".format(viewpoint_camera, viewpoint_angle) 342 | depthFilename = os.path.join(ycb_data_folder+ target_object, basename + ".h5") 343 | rgbFilename = os.path.join(ycb_data_folder + target_object, basename + ".jpg") 344 | maskFilename = os.path.join(ycb_data_folder + target_object + '/masks', basename + "_mask.pbm") 345 | 346 | calibrationFilename = os.path.join(ycb_data_folder+target_object, "calibration.h5") 347 | calibration = h5.File(calibrationFilename) 348 | 349 | if not os.path.isfile(rgbFilename): 350 | print("The rgbd data is not available for the target object \"%s\". Please download the data first." % target_object) 351 | exit(1) 352 | rgbImage = imread(rgbFilename) 353 | mask = imread(maskFilename) 354 | depthK = calibration["{0}_depth_K".format(viewpoint_camera)][:] 355 | rgbK = calibration["{0}_rgb_K".format(viewpoint_camera)][:] 356 | depthScale = np.array(calibration["{0}_ir_depth_scale".format(viewpoint_camera)]) * .0001 # 100um to meters 357 | H_RGBFromDepth = getRGBFromDepthTransform(calibration, viewpoint_camera, reference_camera) 358 | 359 | unregisteredDepthMap = h5.File(depthFilename)["depth"][:] 360 | unregisteredDepthMap = filterDiscontinuities(unregisteredDepthMap) * depthScale 361 | 362 | registeredDepthMap = registerDepthMap(unregisteredDepthMap, 363 | rgbImage, 364 | depthK, 365 | rgbK, 366 | H_RGBFromDepth) 367 | 368 | pointCloud = registeredDepthMapToPointCloud(registeredDepthMap, rgbImage, rgbK, organized=False, mask=mask) 369 | 370 | # writePLY(ycb_data_folder+target_object+"/clouds/pc_"+viewpoint_camera+"_"+referenceCamera+"_"+viewpoint_angle+".ply", pointCloud) 371 | writePCD(pointCloud, ycb_data_folder+target_object+"/clouds/"+viewpoint_camera+"_"+viewpoint_angle+".pcd", ascii=False) 372 | --------------------------------------------------------------------------------