def generate_random_crops(config): crop_size = config['crop_size'] scene_crop_size = config['scene_crop_size'] loss_object = tf.keras.losses.MeanSquaredError() count = 0 pos_count = 0 neg_count = 0 room_count = 0 null_value = 0.0 scan_dirs = [x[0] for x in os.walk(config['in_dir'])] bar = helpers.progbar(len(scan_dirs[1:])) bar.start() with tf.io.TFRecordWriter( config['out_train_file']) as train_writer, tf.io.TFRecordWriter( config['out_test_file']) as test_writer: for i, scan_dir in enumerate(scan_dirs[1:]): room_count += 1 bar.update(room_count) scene_id = scan_dir.split('/')[-1] pts_path = os.path.join( scan_dir, '{}_vh_clean_2.labels.ply'.format(scene_id)) agg_path = os.path.join( scan_dir, '{}_vh_clean.aggregation.json'.format(scene_id)) segmap_path = os.path.join( scan_dir, '{}_vh_clean_2.0.010000.segs.json'.format(scene_id)) with open(agg_path, 'r') as f: agg = json.load(f) segmap = su.load_segmap(segmap_path) room_cloud = su.load_pointcloud(pts_path) objects, obj_means = su.load_objects(room_cloud, agg, segmap, config['label_objects']) object_thetas, object_extents = helpers.get_oabb(objects) for obj_count, obj in enumerate(objects): if i != 5 and np.random.rand() > config['train_ratio']: continue # ignore object obj_mean = obj_means[obj_count] for j in range(config['n_crops_per_object']): crop_x = np.random.uniform( obj_mean[0] - (crop_size[0] / 2), obj_mean[0] + (crop_size[0] / 2)) crop_y = np.random.uniform( obj_mean[1] - (crop_size[1] / 2), obj_mean[1] + (crop_size[1] / 2)) crop_z = np.random.uniform( obj_mean[2] - (crop_size[2] / 2), obj_mean[2] + (crop_size[2] / 2)) scene_bbox = np.array([ crop_x - (scene_crop_size[0] / 2.), crop_y - (scene_crop_size[1] / 2.), crop_z - 1e10, crop_x + (scene_crop_size[0] / 2.), crop_y + (scene_crop_size[1] / 2.), crop_z + 1e10, ]) bbox = np.array([ crop_x - (crop_size[0] / 2.), crop_y - (crop_size[1] / 2.), crop_z - (crop_size[2] / 2.), crop_x + (crop_size[0] / 2.), crop_y + (crop_size[1] / 2.), crop_z + (crop_size[2] / 2.), ]) crop_centre = np.array([crop_x, crop_y, crop_z]) l_dist = [ helpers.euc_dist([crop_x, crop_y, crop_z], c) for c in obj_means ] obj_mean = obj_means[np.argmin(l_dist)] obj_ext = object_extents[np.argmin(l_dist)] obj_theta = object_thetas[np.argmin(l_dist)] curr_obj = objects[np.argmin(l_dist)] norm_bbox = np.array([ bbox[0] - crop_centre[0], bbox[1] - crop_centre[1], bbox[2] - crop_centre[2], bbox[3] - crop_centre[0], bbox[4] - crop_centre[1], bbox[5] - crop_centre[2] ]) label_mean = obj_mean - crop_centre for _ in range(config['n_rotations_per_object']): tmp_room_cloud = helpers.crop_bbox( room_cloud, scene_bbox) _, tmp_room_cloud = helpers.get_fixed_pts( tmp_room_cloud, config['n_scene_pts']) tmp_room_crop = helpers.crop_bbox(tmp_room_cloud, bbox) tmp_room_crop[:, :3] -= crop_centre theta = np.random.uniform() * 2.0 * np.pi tmp_room_crop[:, :3] = helpers.rotate_euler( tmp_room_crop[:, :3], theta) label_mean_r = helpers.rotate_euler( label_mean.reshape(1, 3), theta)[0] new_angle = np.arctan2(obj_theta[1], obj_theta[0]) - theta theta_r = np.array( [np.cos(new_angle), np.sin(new_angle)]) tmp_room_crop = helpers.crop_bbox( tmp_room_crop, norm_bbox) check = helpers.check_occupancy( curr_obj[:, :3] - crop_centre, norm_bbox) if check == 0: continue dist = np.min(l_dist) dir = (label_mean_r - np.array([0, 0, 0])) / dist pred_dummy = tf.Variable([0., 0., 0.]) with tf.GradientTape() as tape: tape.watch(pred_dummy) loss = loss_object(tf.Variable(label_mean_r), pred_dummy) grad = tape.gradient(loss, pred_dummy).numpy() ret, tmp_room_crop = helpers.get_fixed_pts( tmp_room_crop, config['n_pts']) if ret == False: continue label = np.array([ 1, dist, grad[0], grad[1], grad[2], obj_ext[0], obj_ext[1], obj_ext[2], theta_r[0], theta_r[1] ]) # Uncomment to visualise training data # plot(tmp_room_crop[:, :3], tmp_room_crop[:, 3:6], -grad, label_mean_r, label) tf_example = create_example(tmp_room_crop[:, :3], tmp_room_crop[:, 3:6], label) assert tmp_room_crop.shape[0] == config[ 'n_pts'], '{}'.format(tmp_room_crop.shape) # 80/20 train/test split if count % 5 != 0: train_writer.write(tf_example.SerializeToString()) if count % 5 == 0: test_writer.write(tf_example.SerializeToString()) count += 1 pos_count += 1 for j in range(config['n_negative_samples']): for _ in range(1000): crop_x = np.random.uniform(np.min(room_cloud[:, 0]), np.max(room_cloud[:, 0])) crop_y = np.random.uniform(np.min(room_cloud[:, 1]), np.max(room_cloud[:, 1])) crop_z = np.random.uniform(np.min(room_cloud[:, 2]), np.max(room_cloud[:, 2])) crop_point = np.array([crop_x, crop_y, crop_z]) scene_bbox = np.array([ crop_point[0] - (scene_crop_size[0] / 2.), crop_point[1] - (scene_crop_size[1] / 2.), crop_point[2] - (scene_crop_size[2] / 2.), crop_point[0] + (scene_crop_size[0] / 2.), crop_point[1] + (scene_crop_size[1] / 2.), crop_point[2] + (scene_crop_size[2] / 2.), ]) bbox = np.array([ crop_point[0] - (crop_size[0] / 2), crop_point[1] - (crop_size[1] / 2), crop_point[2] - (crop_size[2] / 2), crop_point[0] + (crop_size[0] / 2), crop_point[1] + (crop_size[1] / 2), crop_point[2] + (crop_size[2] / 2), ]) check = np.array([ helpers.check_occupancy(obj_pts, bbox) for obj_pts in objects ]) check = check.astype(np.bool) if True in check: continue norm_bbox = np.array([ bbox[0] - crop_point[0], bbox[1] - crop_point[1], bbox[2] - crop_point[2], bbox[3] - crop_point[0], bbox[4] - crop_point[1], bbox[5] - crop_point[2] ]) for _ in range(config['n_rotations_per_object']): tmp_room_cloud = helpers.crop_bbox( room_cloud, scene_bbox) ret, tmp_room_cloud, = helpers.get_fixed_pts( tmp_room_cloud, config['n_scene_pts']) room_crop = helpers.crop_bbox(tmp_room_cloud, bbox) room_crop[:, :3] -= crop_point[:3] tmp_room_crop = room_crop.copy() theta = np.random.uniform() * 2.0 * np.pi tmp_room_crop[:, :3] = helpers.rotate_euler( tmp_room_crop[:, :3], theta) tmp_room_crop = helpers.crop_bbox( tmp_room_crop, norm_bbox) ret, tmp_room_crop = helpers.get_fixed_pts( tmp_room_crop, config['n_pts']) label = np.array([ 0, 10, null_value, null_value, null_value, null_value, null_value, null_value, null_value, null_value ]) tf_example = create_example( tmp_room_crop[:, :3], tmp_room_crop[:, 3:6], label) assert tmp_room_crop.shape[0] == config[ 'n_pts'], '{}'.format( tmp_room_crop.shape[0].shape) if count % 5 != 0: train_writer.write( tf_example.SerializeToString()) if count % 5 == 0: test_writer.write( tf_example.SerializeToString()) count += 1 neg_count += 1 break bar.finish() print('[info] total scenes saved: {}, {} positive, {} negative'.format( count, pos_count, neg_count))
def create_records(): dataset = ['training', 'testing'] max_objects = cfg['max_objects'] train = True if cfg['dataset'] == 'training' else False img_dir = os.path.join(cfg['in_dir'], 'data_object_image_2', cfg['dataset'], 'image_2') scan_dir = os.path.join(cfg['in_dir'], 'data_object_velodyne', cfg['dataset'], 'velodyne') calib_dir = os.path.join(cfg['in_dir'], 'data_object_calib', cfg['dataset'], 'calib') label_dir = os.path.join(cfg['in_dir'], 'data_object_label_2', 'training', 'label_2') img_files = glob.glob(os.path.join(img_dir, '*.png')) n_scenes = cfg['n_scenes'] if cfg['n_scenes'] > 0 else len(img_files) bar = helpers.progbar(n_scenes) bar.start() with tf.io.TFRecordWriter(cfg['out_train']) as train_writer, tf.io.TFRecordWriter(cfg['out_val']) as val_writer: for scene_id, img_file in enumerate(img_files): if scene_id == n_scenes: break bar.update(scene_id) base = os.path.splitext(os.path.basename(img_file))[0] scan_file = os.path.join(scan_dir, base+'.bin') calib_file = os.path.join(calib_dir, base+'.txt') label_file = os.path.join(label_dir, base+'.txt') img_arr = kitti_utils.load_img(img_file) orig_img_size = img_arr.shape[:2] img_arr = cv.resize(img_arr, (cfg['img_size'][1], cfg['img_size'][0]), interpolation=cv.INTER_CUBIC) _, img = cv.imencode('.png', img_arr) img = img.tobytes() scan = kitti_utils.load_scan(scan_file) calib = kitti_utils.load_calib(calib_file) label_k = kitti_utils.load_label(label_file) if label_k.shape[0] == 0: continue scan = kitti_utils.scan_to_camera_coords(scan, calib) _, scan = helpers.get_fixed_pts(scan, cfg['n_points']) label = {} n_inst = label_k.shape[0] label['calib'] = calib['P2'] label['orig'] = np.array(orig_img_size).astype(np.float32) label['clf'] = label_k[:, 0, None].astype(int) label['clf'] = np.pad(label['clf'], [[0, max_objects-n_inst],[0, 0]], 'constant', constant_values=8) label['c_3d'] = label_k[:, 11:14] label['c_3d'] = np.pad(label['c_3d'], [[0, max_objects-n_inst],[0, 0]]) label['extent'] = label_k[:, 8:11] label['extent'] = np.pad(label['extent'], [[0, max_objects-n_inst],[0, 0]]) label['bbox_3d'] = kitti_utils.extract_3d_boxes(label_k) c_2d, bbox_2d = kitti_utils.extract_2d_boxes(label['bbox_3d'], calib, orig_img_size) label['c_2d'] = np.pad(c_2d, [[0, max_objects-n_inst],[0, 0]]).astype(np.float32) label['bbox_2d'] = np.pad(bbox_2d, [[0, max_objects-n_inst],[0, 0]]).astype(np.float32) label['bbox_3d'] = np.concatenate( [label['bbox_3d'], np.zeros((max_objects-n_inst, 8, 3))] ) ri = np.cos(label_k[:, 14]) rj = np.sin(label_k[:, 14]) label['ri'] = np.pad(ri.reshape(-1, 1), [[0, max_objects-n_inst], [0, 0]]) label['rj'] = np.pad(rj.reshape(-1, 1), [[0, max_objects-n_inst], [0, 0]]) label = kitti_utils.remove_dontcare(label) tf_example = create_example(img, scan, label) if scene_id % 5 != 0: train_writer.write(tf_example.SerializeToString()) else: val_writer.write(tf_example.SerializeToString())
def crop_s3dis(): filelist = glob.glob(os.path.join(config['in_dir'], '*.npy')) box_size = config['box_size'] overlap = config['overlap'] saved = 0 with tf.io.TFRecordWriter( config['out_train_file']) as train_writer, tf.io.TFRecordWriter( config['out_test_file']) as test_writer: bar = helpers.progbar(len(filelist)) bar.start() max_labels = 0 rotations = np.radians(np.array( [0, 90, 180, 270])) if config['rotate'] == True else np.array([0.]) for i, f in enumerate(filelist): bar.update(i + 1) scene = np.load(f) area = '_'.join(f.split('/')[-1].split('_')[:2]) room = '_'.join(f.split('/')[-1].split('.')[0].split('_')[2:]) area_n = int(f.split('/')[-1].split('_')[1]) object_paths = glob.glob( os.path.join(config['root_dir'], area, room, 'Annotations', '*{}*.npy'.format(config['label_object']))) objects = np.array([np.load(o_f)[:, :3] for o_f in object_paths]) object_means_orig = np.array([np.mean(o, axis=0) for o in objects]) if object_means_orig.shape[0] == 0: continue object_thetas_orig, object_extents = helpers.get_oabb(objects) area = int(f.split('/')[-1].split('_')[1]) scene_extent = [ np.min(scene[:, 0]), np.min(scene[:, 1]), np.min(scene[:, 2]), np.max(scene[:, 0]), np.max(scene[:, 1]), np.max(scene[:, 2]) ] x_stride_len = box_size[0] y_stride_len = box_size[1] num_xstrides = int( np.ceil((scene_extent[3] - scene_extent[0]) / box_size[0])) num_ystrides = int( np.ceil((scene_extent[4] - scene_extent[1]) / box_size[1])) for x_stride in range(num_xstrides): for y_stride in range(num_ystrides): bbox = [ scene_extent[0] + (x_stride * x_stride_len) - overlap[0] / 2, scene_extent[1] + (y_stride * y_stride_len) - overlap[0] / 2, -1e10, scene_extent[0] + ((x_stride * x_stride_len) + x_stride_len) + overlap[0] / 2, scene_extent[1] + ((y_stride * y_stride_len) + y_stride_len) + overlap[0] / 2, 1e10 ] scene_crop_orig = helpers.crop_bbox(scene, bbox) if scene_crop_orig.shape[0] < config['n_pts'] / 2: continue for angle in rotations: _, scene_crop = helpers.get_fixed_pts( scene_crop_orig, config['n_pts']) object_means = object_means_orig.copy() object_thetas = object_thetas_orig.copy() scene_crop[:, :3] = helpers.rotate_euler( scene_crop[:, :3], angle) object_means = helpers.rotate_euler( object_means_orig, angle) radians = np.arctan2(object_thetas[:, 1], object_thetas[:, 0]) radians -= angle object_thetas[:, 0] = np.cos(radians) object_thetas[:, 1] = np.sin(radians) pts = scene_crop[:, :3] scene_mean = np.mean(pts, axis=0) pts -= scene_mean colors = scene_crop[:, 3:6] / 255. obj_occ = np.array([ helpers.check_occupancy(obj_pts, bbox) for obj_pts in objects ]) obj_occ[obj_occ < 1000] = 0 obj_occ = obj_occ.astype(np.bool) if True in obj_occ: try: labels = object_means[np.where( obj_occ == True)] labels -= scene_mean labels = np.hstack( (labels, object_extents[np.where( obj_occ == True)])) labels = np.hstack( (labels, object_thetas[np.where(obj_occ == True)])) max_labels = labels.shape[0] if labels.shape[ 0] > max_labels else max_labels labels = np.pad(labels, [[ 0, config['max_labels'] - labels.shape[0] ], [0, 0]]) except: print(labels.shape) continue else: continue # Uncomment to visualise training data # plot(pts, colors, labels) tf_example = create_example(pts, colors, labels) if area_n != 5: train_writer.write(tf_example.SerializeToString()) else: test_writer.write(tf_example.SerializeToString()) saved += 1 bar.finish() print('[info] total scenes generated: {}'.format(saved)) print('[info] max label count: {}'.format(max_labels))
def generate_random_crops(config): crop_size = config['crop_size'] scene_crop_size = config['scene_crop_size'] loss_object = tf.keras.losses.MeanSquaredError() pos_count = 0 neg_count = 0 room_count = 0 n_rooms = np.array([len(os.listdir(os.path.join(config['root_dir'], 'Area_{}'.format(i+1)))) for i in range(6)]).sum() bar = helpers.progbar(n_rooms) bar.start() with tf.io.TFRecordWriter(config['out_train_file']) as train_writer, tf.io.TFRecordWriter(config['out_test_file']) as test_writer: for i in range(1, 7): area_path = os.path.join(config['root_dir'], 'Area_{}'.format(i)) for room_folder in os.listdir(area_path): room_path = os.path.join(area_path, room_folder) if not os.path.isdir(room_path): continue room_count += 1 bar.update(room_count) # You can speed this script up by pre loading .txt files and saving as .npy room_cloud = np.loadtxt(glob.glob(room_path+'/*.txt')[0], delimiter=' ') obj_paths = glob.glob(os.path.join(room_path, 'Annotations', '*{}*.npy'.format(config['object']))) objects = np.array([np.load(obj_path) for obj_path in obj_paths]) object_thetas, object_extents = helpers.get_oabb(objects) obj_means = np.array([np.mean(p[:, :3], axis=0) for p in objects]) for obj_count, obj in enumerate(objects): if i != 5 and np.random.rand() > config['train_ratio']: continue # ignore object obj_mean = obj_means[obj_count] for j in range(config['n_crops_per_object']): crop_x = np.random.uniform(obj_mean[0]-(crop_size[0]/2), obj_mean[0]+(crop_size[0]/2)) crop_y = np.random.uniform(obj_mean[1]-(crop_size[1]/2), obj_mean[1]+(crop_size[1]/2)) crop_z = np.random.uniform(obj_mean[2]-(crop_size[2]/2), obj_mean[2]+(crop_size[2]/2)) scene_bbox = np.array([ crop_x - (scene_crop_size[0]/2.), crop_y - (scene_crop_size[1]/2.), crop_z - 1e10, crop_x + (scene_crop_size[0]/2.), crop_y + (scene_crop_size[1]/2.), crop_z + 1e10, ]) bbox = np.array([ crop_x - (crop_size[0]/2.), crop_y - (crop_size[1]/2.), crop_z - (crop_size[2]/2.), crop_x + (crop_size[0]/2.), crop_y + (crop_size[1]/2.), crop_z + (crop_size[2]/2.), ]) crop_centre = np.array([crop_x, crop_y, crop_z]) l_dist = [helpers.euc_dist([crop_x, crop_y, crop_z], c) for c in obj_means] obj_mean = obj_means[np.argmin(l_dist)] obj_ext = object_extents[np.argmin(l_dist)] obj_theta = object_thetas[np.argmin(l_dist)] curr_obj = objects[np.argmin(l_dist)] norm_bbox = np.array([ bbox[0] - crop_centre[0], bbox[1] - crop_centre[1], bbox[2] - crop_centre[2], bbox[3] - crop_centre[0], bbox[4] - crop_centre[1], bbox[5] - crop_centre[2] ]) label_mean = obj_mean - crop_centre for _ in range(config['n_rotations_per_object']): tmp_room_cloud = helpers.crop_bbox(room_cloud, scene_bbox) _, tmp_room_cloud = helpers.get_fixed_pts(tmp_room_cloud, config['n_scene_pts']) tmp_room_crop = helpers.crop_bbox(tmp_room_cloud, bbox) tmp_room_crop[:, :3] -= crop_centre theta = np.random.uniform() * 2.0 * np.pi tmp_room_crop[:, :3] = helpers.rotate_euler(tmp_room_crop[:, :3], theta) label_mean_r = helpers.rotate_euler(label_mean.reshape(1, 3), theta)[0] new_angle = np.arctan2(obj_theta[1], obj_theta[0]) - theta theta_r = np.array([np.cos(new_angle), np.sin(new_angle)]) tmp_room_crop = helpers.crop_bbox(tmp_room_crop, norm_bbox) check = helpers.check_occupancy(curr_obj[:, :3]-crop_centre, norm_bbox) if check == 0: continue dist = np.min(l_dist) dir = (label_mean_r - np.array([0, 0, 0])) / dist pred_dummy = tf.Variable([0., 0., 0.]) with tf.GradientTape() as tape: tape.watch(pred_dummy) loss = loss_object(tf.Variable(label_mean_r), pred_dummy) grad = tape.gradient(loss, pred_dummy).numpy() ret, tmp_room_crop = helpers.get_fixed_pts(tmp_room_crop, config['n_pts']) if ret == False: continue label = np.array([ 1, dist, grad[0], grad[1], grad[2], obj_ext[0], obj_ext[1], obj_ext[2], theta_r[0], theta_r[1] ]) # Uncomment to visualise training data # plot(tmp_room_crop[:, :3], tmp_room_crop[:, 3:6], -grad, label_mean_r, label) tf_example = create_example(tmp_room_crop[:, :3], tmp_room_crop[:, 3:6], label) assert tmp_room_crop.shape[0] == config['n_pts'], '{}'.format(tmp_room_crop.shape) if i != 5: train_writer.write(tf_example.SerializeToString()) if i == 5: test_writer.write(tf_example.SerializeToString()) pos_count += 1 for j in range(config['n_negative_samples']): for _ in range(1000): crop_x = np.random.uniform(np.min(room_cloud[:, 0]), np.max(room_cloud[:, 0])) crop_y = np.random.uniform(np.min(room_cloud[:, 1]), np.max(room_cloud[:, 1])) crop_z = np.random.uniform(np.min(room_cloud[:, 2]), np.max(room_cloud[:, 2])) crop_point = np.array([crop_x, crop_y, crop_z]) scene_bbox = np.array([ crop_point[0] - (scene_crop_size[0]/2.), crop_point[1] - (scene_crop_size[1]/2.), crop_point[2] - (scene_crop_size[2]/2.), crop_point[0] + (scene_crop_size[0]/2.), crop_point[1] + (scene_crop_size[1]/2.), crop_point[2] + (scene_crop_size[2]/2.), ]) bbox = np.array([ crop_point[0] - (crop_size[0]/2), crop_point[1] - (crop_size[1]/2), crop_point[2] - (crop_size[2]/2), crop_point[0] + (crop_size[0]/2), crop_point[1] + (crop_size[1]/2), crop_point[2] + (crop_size[2]/2), ]) check = np.array([helpers.check_occupancy(obj_pts, bbox) for obj_pts in objects]) check = check.astype(np.bool) if True in check: continue norm_bbox = np.array([ bbox[0]-crop_point[0], bbox[1]-crop_point[1], bbox[2]-crop_point[2], bbox[3]-crop_point[0], bbox[4]-crop_point[1], bbox[5]-crop_point[2] ]) for _ in range(config['n_rotations_per_object']): tmp_room_cloud = helpers.crop_bbox(room_cloud, scene_bbox) ret, tmp_room_cloud, = helpers.get_fixed_pts(tmp_room_cloud, config['n_scene_pts']) room_crop = helpers.crop_bbox(tmp_room_cloud, bbox) room_crop[:, :3] -= crop_point[:3] tmp_room_crop = room_crop.copy() theta = np.random.uniform() * 2.0 * np.pi tmp_room_crop[:, :3] = helpers.rotate_euler(tmp_room_crop[:, :3], theta) tmp_room_crop = helpers.crop_bbox(tmp_room_crop, norm_bbox) ret, tmp_room_crop = helpers.get_fixed_pts(tmp_room_crop, config['n_pts']) null_value = 0.0 label = np.array([ 0, 10, null_value, null_value, null_value, null_value, null_value, null_value, null_value, null_value ]) tf_example = create_example(tmp_room_crop[:, :3], tmp_room_crop[:, 3:6], label) assert tmp_room_crop.shape[0] == config['n_pts'], '{}'.format(tmp_room_crop.shape[0].shape) if i != 5: train_writer.write(tf_example.SerializeToString()) if i == 5: test_writer.write(tf_example.SerializeToString()) neg_count += 1 break bar.finish() print('[info] total scenes saved: {}, {} positive, {} negative'.format( pos_count+neg_count, pos_count, neg_count) )
def crop_scannet(config): scan_dirs = [x[0] for x in os.walk(config['in_dir'])] box_size = config['box_size'] overlap = config['overlap'] max_labels = 0 with tf.io.TFRecordWriter( config['out_train_file']) as train_writer, tf.io.TFRecordWriter( config['out_test_file']) as test_writer: bar = helpers.progbar(len(scan_dirs[1:])) bar.start() save_count = 0 train_count = 0 test_count = 0 max_labels = 0 rotations = np.radians(np.array( [0, 90, 180, 270])) if config['rotate'] == True else np.array([0.]) for i, scan_dir in enumerate(scan_dirs[1:]): bar.update(i + 1) scene_id = scan_dir.split('/')[-1] pts_path = os.path.join( scan_dir, '{}_vh_clean_2.labels.ply'.format(scene_id)) agg_path = os.path.join( scan_dir, '{}_vh_clean.aggregation.json'.format(scene_id)) segmap_path = os.path.join( scan_dir, '{}_vh_clean_2.0.010000.segs.json'.format(scene_id)) with open(agg_path, 'r') as f: agg = json.load(f) segmap = su.load_segmap(segmap_path) scene = su.load_pointcloud(pts_path) objects, object_means_orig = su.load_objects( scene, agg, segmap, config['label_objects']) if object_means_orig.shape[0] == 0: continue object_thetas_orig, object_extents = helpers.get_oabb(objects) scene_extent = [ np.min(scene[:, 0]), np.min(scene[:, 1]), np.min(scene[:, 2]), np.max(scene[:, 0]), np.max(scene[:, 1]), np.max(scene[:, 2]) ] x_stride_len = box_size[0] y_stride_len = box_size[1] num_xstrides = int( np.ceil((scene_extent[3] - scene_extent[0]) / box_size[0])) num_ystrides = int( np.ceil((scene_extent[4] - scene_extent[1]) / box_size[1])) for x_stride in range(num_xstrides): for y_stride in range(num_ystrides): bbox = [ scene_extent[0] + (x_stride * x_stride_len) - overlap[0] / 2, scene_extent[1] + (y_stride * y_stride_len) - overlap[0] / 2, -1e10, scene_extent[0] + ((x_stride * x_stride_len) + x_stride_len) + overlap[0] / 2, scene_extent[1] + ((y_stride * y_stride_len) + y_stride_len) + overlap[0] / 2, 1e10 ] scene_crop_orig = helpers.crop_bbox(scene, bbox) if scene_crop_orig.shape[0] < config['n_pts'] / 2: continue for angle in rotations: _, scene_crop = helpers.get_fixed_pts( scene_crop_orig, config['n_pts']) object_means = object_means_orig.copy() object_thetas = object_thetas_orig.copy() scene_crop[:, :3] = helpers.rotate_euler( scene_crop[:, :3], angle) object_means = helpers.rotate_euler( object_means_orig, angle) radians = np.arctan2(object_thetas[:, 1], object_thetas[:, 0]) radians -= angle object_thetas[:, 0] = np.cos(radians) object_thetas[:, 1] = np.sin(radians) pts = scene_crop[:, :3] scene_mean = np.mean(pts, axis=0) pts -= scene_mean colors = scene_crop[:, 3:6] / 255. obj_occ = np.array([ helpers.check_occupancy(obj_pts, bbox) for obj_pts in objects ]) obj_occ[obj_occ < 1000] = 0 obj_occ = obj_occ.astype(np.bool) if True in obj_occ: labels = object_means[np.where(obj_occ == True)] labels -= scene_mean labels = np.hstack( (labels, object_extents[np.where(obj_occ == True)])) labels = np.hstack( (labels, object_thetas[np.where(obj_occ == True)])) max_labels = labels.shape[0] if labels.shape[ 0] > max_labels else max_labels labels = np.pad( labels, [[0, config['max_labels'] - labels.shape[0]], [0, 0]]) else: continue # Uncomment to visualise training data # plot(pts, colors, labels) tf_example = create_example(pts, colors, labels) if i % 5 != 0: train_writer.write(tf_example.SerializeToString()) train_count += 1 else: test_writer.write(tf_example.SerializeToString()) test_count += 1 save_count += 1 bar.finish() print('[info] saved {} examples, {} train, {} test'.format( save_count, train_count, test_count)) print('[info] max label count: {}'.format(max_labels))