def test(args): inputs = tf.placeholder(tf.float32, (1, None, 3)) npts = tf.placeholder(tf.int32, (1, )) print('AAAAAAAAAAAAAAAA') print(args.num_gt_points) print('AAAAAAAAAAAAAAAAAAAAAAAAAa') gt = tf.placeholder(tf.float32, (1, args.num_gt_points, 3)) model_module = importlib.import_module('.%s' % args.model_type, 'models') model = model_module.Model(inputs, npts, gt, tf.constant(1.0)) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True sess = tf.Session(config=config) saver = tf.train.Saver() saver.restore(sess, args.checkpoint) t0 = time.time() partial = read_pcd(args.pcd_file) bbox = np.loadtxt(args.bbox_file) # Calculate center, rotation and scale center = (bbox.min(0) + bbox.max(0)) / 2 bbox -= center yaw = np.arctan2(bbox[3, 1] - bbox[0, 1], bbox[3, 0] - bbox[0, 0]) rotation = np.array([[np.cos(yaw), -np.sin(yaw), 0], [np.sin(yaw), np.cos(yaw), 0], [0, 0, 1]]) bbox = np.dot(bbox, rotation) scale = bbox[3, 0] - bbox[0, 0] bbox /= scale partial = np.dot(partial - center, rotation) / scale partial = np.dot(partial, [[1, 0, 0], [0, 0, 1], [0, 1, 0]]) start = time.time() completion = sess.run(model.outputs, feed_dict={ inputs: [partial], npts: [partial.shape[0]] }) completion = completion[0] completion_w = np.dot(completion, [[1, 0, 0], [0, 0, 1], [0, 1, 0]]) completion_w = np.dot(completion_w * scale, rotation.T) + center dummy = args.pcd_file.split('/')[-1] result_path = os.path.join(args.results_dir, dummy) save_pcd(result_path, completion_w) plot_path = os.path.join(args.results_dir, 'plots.png') plot_pcd_three_views(plot_path, [partial, completion], ['input', 'output'], '%d input points' % partial.shape[0], [5, 0.5]) sess.close() tf = time.time() print('Total time: {}'.format(tf - t0))
def test(args): inputs = tf.placeholder(tf.float32, (1, None, 3)) npts = tf.placeholder(tf.int32, (1,)) gt = tf.placeholder(tf.float32, (1, args.num_gt_points, 3)) model_module = importlib.import_module('.%s' % args.model_type, 'models') model = model_module.Model(inputs, npts, gt, tf.constant(1.0)) os.makedirs(os.path.join(args.results_dir, 'plots'), exist_ok=True) os.makedirs(os.path.join(args.results_dir, 'completions'), exist_ok=True) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True sess = tf.Session(config=config) saver = tf.train.Saver() saver.restore(sess, args.checkpoint) car_ids = [filename.split('.')[0] for filename in os.listdir(args.pcd_dir)] total_time = 0 total_points = 0 for i, car_id in enumerate(car_ids): partial = read_pcd(os.path.join(args.pcd_dir, '%s.pcd' % car_id)) bbox = np.loadtxt(os.path.join(args.bbox_dir, '%s.txt' % car_id)) total_points += partial.shape[0] # Calculate center, rotation and scale center = (bbox.min(0) + bbox.max(0)) / 2 bbox -= center yaw = np.arctan2(bbox[3, 1] - bbox[0, 1], bbox[3, 0] - bbox[0, 0]) rotation = np.array([[np.cos(yaw), -np.sin(yaw), 0], [np.sin(yaw), np.cos(yaw), 0], [0, 0, 1]]) bbox = np.dot(bbox, rotation) scale = bbox[3, 0] - bbox[0, 0] bbox /= scale partial = np.dot(partial - center, rotation) / scale partial = np.dot(partial, [[1, 0, 0], [0, 0, 1], [0, 1, 0]]) start = time.time() completion = sess.run(model.outputs, feed_dict={inputs: [partial], npts: [partial.shape[0]]}) total_time += time.time() - start completion = completion[0] completion_w = np.dot(completion, [[1, 0, 0], [0, 0, 1], [0, 1, 0]]) completion_w = np.dot(completion_w * scale, rotation.T) + center pcd_path = os.path.join(args.results_dir, 'completions', '%s.pcd' % car_id) save_pcd(pcd_path, completion_w) if i % args.plot_freq == 0: plot_path = os.path.join(args.results_dir, 'plots', '%s.png' % car_id) plot_pcd_three_views(plot_path, [partial, completion], ['input', 'output'], '%d input points' % partial.shape[0], [5, 0.5]) print('Average # input points:', total_points / len(car_ids)) print('Average time:', total_time / len(car_ids)) sess.close()
def test(args): inputs = tf.placeholder(tf.float32, (1, None, 3)) npts = tf.placeholder(tf.int32, (1, )) gt = tf.placeholder(tf.float32, (1, args.num_gt_points, 3)) model_module = importlib.import_module('.%s' % args.model_type, 'models') model = model_module.Model(inputs, npts, gt, tf.constant(1.0)) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True sess = tf.Session(config=config) saver = tf.train.Saver() saver.restore(sess, args.checkpoint) an = Analyzer(dr.get_PLIDAR_predicted_path) argument_list = an.get_datatype_trackno_carno(data_types=["train"]) for i in range(len(argument_list)): data_type, track_no, car_no = argument_list[i] # Input - partial point cloud .pcd pcd_file = \ dr.get_pcn_lidar_reference_partial_path(data_type, track_no, car_no, extension='pcd') # Input - bbox 8 corners .txt bbox_file = dr.get_pcn_bbox_lidar_path(data_type, track_no, car_no) # Result - complete point cloud .pcd result_path = dr.get_pcn_lidar_reference_complete_path( data_type, track_no, car_no) # Result - plot of input and output of pointcloud plot_path = dr.get_pcn_plot_lidar_path(data_type, track_no, car_no) partial = read_pcd(pcd_file) bbox = np.loadtxt(bbox_file) # Calculate center, rotation and scale center = (bbox.min(0) + bbox.max(0)) / 2 bbox -= center yaw = np.arctan2(bbox[3, 1] - bbox[0, 1], bbox[3, 0] - bbox[0, 0]) rotation = np.array([[np.cos(yaw), -np.sin(yaw), 0], [np.sin(yaw), np.cos(yaw), 0], [0, 0, 1]]) bbox = np.dot(bbox, rotation) scale = bbox[3, 0] - bbox[0, 0] bbox /= scale partial = np.dot(partial - center, rotation) / scale partial = np.dot(partial, [[1, 0, 0], [0, 0, 1], [0, 1, 0]]) completion = sess.run(model.outputs, feed_dict={ inputs: [partial], npts: [partial.shape[0]] }) completion = completion[0] completion_w = np.dot(completion, [[1, 0, 0], [0, 0, 1], [0, 1, 0]]) completion_w = np.dot(completion_w * scale, rotation.T) + center save_pcd(result_path, completion_w) plot_pcd_three_views(plot_path, [partial, completion], ['input', 'output'], '%d input points' % partial.shape[0], [5, 0.5]) print('Finish {}/{}'.format(i + 1, len(argument_list))) sess.close()
def train(args): min_loss_fine = 1.0 is_training_pl = tf.placeholder(tf.bool, shape=(), name='is_training') global_step = tf.Variable(0, trainable=False, name='global_step') alpha = tf.train.piecewise_constant(global_step, [10000, 20000, 50000], [0.01, 0.1, 0.5, 1.0], 'alpha_op') provider = TrainProvider(args, is_training_pl) ids, inputs, gt = provider.batch_data num_eval_steps = provider.num_valid // args.batch_size model_module = importlib.import_module('.%s' % args.model_type, 'models') model = model_module.Model(inputs, gt, alpha, is_training_pl) add_train_summary('alpha', alpha) if args.lr_decay: learning_rate = tf.train.exponential_decay(args.base_lr, global_step, args.lr_decay_steps, args.lr_decay_rate, staircase=True, name='lr') learning_rate = tf.maximum(learning_rate, args.lr_clip) add_train_summary('learning_rate', learning_rate) else: learning_rate = tf.constant(args.base_lr, name='lr') trainer = tf.train.AdamOptimizer(learning_rate) train_op = trainer.minimize(model.loss, global_step) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True sess = tf.Session(config=config) sess.run(tf.global_variables_initializer()) saver = tf.train.Saver(max_to_keep=15) if args.restore: saver.restore(sess, tf.train.latest_checkpoint(args.log_dir)) #saver.restore(sess, 'data/trained_models/pcn_cd') else: if os.path.exists(args.log_dir): delete_key = input( colored('%s exists. Delete? [y (or enter)/N]' % args.log_dir, 'white', 'on_red')) if delete_key == 'y' or delete_key == "": os.system('rm -rf %s/*' % args.log_dir) os.makedirs(os.path.join(args.log_dir, 'plots')) else: os.makedirs(os.path.join(args.log_dir, 'plots')) with open(os.path.join(args.log_dir, 'args.txt'), 'w') as log: for arg in sorted(vars(args)): log.write(arg + ': ' + str(getattr(args, arg)) + '\n') # log of arguments os.system('cp models/%s.py %s' % (args.model_type, args.log_dir)) # bkp of model def os.system('cp train.py %s' % args.log_dir) # bkp of train procedure train_summary = tf.summary.merge_all('train_summary') valid_summary = tf.summary.merge_all('valid_summary') writer = tf.summary.FileWriter(args.log_dir, sess.graph) coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess=sess, coord=coord) total_time = 0 train_start = time.time() step = sess.run(global_step) while not coord.should_stop(): step += 1 epoch = step * args.batch_size // provider.num_train + 1 start = time.time() _, loss, loss_fine, summary = sess.run( [train_op, model.loss, model.loss_fine, train_summary], feed_dict={is_training_pl: True}) total_time += time.time() - start writer.add_summary(summary, step) if step % args.steps_per_print == 0: print( 'epoch %d step %d loss %.8f loss_fine %.8f - time per batch %.4f' % (epoch, step, loss, loss_fine, total_time / args.steps_per_print)) total_time = 0 if step < 100000: steps_per_eval = args.steps_per_eval * 10 else: steps_per_eval = args.steps_per_eval if step % steps_per_eval == 0: print(colored('Testing...', 'grey', 'on_green')) total_loss = 0 total_time = 0 total_loss_fine = 0 sess.run(tf.local_variables_initializer()) for i in range(num_eval_steps): start = time.time() loss, loss_fine, _ = sess.run( [model.loss, model.loss_fine, model.update], feed_dict={is_training_pl: False}) total_loss += loss total_loss_fine += loss_fine total_time += time.time() - start summary = sess.run(valid_summary, feed_dict={is_training_pl: False}) writer.add_summary(summary, step) print( colored( 'epoch %d step %d loss %.8f loss_fine %.8f - time per batch %.4f' % (epoch, step, total_loss / num_eval_steps, total_loss_fine / num_eval_steps, total_time / num_eval_steps), 'grey', 'on_green')) total_time = 0 if (total_loss_fine / num_eval_steps) < min_loss_fine: min_loss_fine = total_loss_fine / num_eval_steps saver.save(sess, os.path.join(args.log_dir, 'model'), step) print( colored('Model saved at %s' % args.log_dir, 'white', 'on_blue')) if step % args.steps_per_visu == 0: model_id, pcds = sess.run([ids[0], model.visualize_ops], feed_dict={is_training_pl: True}) model_id = model_id.decode('utf-8') plot_path = os.path.join( args.log_dir, 'plots', 'epoch_%d_step_%d_%s.png' % (epoch, step, model_id)) plot_pcd_three_views(plot_path, pcds, model.visualize_titles) #if step % args.steps_per_save == 0: # saver.save(sess, os.path.join(args.log_dir, 'model'), step) # print(colored('Model saved at %s' % args.log_dir, 'white', 'on_blue')) if step >= args.max_step: break print('Total time', datetime.timedelta(seconds=time.time() - train_start)) coord.request_stop() coord.join(threads) sess.close()
def test(args): inputs = tf.placeholder(tf.float32, (1, None, 3)) gt = tf.placeholder(tf.float32, (1, args.num_gt_points, 3)) model_module = importlib.import_module('.%s' % args.model_type, 'models') model = model_module.Model(inputs, gt, tf.constant(1.0)) output = tf.placeholder(tf.float32, (1, args.num_gt_points, 3)) cd_op = chamfer(output, gt) emd_op = earth_mover(output, gt) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True sess = tf.Session(config=config) saver = tf.train.Saver() saver.restore(sess, args.checkpoint) os.makedirs(args.results_dir, exist_ok=True) csv_file = open(os.path.join(args.results_dir, 'results.csv'), 'w') writer = csv.writer(csv_file) writer.writerow(['id', 'cd', 'emd']) with open(args.list_path) as file: model_list = file.read().splitlines() total_time = 0 total_cd = 0 total_emd = 0 cd_per_cat = {} emd_per_cat = {} for i, model_id in enumerate(model_list): partial = read_pcd( os.path.join(args.data_dir, 'partial', '%s.pcd' % model_id)) complete = read_pcd( os.path.join(args.data_dir, 'complete', '%s.pcd' % model_id)) start = time.time() completion = sess.run(model.outputs, feed_dict={inputs: [partial]}) total_time += time.time() - start cd, emd = sess.run([cd_op, emd_op], feed_dict={ output: completion, gt: [complete] }) total_cd += cd total_emd += emd writer.writerow([model_id, cd, emd]) synset_id, model_id = model_id.split('/') if not cd_per_cat.get(synset_id): cd_per_cat[synset_id] = [] if not emd_per_cat.get(synset_id): emd_per_cat[synset_id] = [] cd_per_cat[synset_id].append(cd) emd_per_cat[synset_id].append(emd) if i % args.plot_freq == 0: os.makedirs(os.path.join(args.results_dir, 'plots', synset_id), exist_ok=True) plot_path = os.path.join(args.results_dir, 'plots', synset_id, '%s.png' % model_id) plot_pcd_three_views(plot_path, [partial, completion[0], complete], ['input', 'output', 'ground truth'], 'CD %.4f EMD %.4f' % (cd, emd), [5, 0.5, 0.5]) if args.save_pcd: os.makedirs(os.path.join(args.results_dir, 'pcds', synset_id), exist_ok=True) save_pcd( os.path.join(args.results_dir, 'pcds', '%s.pcd' % model_id), completion[0]) csv_file.close() sess.close() print('Average time: %f' % (total_time / len(model_list))) print('Average Chamfer distance: %f' % (total_cd / len(model_list))) print('Average Earth mover distance: %f' % (total_emd / len(model_list))) print('Chamfer distance per category') for synset_id in cd_per_cat.keys(): print(synset_id, '%f' % np.mean(cd_per_cat[synset_id])) print('Earth mover distance per category') for synset_id in emd_per_cat.keys(): print(synset_id, '%f' % np.mean(emd_per_cat[synset_id]))
def train(args): is_training_pl = tf.placeholder(tf.bool, shape=(), name='is_training') global_step = tf.Variable(0, trainable=False, name='global_step') alpha = tf.train.piecewise_constant(global_step, [10000, 20000, 50000], [0.01, 0.1, 0.5, 1.0], 'alpha_op') inputs_pl = tf.placeholder(tf.float32, (1, None, 3), 'inputs') npts_pl = tf.placeholder(tf.int32, (args.batch_size, ), 'num_points') gt_pl = tf.placeholder(tf.float32, (args.batch_size, args.num_gt_points, 3), 'ground_truths') model_module = importlib.import_module('.%s' % args.model_type, 'models') model = model_module.Model(inputs_pl, npts_pl, gt_pl, alpha) add_train_summary('alpha', alpha) if args.lr_decay: learning_rate = tf.train.exponential_decay(args.base_lr, global_step, args.lr_decay_steps, args.lr_decay_rate, staircase=True, name='lr') learning_rate = tf.maximum(learning_rate, args.lr_clip) add_train_summary('learning_rate', learning_rate) else: learning_rate = tf.constant(args.base_lr, name='lr') train_summary = tf.summary.merge_all('train_summary') valid_summary = tf.summary.merge_all('valid_summary') trainer = tf.train.AdamOptimizer(learning_rate) train_op = trainer.minimize(model.loss, global_step) df_train, num_train = lmdb_dataflow(args.lmdb_train, args.batch_size, args.num_input_points, args.num_gt_points, is_training=True) train_gen = df_train.get_data() df_valid, num_valid = lmdb_dataflow(args.lmdb_valid, args.batch_size, args.num_input_points, args.num_gt_points, is_training=False) valid_gen = df_valid.get_data() config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True sess = tf.Session(config=config) saver = tf.train.Saver() if args.restore: saver.restore(sess, tf.train.latest_checkpoint(args.log_dir)) writer = tf.summary.FileWriter(args.log_dir) else: sess.run(tf.global_variables_initializer()) if os.path.exists(args.log_dir): delete_key = input( colored('%s exists. Delete? [y (or enter)/N]' % args.log_dir, 'white', 'on_red')) if delete_key == 'y' or delete_key == "": os.system('rm -rf %s/*' % args.log_dir) os.makedirs(os.path.join(args.log_dir, 'plots')) else: os.makedirs(os.path.join(args.log_dir, 'plots')) with open(os.path.join(args.log_dir, 'args.txt'), 'w') as log: for arg in sorted(vars(args)): log.write(arg + ': ' + str(getattr(args, arg)) + '\n') # log of arguments os.system('cp models/%s.py %s' % (args.model_type, args.log_dir)) # bkp of model def os.system('cp train.py %s' % args.log_dir) # bkp of train procedure writer = tf.summary.FileWriter(args.log_dir, sess.graph) total_time = 0 train_start = time.time() init_step = sess.run(global_step) for step in range(init_step + 1, args.max_step + 1): epoch = step * args.batch_size // num_train + 1 ids, inputs, npts, gt = next(train_gen) start = time.time() feed_dict = { inputs_pl: inputs, npts_pl: npts, gt_pl: gt, is_training_pl: True } _, loss, summary = sess.run([train_op, model.loss, train_summary], feed_dict=feed_dict) total_time += time.time() - start writer.add_summary(summary, step) if step % args.steps_per_print == 0: print('epoch %d step %d loss %.8f - time per batch %.4f' % (epoch, step, loss, total_time / args.steps_per_print)) total_time = 0 if step % args.steps_per_eval == 0: print(colored('Testing...', 'grey', 'on_green')) num_eval_steps = num_valid // args.batch_size total_loss = 0 total_time = 0 sess.run(tf.local_variables_initializer()) for i in range(num_eval_steps): start = time.time() ids, inputs, npts, gt = next(valid_gen) feed_dict = { inputs_pl: inputs, npts_pl: npts, gt_pl: gt, is_training_pl: False } loss, _ = sess.run([model.loss, model.update], feed_dict=feed_dict) total_loss += loss total_time += time.time() - start summary = sess.run(valid_summary, feed_dict={is_training_pl: False}) writer.add_summary(summary, step) print( colored( 'epoch %d step %d loss %.8f - time per batch %.4f' % (epoch, step, total_loss / num_eval_steps, total_time / num_eval_steps), 'grey', 'on_green')) total_time = 0 if step % args.steps_per_visu == 0: all_pcds = sess.run(model.visualize_ops, feed_dict=feed_dict) for i in range(0, args.batch_size, args.visu_freq): plot_path = os.path.join( args.log_dir, 'plots', 'epoch_%d_step_%d_%s.png' % (epoch, step, ids[i])) pcds = [x[i] for x in all_pcds] plot_pcd_three_views(plot_path, pcds, model.visualize_titles) if step % args.steps_per_save == 0: saver.save(sess, os.path.join(args.log_dir, 'model'), step) print( colored('Model saved at %s' % args.log_dir, 'white', 'on_blue')) print('Total time', datetime.timedelta(seconds=time.time() - train_start)) sess.close()
def train(args): is_training_pl = tf.placeholder(tf.bool, shape=(), name='is_training') global_step = tf.Variable(0, trainable=False, name='global_step') provider = TrainProvider(args, is_training_pl) ids, inputs, gt_pose = provider.batch_data num_eval_steps = provider.num_valid // args.batch_size model = itn.ITN(inputs, gt_pose, args.iterations, args.validation_iterations, args.no_batchnorm, args.rot_representation, is_training_pl) if not args.no_lr_decay: learning_rate = tf.train.exponential_decay(args.base_lr, global_step, args.lr_decay_steps, args.lr_decay_rate, staircase=True, name='lr') learning_rate = tf.maximum(learning_rate, args.lr_clip) add_train_summary('learning_rate', learning_rate) else: learning_rate = tf.constant(args.base_lr, name='lr') trainer = tf.train.AdamOptimizer(learning_rate) train_op = trainer.minimize(model.loss, global_step) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True sess = tf.Session(config=config) sess.run(tf.global_variables_initializer()) # sess = tf_debug.LocalCLIDebugWrapperSession(sess) saver = tf.train.Saver() now = datetime.datetime.now(pytz.timezone('US/Pacific')) hr = '{:02d}'.format(now.hour) mn = '{:02d}'.format(now.minute) dy = '{:02d}'.format(now.day) mt = '{:02d}'.format(now.month) yr = '{:04d}'.format(now.year) log_dir = args.log_dir + '_'.join(['', hr, mn, dy, mt, yr]) if args.restore: saver.restore(sess, tf.train.latest_checkpoint(log_dir)) else: if os.path.exists(log_dir): delete_key = input( colored('%s exists. Delete? [y (or enter)/N]' % log_dir, 'white', 'on_red')) if delete_key == 'y' or delete_key == "": os.system('rm -rf %s' % log_dir) os.makedirs(log_dir) os.makedirs(os.path.join(log_dir, 'plots')) else: os.makedirs(os.path.join(log_dir, 'plots')) with open(os.path.join(log_dir, 'args.txt'), 'w') as log: for arg in sorted(vars(args)): log.write(arg + ': ' + str(getattr(args, arg)) + '\n') # log of arguments os.system('cp models/itn.py %s' % (log_dir)) # bkp of model def os.system('cp train.py %s' % log_dir) # bkp of train procedure writer = tf.summary.FileWriter(log_dir, sess.graph) train_summary = tf.summary.merge_all('train_summary') valid_summary = tf.summary.merge_all('valid_summary') writer = tf.summary.FileWriter(log_dir, sess.graph) coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess=sess, coord=coord) total_time = 0 train_start = time.time() step = sess.run(global_step) while not coord.should_stop(): step += 1 epoch = step * args.batch_size // provider.num_train + 1 start = time.time() __, loss, summary = sess.run([train_op, model.loss, train_summary], feed_dict={is_training_pl: True}) total_time += time.time() - start writer.add_summary(summary, step) if step % args.steps_per_print == 0: print('epoch %d step %d loss %.8f - time per batch %.4f' % (epoch, step, loss, total_time / args.steps_per_print)) total_time = 0 if step % args.steps_per_eval == 0: print(colored('Testing...', 'grey', 'on_green')) total_loss = 0 total_time = 0 sess.run(tf.local_variables_initializer()) for i in range(num_eval_steps): start = time.time() loss, _ = sess.run([model.loss, model.update], feed_dict={is_training_pl: False}) total_loss += loss total_time += time.time() - start summary = sess.run(valid_summary, feed_dict={is_training_pl: False}) writer.add_summary(summary, step) print( colored( 'epoch %d step %d loss %.8f - time per batch %.4f' % (epoch, step, total_loss / num_eval_steps, total_time / num_eval_steps), 'grey', 'on_green')) total_time = 0 if step % args.steps_per_visu == 0: model_id, pcds = sess.run([ids[0], model.visualize_ops], feed_dict={is_training_pl: True}) model_id = model_id.decode('utf-8') plot_path = os.path.join( log_dir, 'plots', 'epoch_%d_step_%d_%s.png' % (epoch, step, model_id)) plot_pcd_three_views(plot_path, pcds, model.visualize_titles) if step % args.steps_per_save == 0: saver.save(sess, os.path.join(log_dir, 'model'), step) print(colored('Model saved at %s' % log_dir, 'white', 'on_blue')) if step >= args.max_step: break print('Total time', datetime.timedelta(seconds=time.time() - train_start)) coord.request_stop() coord.join(threads) sess.close()
def test(args): inputs = tf.placeholder(tf.float32, (1, None, 3)) npts = tf.placeholder(tf.int32, (1, )) gt = tf.placeholder(tf.float32, (1, args.num_gt_points, 6)) model_module = importlib.import_module('.%s' % args.model_type, 'models') model = model_module.Model(inputs, npts, gt, tf.constant(1.0), args.num_channel) output = tf.placeholder(tf.float32, (1, args.num_gt_points, 3 + args.num_channel)) cd_op = chamfer(output[:, :, 0:3], gt[:, :, 0:3]) emd_op = earth_mover(output[:, :, 0:3], gt[:, :, 0:3]) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True sess = tf.Session(config=config) saver = tf.train.Saver() saver.restore(sess, args.checkpoint) os.makedirs(args.results_dir, exist_ok=True) csv_file = open(os.path.join(args.results_dir, 'results.csv'), 'w') writer = csv.writer(csv_file) writer.writerow(['id', 'cd', 'emd']) with open(args.list_path) as file: model_list = file.read().splitlines() total_time = 0 total_cd = 0 total_emd = 0 cd_per_cat = {} emd_per_cat = {} np.random.seed(1) for i, model_id in enumerate(model_list): if args.experiment == 'shapenet': synset_id, model_id = model_id.split('/') partial = read_pcd( os.path.join(args.data_dir, 'partial', synset_id, '%s.pcd' % model_id)) complete = read_pcd( os.path.join(args.data_dir, 'complete', synset_id, '%s.pcd' % model_id)) elif args.experiment == 'suncg': synset_id = 'all_rooms' partial = read_pcd( os.path.join(args.data_dir, 'pcd_partial', '%s.pcd' % model_id)) complete = read_pcd( os.path.join(args.data_dir, 'pcd_complete', '%s.pcd' % model_id)) if args.rotate: angle = np.random.rand(1) * 2 * np.pi partial = np.stack([ np.cos(angle) * partial[:, 0] - np.sin(angle) * partial[:, 2], partial[:, 1], np.sin(angle) * partial[:, 0] + np.cos(angle) * partial[:, 2] ], axis=-1) complete = np.stack([ np.cos(angle) * complete[:, 0] - np.sin(angle) * complete[:, 2], complete[:, 1], np.sin(angle) * complete[:, 0] + np.cos(angle) * complete[:, 2], complete[:, 3], complete[:, 4], complete[:, 5] ], axis=-1) partial = partial[:, :3] complete = resample_pcd(complete, 16384) start = time.time() completion1, completion2, mesh_out = sess.run( [model.outputs1, model.outputs2, model.gt_can], feed_dict={ inputs: [partial], npts: [partial.shape[0]], gt: [complete] }) completion1[0][:, (3 + args.num_channel):] *= 0 completion2[0][:, (3 + args.num_channel):] *= 0 mesh_out[0][:, (3 + args.num_channel):] *= 0 total_time += time.time() - start # cd, emd = sess.run([cd_op, emd_op], cd, emd = sess.run([cd_op, cd_op], feed_dict={ output: completion2, gt: [complete] }) total_cd += cd total_emd += emd if not cd_per_cat.get(synset_id): cd_per_cat[synset_id] = [] if not emd_per_cat.get(synset_id): emd_per_cat[synset_id] = [] cd_per_cat[synset_id].append(cd) emd_per_cat[synset_id].append(emd) writer.writerow([model_id, cd, emd]) if i % args.plot_freq == 0: os.makedirs( os.path.join(args.results_dir, 'plots', synset_id), exist_ok=True) plot_path = os.path.join(args.results_dir, 'plots', synset_id, '%s.png' % model_id) plot_pcd_three_views( plot_path, [ partial, completion1[0], completion2[0], mesh_out[0], complete ], ['input', 'coarse', 'fine', 'mesh', 'ground truth'], 'CD %.4f EMD %.4f' % (cd, emd), [5, 0.5, 0.5, 0.5, 0.5], num_channel=args.num_channel) if args.save_pcd: os.makedirs( os.path.join(args.results_dir, 'input', synset_id), exist_ok=True) pts_coord = partial[:, 0:3] pts_color = matplotlib.cm.cool((partial[:, 1]))[:, 0:3] # save_pcd(os.path.join(args.results_dir, 'input', synset_id, '%s.ply' % model_id), np.concatenate((pts_coord, pts_color), -1)) pcd = PointCloud() pcd.points = Vector3dVector(pts_coord) pcd.colors = Vector3dVector(pts_color) write_point_cloud( os.path.join(args.results_dir, 'input', synset_id, '%s.ply' % model_id), pcd, write_ascii=True) os.makedirs( os.path.join(args.results_dir, 'output1', synset_id), exist_ok=True) pts_coord = completion1[0][:, 0:3] pts_color = matplotlib.cm.Set1( (np.argmax(completion1[0][:, 3:3 + args.num_channel], -1) + 1) / args.num_channel - 0.5 / args.num_channel)[:, 0:3] # pts_color = matplotlib.cm.tab20((np.argmax(completion1[0][:, 3:3+args.num_channel], -1) + 1)/args.num_channel - 0.5/args.num_channel)[:,0:3] # save_pcd(os.path.join(args.results_dir, 'output1', synset_id, '%s.ply' % model_id), np.concatenate((pts_coord, pts_color), -1)) pcd.points = Vector3dVector(pts_coord) pcd.colors = Vector3dVector(pts_color) write_point_cloud( os.path.join(args.results_dir, 'output1', synset_id, '%s.ply' % model_id), pcd, write_ascii=True) os.makedirs( os.path.join(args.results_dir, 'output2', synset_id), exist_ok=True) pts_coord = completion2[0][:, 0:3] pts_color = matplotlib.cm.Set1( (np.argmax(completion2[0][:, 3:3 + args.num_channel], -1) + 1) / args.num_channel - 0.5 / args.num_channel)[:, 0:3] # pts_color = matplotlib.cm.tab20((np.argmax(completion2[0][:, 3:3+args.num_channel], -1) + 1)/args.num_channel - 0.5/args.num_channel)[:,0:3] # save_pcd(os.path.join(args.results_dir, 'output2', synset_id, '%s.ply' % model_id), np.concatenate((pts_coord, pts_color), -1)) pcd.points = Vector3dVector(pts_coord) pcd.colors = Vector3dVector(pts_color) write_point_cloud( os.path.join(args.results_dir, 'output2', synset_id, '%s.ply' % model_id), pcd, write_ascii=True) ####### os.makedirs( os.path.join(args.results_dir, 'regions', synset_id), exist_ok=True) for idx in range(3, 3 + args.num_channel): val_min = np.min(completion2[0][:, idx]) val_max = np.max(completion2[0][:, idx]) pts_color = 0.8 * matplotlib.cm.Reds( (completion2[0][:, idx] - val_min) / (val_max - val_min))[:, 0:3] pts_color += 0.2 * matplotlib.cm.gist_gray( (completion2[0][:, idx] - val_min) / (val_max - val_min))[:, 0:3] pcd.colors = Vector3dVector(pts_color) write_point_cloud( os.path.join(args.results_dir, 'regions', synset_id, '%s_%s.ply' % (model_id, idx - 3)), pcd, write_ascii=True) os.makedirs( os.path.join(args.results_dir, 'gt', synset_id), exist_ok=True) pts_coord = complete[:, 0:3] if args.experiment == 'shapenet': pts_color = matplotlib.cm.cool(complete[:, 1])[:, 0:3] elif args.experiment == 'suncg': pts_color = matplotlib.cm.Set1(complete[:, 3] - 0.5 / args.num_channel)[:, 0:3] # save_pcd(os.path.join(args.results_dir, 'gt', synset_id, '%s.ply' % model_id), np.concatenate((pts_coord, pts_color), -1)) pcd.points = Vector3dVector(pts_coord) pcd.colors = Vector3dVector(pts_color) write_point_cloud( os.path.join(args.results_dir, 'gt', synset_id, '%s.ply' % model_id), pcd, write_ascii=True) sess.close() print('Average time: %f' % (total_time / len(model_list))) print('Average Chamfer distance: %f' % (total_cd / len(model_list))) print('Average Earth mover distance: %f' % (total_emd / len(model_list))) writer.writerow([ total_time / len(model_list), total_cd / len(model_list), total_emd / len(model_list) ]) print('Chamfer distance per category') for synset_id in cd_per_cat.keys(): print(synset_id, '%f' % np.mean(cd_per_cat[synset_id])) writer.writerow([synset_id, np.mean(cd_per_cat[synset_id])]) print('Earth mover distance per category') for synset_id in emd_per_cat.keys(): print(synset_id, '%f' % np.mean(emd_per_cat[synset_id])) writer.writerow([synset_id, np.mean(emd_per_cat[synset_id])]) csv_file.close()
def train(args): is_training_pl = tf.placeholder(tf.bool, shape=(), name='is_training') global_step = tf.Variable(0, trainable=False, name='global_step') #Note that theta is a parameter used for progressive training theta = tf.train.piecewise_constant(global_step, [10000, 20000, 50000], [0.01, 0.1, 0.5, 1.0], 'theta_op') provider = TrainProvider(args, is_training_pl) ids, inputs, gt = provider.batch_data num_eval_steps = provider.num_test // args.batch_size print('provider.num_valid', provider.num_test) print('num_eval_steps', num_eval_steps) model_module = importlib.import_module('.%s' % args.model_type, 'models') model = model_module.Model(inputs, gt, theta, False) add_train_summary('alpha', theta) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True sess = tf.Session(config=config) sess.run(tf.global_variables_initializer()) saver = tf.train.Saver(max_to_keep=10) saver.restore(sess, args.model_path) coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess=sess, coord=coord) start_time = time.time() while not coord.should_stop(): print(colored('Testing...', 'grey', 'on_green')) total_time = 0 total_loss_fine = 0 cd_per_cat = {} sess.run(tf.local_variables_initializer()) for j in range(num_eval_steps): start = time.time() ids_eval, inputs_eval, gt_eval, loss_fine, fine = sess.run( [ids, inputs, gt, model.loss_fine, model.fine], feed_dict={is_training_pl: False}) synset_id = str(ids_eval[0]).split('_')[0].split('\'')[1] total_loss_fine += loss_fine total_time += time.time() - start if not cd_per_cat.get(synset_id): cd_per_cat[synset_id] = [] cd_per_cat[synset_id].append(loss_fine) if args.plot: for i in range(args.batch_size): model_id = str(ids_eval[i]).split('_')[1] os.makedirs(os.path.join(args.save_path, 'plots', synset_id), exist_ok=True) plot_path = os.path.join(args.save_path, 'plots', synset_id, '%s.png' % model_id) plot_pcd_three_views(plot_path, [inputs_eval[i], fine[i], gt_eval[i]], ['input', 'output', 'ground truth'], 'CD %.4f' % (loss_fine), [0.5, 0.5, 0.5]) print('Average Chamfer distance: %f' % (total_loss_fine / num_eval_steps)) print('Chamfer distance per category') dict_novel = { '02924116': 'Bus', '02818832': 'Bed', '02871439': 'bookshelf', '02828884': 'bench', '03467517': 'guitar', '03790512': 'motorbike', '04225987': 'skateboard', '03948459': 'pistol' } temp_loss = 0 for synset_id in dict_novel.keys(): temp_loss += np.mean(cd_per_cat[synset_id]) print(dict_novel[synset_id], ' %f' % np.mean(cd_per_cat[synset_id])) break print('Total time', datetime.timedelta(seconds=time.time() - start_time)) coord.request_stop() coord.join(threads) sess.close()