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 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_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 parse_room():
model = ObjDetectorModel(1, config['n_pred'])
model(tf.zeros((1, config['n_pts'], 3), tf.float32))
model.load_weights(config['weights'])
if config['dataset'] == 's3d':
room = 'Area_' + str(config['area']) + '_' + config['room'] + '.npy'
scene = np.load(os.path.join(config['dataset_dir'], 'processed', room))
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])
]
object_paths = glob.glob(
os.path.join(config['dataset_dir'], 'Area_' + str(config['area']),
config['room'], 'Annotations', '*chair*.npy'))
objects = np.array([np.load(o_f)[:, :3] for o_f in object_paths])
gt_pts = np.array([np.mean(o, axis=0) for o in objects])
gt_theta, gt_ext = helpers.get_oabb(objects)
gt_ext = np.hstack((gt_ext, gt_theta))
x_stride_len = config['box_size'][0]
y_stride_len = config['box_size'][1]
num_xstrides = int(
np.ceil((scene_extent[3] - scene_extent[0]) / x_stride_len))
num_ystrides = int(
np.ceil((scene_extent[4] - scene_extent[1]) / y_stride_len))
scene_pts = []
scene_ext = []
for x_stride in range(num_xstrides):
for y_stride in range(num_ystrides):
bbox = [
scene_extent[0] + (x_stride * x_stride_len),
scene_extent[1] + (y_stride * y_stride_len), -1e10,
scene_extent[0] + ((x_stride * x_stride_len) + x_stride_len),
scene_extent[1] + ((y_stride * y_stride_len) + y_stride_len),
1e10
]
scene_crop = helpers.crop_bbox(scene, bbox)
_, scene_crop = helpers.get_fixed_pts(scene_crop, config['n_pts'])
pts = scene_crop[:, 0:3]
cols = scene_crop[:, 3:6] / 256.
pts_mean = np.mean(pts, axis=0)
pts -= pts_mean
xyz, score, ext, _ = model(tf.expand_dims(pts, 0))
xyz, score, ext = tf_utils.objectness_mask(xyz, score[:, :,
1], ext,
config['score_thresh'])
labels = gt_pts[(gt_pts[:, 0] >= bbox[0])
& (gt_pts[:, 0] <= bbox[3]) &
(gt_pts[:, 1] >= bbox[1]) &
(gt_pts[:, 1] <= bbox[4]) &
(gt_pts[:, 2] >= bbox[2]) &
(gt_pts[:, 2] <= bbox[5])]
labels = tf.expand_dims(labels - pts_mean, 0)
if xyz.shape[1] > 1 and config['nms'] == True:
boxes = helpers.make_boxes(xyz, ext)
nms_inds = tf_utils.nms(xyz, boxes, score, 15, 0,
config['score_thresh'])
xyz = tf.gather(xyz, nms_inds, batch_dims=1)
score = tf.gather(score, nms_inds, batch_dims=1)
boxes = tf.gather(boxes, nms_inds, batch_dims=1)
ext = tf.gather(ext, nms_inds, batch_dims=1)
if xyz.shape[1] > 0:
for i, pred in enumerate(xyz[0]):
scene_pts.append((pred.numpy() + pts_mean).tolist())
scene_ext.append(ext[0, i].numpy().tolist())
scene_pts = np.array(scene_pts)
scene_ext = np.array(scene_ext)
plot_scene(scene, scene_pts, scene_ext, gt_pts, gt_ext)
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))