def train():
cfg = app_config
setup_environment(cfg)
o = np.ones(3)
z = np.zeros(3)
v = np.stack([o, z])
np.random.seed(0)
cam = torch.from_numpy(np.random.random((128, 4))).float()
myarr = torch.from_numpy(np.random.random((128,140,3))).float()
scaling_factor = torch.from_numpy(np.random.random((128,1))).float()
pc = myarr
global_step=0
pc.requires_grad=True
from util.point_cloud_to import smoothen_voxels3d, pointcloud_project_fast
model = model_pc.ModelPointCloud(cfg)
_sigma_rel, _gauss_sigma, _gauss_kernel = model.setup_sigma(None,global_step)
pc_out = pointcloud_project_fast(cfg,pc,cam,None,None, _gauss_kernel, scaling_factor=scaling_factor)
proj = pc_out['proj']
vx = pc_out['voxels']
tr_pc = pc_out['tr_pc']
drc_pr = pc_out['drc_probs']
dp = pc_out['proj_depth']
print('Proj ',proj.sum())
print('Voxels ',vx.sum())
print('Tr pc ',tr_pc.sum())
print('DRC ',drc_pr.sum())
print('Depth', dp.sum())
import pdb
pdb.set_trace()
def train():
cfg = app_config
setup_environment(cfg)
train_dir = cfg.checkpoint_dir
mkdir_if_missing(train_dir)
# tf.logging.set_verbosity(tf.logging.INFO)
split_name = "val" #train
dataset_file = os.path.join(cfg.inp_dir,
f"{cfg.synth_set}_{split_name}.pkl")
dataset = Chair_dataset(dataset_file, cfg)
if cfg.shuffle_dataset:
torch.manual_seed(7000)
print("*" * 30)
print('creating dataloader')
train_loader = DataLoader(dataset=dataset,
batch_size=cfg.batch_size,
num_workers=8,
shuffle=cfg.shuffle_dataset)
for epoch in tqdm(range(cfg.max_number_of_steps), desc='Epoch', ncols=100):
train_size = len(train_loader)
ts = time.time()
print_now = 0
for batch_idx, train_data in tqdm(enumerate(train_loader),
desc='Batch',
total=train_size,
ncols=100):
# global_step = tf.train.get_or_create_global_step()
# model = model_pc_pytorch.ModelPointCloud(cfg, global_step)
model = model_pc_pytorch.ModelPointCloud(cfg)
inputs = preprocess(cfg, train_data)
# print('inputs shape')
# for i in inputs:
# print(i, inputs[i].shape)
# Call Forward of model
# outputs = model(inputs)
# task_loss = model.get_loss(inputs, outputs)
# print(inputs.keys())
loss = model.optimize_parameters(inputs)
if print_now % 200 == 0:
print("Epoch: %d, Step: %d, Loss: %f" %
(epoch, print_now, loss.item()))
print_now += 1
#reg_loss = regularization_loss(train_scopes, cfg)
#loss = task_loss + reg_loss
# break
# break
print("Training Complete!")
'''
def eval():
cfg = app_config
setup_environment(cfg)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
train_dir = cfg.checkpoint_dir
split_name = "eval"
dataset_folder = cfg.inp_dir
dataset = ShapeRecords(dataset_folder, cfg, 'test')
dataset_loader = torch.utils.data.DataLoader(dataset,
batch_size=cfg.batch_size, shuffle=cfg.shuffle_dataset,
num_workers=4, drop_last=True)
run_eval()
def train():
cfg = app_config
setup_environment(cfg)
train_dir = cfg.checkpoint_dir
mkdir_if_missing(train_dir)
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.INFO)
# dataset = tf.data.Dataset.from_tensor_slices(data)
# dataset = dataset.batch(1)
# iterator = dataset.make_one_shot_iterator()
# train_data = iterator.get_next()
# print("SHape2", train_data.get_shape())
summary_writer = tfsum.create_file_writer(train_dir, flush_millis=10000)
global_step = tf.compat.v1.train.get_or_create_global_step()
model_fn = model_p.ModelPointCloud(cfg, global_step)
optimizer = tf.keras.optimizers.Adam(0.01)
train_loss = tf.keras.metrics.Mean(name='train_loss')
train_scopes = ['transform']
# var_list = get_trainable_variables(train_scopes)
def train_step(data, ref):
with tf.GradientTape() as tape:
outputs = model_fn(data, ref)
task_loss = model_fn.get_loss(outputs)
reg_loss = regularization_loss(train_scopes, cfg)
loss = task_loss + reg_loss
gradients = tape.gradient(loss, model_fn.trainable_variables)
optimizer.apply_gradients(zip(gradients, model_fn.trainable_variables))
train_loss(loss)
global_step_val = 0
while global_step_val < cfg.max_number_of_steps:
for images in train_ds:
train_step(images, base_pc)
print('step: {0}'
'loss = {loss_val:.4f}'.format(global_step_val,
loss_val=train_loss.result()))
def train():
cfg = app_config
setup_environment(cfg)
o = np.ones(3)
z = np.zeros(3)
v = np.stack([o, z])
np.random.seed(0)
myarr = torch.from_numpy(np.repeat(v, 8960, axis=0).reshape((128, 140, 3)))
myarr = torch.from_numpy(np.random.random((128, 140, 3)))
pc = myarr
pc.requires_grad = True
from util.point_cloud_to import pointcloud2voxels3d_fast
pc_out = pointcloud2voxels3d_fast(cfg, pc, None)
vx = pc_out[0]
loss = torch.sum(vx**2) / 2.
vx.register_hook(lambda x: print('output_grads_sum ', x.sum()))
pc.register_hook(lambda x: print('input_grads_sum ', x.sum()))
loss.backward()
print('loss ', loss)
import pdb
pdb.set_trace()
yaw, pitch, roll = ypr_from_campos(pos[0], pos[1], pos[2])
yaw = yaw + np.pi
# get everything from 0 to 2*pi
yaw = yaw%(2*np.pi)+0.00000001
pitch = pitch%(2*np.pi)+0.00000001
roll = roll%(2*np.pi) + 0.00000001
q1 = np.ceil(scale*yaw/np.pi)-1
q2 = np.ceil(scale*pitch/np.pi)-1
q3 = np.ceil(scale*roll/np.pi)-1
return q1*np.square(divs)+q2*divs+q3
cfg = app_config
setup_environment(cfg)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
train_dir = cfg.checkpoint_dir
split_name = "eval"
dataset_folder = cfg.inp_dir
dataset = ShapeRecords(dataset_folder, cfg, 'test')
dataset_loader = torch.utils.data.DataLoader(dataset,
batch_size=cfg.batch_size, shuffle=cfg.shuffle_dataset,
num_workers=4, drop_last=True)
divs = 4
cfg = app_config
exp_dir = cfg.checkpoint_dir
num_views = cfg.num_views
def compute_predictions():
cfg = app_config
setup_environment(cfg)
exp_dir = cfg.checkpoint_dir
cfg.batch_size = 1
cfg.step_size = 1
pc_num_points = cfg.pc_num_points
vox_size = cfg.vox_size
save_pred = cfg.save_predictions
save_voxels = cfg.save_voxels
fast_conversion = True
pose_student = cfg.pose_predictor_student and cfg.predict_pose
g = tf.Graph()
with g.as_default():
model = model_pc.ModelPointCloud(cfg)
out = build_model(model)
input_image = out["inputs"]
cam_matrix = out["camera_extr_src"]
cam_quaternion = out["cam_quaternion"]
point_cloud = out["points_1"]
rgb = out["rgb_1"] if cfg.pc_rgb else tf.no_op()
projs = out["projs"]
projs_rgb = out["projs_rgb"]
projs_depth = out["projs_depth"]
cam_transform = out["cam_transform"]
z_latent = out["z_latent"]
if pose_student:
proj_student, camera_pose_student = model_student(
input_image, model)
input_pc = tf.placeholder(tf.float32, [cfg.batch_size, None, 3])
if save_voxels:
if fast_conversion:
voxels, _ = pointcloud2voxels3d_fast(cfg, input_pc, None)
voxels = tf.expand_dims(voxels, axis=-1)
voxels = smoothen_voxels3d(cfg, voxels, model.gauss_kernel())
else:
voxels = pointcloud2voxels(cfg, input_pc, model.gauss_sigma())
q_inp = tf.placeholder(tf.float32, [1, 4])
q_matrix = as_rotation_matrix(q_inp)
input_pc, pred_quat, gt_quat, pc_unrot = model_unrotate_points(cfg)
pc_rot = quaternion_rotate(input_pc, pred_quat)
config = tf.ConfigProto(device_count={'GPU': 1})
config.gpu_options.per_process_gpu_memory_fraction = cfg.per_process_gpu_memory_fraction
sess = tf.Session(config=config)
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
variables_to_restore = slim.get_variables_to_restore(exclude=["meta"])
restorer = tf.train.Saver(variables_to_restore)
checkpoint_file = tf.train.latest_checkpoint(exp_dir)
print("restoring checkpoint", checkpoint_file)
restorer.restore(sess, checkpoint_file)
save_dir = os.path.join(exp_dir,
'{}_vis_proj'.format(cfg.save_predictions_dir))
mkdir_if_missing(save_dir)
save_pred_dir = os.path.join(exp_dir, cfg.save_predictions_dir)
mkdir_if_missing(save_pred_dir)
vis_size = cfg.vis_size
dataset = Dataset3D(cfg)
pose_num_candidates = cfg.pose_predict_num_candidates
num_views = cfg.num_views
plot_h = 4
plot_w = 6
num_views = int(min(num_views, plot_h * plot_w / 2))
if cfg.models_list:
model_names = parse_lines(cfg.models_list)
else:
model_names = [sample.name for sample in dataset.data]
num_models = len(model_names)
for k in range(num_models):
model_name = model_names[k]
sample = dataset.sample_by_name(model_name)
images = sample.image
masks = sample.mask
if cfg.saved_camera:
cameras = sample.camera
cam_pos = sample.cam_pos
if cfg.vis_depth_projs:
depths = sample.depth
if cfg.variable_num_views:
num_views = sample.num_views
print("{}/{} {}".format(k, num_models, model_name))
if pose_num_candidates == 1:
grid = np.empty((plot_h, plot_w), dtype=object)
else:
plot_w = pose_num_candidates + 1
if pose_student:
plot_w += 1
grid = np.empty((num_views, plot_w), dtype=object)
if save_pred:
all_pcs = np.zeros((num_views, pc_num_points, 3))
all_cameras = np.zeros((num_views, 4))
all_voxels = np.zeros((num_views, vox_size, vox_size, vox_size))
all_z_latent = np.zeros((num_views, cfg.fc_dim))
for view_idx in range(num_views):
input_image_np = images[[view_idx], :, :, :]
gt_mask_np = masks[[view_idx], :, :, :]
if cfg.saved_camera:
extr_mtr = cameras[view_idx, :, :]
cam_quaternion_np = quaternion_from_campos(
cam_pos[view_idx, :])
cam_quaternion_np = np.expand_dims(cam_quaternion_np, axis=0)
else:
extr_mtr = np.zeros((4, 4))
if cfg.pc_rgb:
proj_tensor = projs_rgb
elif cfg.vis_depth_projs:
proj_tensor = projs_depth
else:
proj_tensor = projs
(pc_np, rgb_np, proj_np, cam_transf_np, z_latent_np) = sess.run(
[point_cloud, rgb, proj_tensor, cam_transform, z_latent],
feed_dict={
input_image: input_image_np,
cam_matrix: extr_mtr,
cam_quaternion: cam_quaternion_np
})
if pose_student:
(proj_student_np, camera_student_np) = sess.run(
[proj_student, camera_pose_student],
feed_dict={input_image: input_image_np})
predicted_camera = camera_student_np
else:
predicted_camera = cam_transf_np
if cfg.vis_depth_projs:
proj_np = normalise_depthmap(proj_np)
if depths is not None:
depth_np = depths[view_idx, :, :, :]
depth_np = normalise_depthmap(depth_np)
else:
depth_np = 1.0 - np.squeeze(gt_mask_np)
if pose_student:
proj_student_np = normalise_depthmap(proj_student_np)
if cfg.predict_pose:
if cfg.save_rotated_points:
ref_rot = scipy.io.loadmat(
"{}/final_reference_rotation.mat".format(exp_dir))
ref_rot = ref_rot["rotation"]
pc_np_unrot = sess.run(pc_rot,
feed_dict={
input_pc: pc_np,
pred_quat: ref_rot
})
pc_np = pc_np_unrot
if cfg.pc_rgb:
gt_image = input_image_np
elif cfg.vis_depth_projs:
gt_image = depth_np
else:
gt_image = gt_mask_np
if pose_num_candidates == 1:
view_j = view_idx * 2 // plot_w
view_i = view_idx * 2 % plot_w
gt_image = np.squeeze(gt_image)
grid[view_j, view_i] = mask4vis(cfg, gt_image, vis_size)
curr_img = np.squeeze(proj_np)
grid[view_j, view_i + 1] = mask4vis(cfg, curr_img, vis_size)
if cfg.save_individual_images:
curr_dir = os.path.join(save_dir, sample.name)
if not os.path.exists(curr_dir):
os.makedirs(curr_dir)
imageio.imwrite(
os.path.join(curr_dir,
'{}_{}.png'.format(view_idx, 'rgb_gt')),
mask4vis(cfg, np.squeeze(input_image_np), vis_size))
imageio.imwrite(
os.path.join(curr_dir,
'{}_{}.png'.format(view_idx,
'mask_pred')),
mask4vis(cfg, np.squeeze(proj_np), vis_size))
else:
view_j = view_idx
gt_image = np.squeeze(gt_image)
grid[view_j, 0] = mask4vis(cfg, gt_image, vis_size)
for kk in range(pose_num_candidates):
curr_img = np.squeeze(proj_np[kk, :, :, :])
grid[view_j, kk + 1] = mask4vis(cfg, curr_img, vis_size)
if cfg.save_individual_images:
curr_dir = os.path.join(save_dir, sample.name)
if not os.path.exists(curr_dir):
os.makedirs(curr_dir)
imageio.imwrite(
os.path.join(
curr_dir,
'{}_{}_{}.png'.format(view_idx, kk,
'mask_pred')),
mask4vis(cfg, np.squeeze(curr_img), vis_size))
if cfg.save_individual_images:
imageio.imwrite(
os.path.join(curr_dir,
'{}_{}.png'.format(view_idx, 'mask_gt')),
mask4vis(cfg, np.squeeze(gt_mask_np), vis_size))
if pose_student:
grid[view_j,
-1] = mask4vis(cfg, np.squeeze(proj_student_np),
vis_size)
if save_pred:
all_pcs[view_idx, :, :] = np.squeeze(pc_np)
all_z_latent[view_idx] = z_latent_np
if cfg.predict_pose:
all_cameras[view_idx, :] = predicted_camera
if save_voxels:
# multiplying by two is necessary because
# pc->voxel conversion expects points in [-1, 1] range
pc_np_range = pc_np
if not fast_conversion:
pc_np_range *= 2.0
voxels_np = sess.run(voxels,
feed_dict={input_pc: pc_np_range})
all_voxels[view_idx, :, :, :] = np.squeeze(voxels_np)
vis_view = view_idx == 0 or cfg.vis_all_views
if cfg.vis_voxels and vis_view:
rgb_np = np.squeeze(rgb_np) if cfg.pc_rgb else None
vis_pc(np.squeeze(pc_np), rgb=rgb_np)
grid_merged = merge_grid(cfg, grid)
imageio.imwrite("{}/{}_proj.png".format(save_dir, sample.name),
grid_merged)
if save_pred:
if cfg.save_as_mat:
save_dict = {"points": all_pcs, "z_latent": all_z_latent}
if cfg.predict_pose:
save_dict["camera_pose"] = all_cameras
scipy.io.savemat("{}/{}_pc".format(save_pred_dir, sample.name),
mdict=save_dict)
else:
np.savez("{}/{}_pc".format(save_pred_dir, sample.name),
all_pcs)
if save_voxels:
np.savez("{}/{}_vox".format(save_pred_dir, sample.name),
all_voxels)
sess.close()
def train():
cfg = app_config
setup_environment(cfg)
train_dir = cfg.checkpoint_dir
mkdir_if_missing(train_dir)
tf.logging.set_verbosity(tf.logging.INFO)
split_name = "train"
dataset_file = os.path.join(cfg.inp_dir, f"{cfg.synth_set}_{split_name}.tfrecords")
dataset = tf.data.TFRecordDataset(dataset_file, compression_type=tf_record_compression(cfg))
if cfg.shuffle_dataset:
dataset = dataset.shuffle(7000)
dataset = dataset.map(lambda rec: parse_tf_records(cfg, rec), num_parallel_calls=3) \
.batch(cfg.batch_size) \
.prefetch(buffer_size=100) \
.repeat()
iterator = dataset.make_one_shot_iterator()
train_data = iterator.get_next()
summary_writer = tfsum.create_file_writer(train_dir, flush_millis=10000)
with summary_writer.as_default(), tfsum.record_summaries_every_n_global_steps(10):
global_step = tf.train.get_or_create_global_step()
print("global step: ", global_step)
model = model_pc.ModelPointCloud(cfg, global_step)
inputs = model.preprocess(train_data, cfg.step_size)
model_fn = model.get_model_fn(
is_training=True, reuse=False, run_projection=True)
outputs = model_fn(inputs)
# train_scopes
train_scopes = ["decoder/point_cloud"]
# train_scopes = ["decoder"]
# loss
task_loss = model.get_loss(inputs, outputs)
reg_loss = regularization_loss(train_scopes, cfg)
loss = task_loss + reg_loss
# summary op
summary_op = tfsum.all_summary_ops()
# optimizer
var_list = get_trainable_variables(train_scopes)
optimizer = tf.train.AdamOptimizer(get_learning_rate(cfg, global_step))
train_op = optimizer.minimize(loss, global_step, var_list)
# saver
max_to_keep = 2
saver = tf.train.Saver(max_to_keep=max_to_keep)
session_config = tf.ConfigProto(
log_device_placement=False)
session_config.gpu_options.allow_growth = cfg.gpu_allow_growth
session_config.gpu_options.per_process_gpu_memory_fraction = cfg.per_process_gpu_memory_fraction
sess = tf.Session(config=session_config)
with sess, summary_writer.as_default():
tf.global_variables_initializer().run()
tf.local_variables_initializer().run()
tfsum.initialize(graph=tf.get_default_graph())
## TODO load pretrain
variables_to_restore = slim.get_variables_to_restore(exclude=["meta"])
#, "decoder/point_cloud"
print(variables_to_restore)
restorer = tf.train.Saver(variables_to_restore)
checkpoint_file = "model-800000"
# checkpoint_file = "model-665000"
# checkpoint_file = "model-600000"
#tf.train.latest_checkpoint(cfg.checkpoint_dir)
print("restoring checkpoint", checkpoint_file)
restorer.restore(sess, checkpoint_file)
global_step_val = 0
while global_step_val < cfg.max_number_of_steps:
t0 = time.perf_counter()
_, loss_val, global_step_val, summary = sess.run([train_op, loss, global_step, summary_op])
t1 = time.perf_counter()
dt = t1 - t0
print(f"step: {global_step_val}, loss = {loss_val:.4f} ({dt:.3f} sec/step), lr = {sess.run(optimizer._lr)}")
if global_step_val % 5000 == 0:
saver.save(sess, f"{train_dir}/model", global_step=global_step_val)
def run_eval():
cfg = app_config
setup_environment(cfg)
exp_dir = cfg.checkpoint_dir
num_views = cfg.num_views
dataset_folder = cfg.inp_dir
save_pred_name = "{}_{}".format(cfg.save_predictions_dir, cfg.eval_split)
save_dir = os.path.join(exp_dir, cfg.save_predictions_dir)
reference_rotation = scipy.io.loadmat(
"{}/final_reference_rotation.mat".format(exp_dir))["rotation"]
ref_conj_np = quaternion_conjugate_np(reference_rotation)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
dataset = ShapeRecords(dataset_folder, cfg, 'test')
if cfg.models_list:
model_names = parse_lines(cfg.models_list)
else:
model_names = dataset.file_names
num_models = len(model_names)
angle_error = np.zeros((num_models, num_views), dtype=np.float64)
for k in range(num_models):
sample = dataset.__getitem__(k)
print("{}/{}".format(k, num_models))
print(model_names[k])
mat_filename = "{}/{}_pc.pkl".format(save_dir, model_names[k])
if not os.path.isfile(mat_filename):
continue
with open(mat_filename, 'rb') as handle:
data = pickle.load(handle)
all_cameras = data["camera_pose"]
for view_idx in range(num_views):
cam_pos = sample["cam_pos"][view_idx, :]
gt_quat_np = quaternion_from_campos(cam_pos)
gt_quat_np = np.expand_dims(gt_quat_np, 0)
pred_quat_np = all_cameras[view_idx, :]
pred_quat_np /= np.linalg.norm(pred_quat_np)
pred_quat_np = np.expand_dims(pred_quat_np, 0)
pred_quat_aligned_np = qmul(pred_quat_np, ref_conj_np)
q1 = gt_quat_np
q2 = pred_quat_aligned_np
q1_conj = quaternion_conjugate_np(q1)
q_diff = qmul(q1_conj, q2)
ang_diff = 2 * np.arccos(q_diff[0, 0])
if ang_diff > np.pi:
ang_diff -= 2 * np.pi
angle_error[k, view_idx] = np.fabs(ang_diff)
all_errors = np.reshape(angle_error, (-1))
angle_thresh_rad = cfg.pose_accuracy_threshold / 180.0 * np.pi
correct = all_errors < angle_thresh_rad
num_predictions = correct.shape[0]
accuracy = np.count_nonzero(correct) / num_predictions
median_error = np.sort(all_errors)[num_predictions // 2]
median_error = median_error / np.pi * 180
print("accuracy:", accuracy, "median angular error:", median_error)
scipy.io.savemat(
os.path.join(exp_dir, "pose_error_{}.mat".format(save_pred_name)), {
"angle_error": angle_error,
"accuracy": accuracy,
"median_error": median_error
})
f = open(os.path.join(exp_dir, "pose_error_{}.txt".format(save_pred_name)),
"w")
f.write("{} {}\n".format(accuracy, median_error))
f.close()
def compute_predictions():
cfg = app_config
setup_environment(cfg)
exp_dir = cfg.checkpoint_dir
cfg.batch_size = 1
cfg.step_size = 1
pc_num_points = cfg.pc_num_points
vox_size = cfg.vox_size
save_pred = cfg.save_predictions
save_voxels = cfg.save_voxels
fast_conversion = True
pose_student = cfg.pose_predictor_student and cfg.predict_pose
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = model_pc.ModelPointCloud(cfg)
model = model.to(device)
log_dir = '../../dpc/run/model_run_data/'
learning_rate = 1e-4
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay = cfg.weight_decay)
global_step = 100000
if global_step>0:
checkpoint_path = os.path.join(log_dir,'model.ckpt_{}.pth'.format(global_step))
print("Loading from path:",checkpoint_path)
checkpoint = torch.load(checkpoint_path)
global_step_val = checkpoint['global_step']
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
loss = checkpoint['loss']
else:
global_step_val = global_step
print('Restored checkpoint at {} with loss {}'.format(global_step, loss))
save_dir = os.path.join(exp_dir, '{}_vis_proj'.format(cfg.save_predictions_dir))
mkdir_if_missing(save_dir)
save_pred_dir = os.path.join(exp_dir, cfg.save_predictions_dir)
mkdir_if_missing(save_pred_dir)
vis_size = cfg.vis_size
split_name = "val"
dataset_folder = cfg.inp_dir
dataset = ShapeRecords(dataset_folder, cfg, split_name)
dataset_loader = torch.utils.data.DataLoader(dataset,
batch_size=cfg.batch_size, shuffle=cfg.shuffle_dataset,
num_workers=4,drop_last=True)
pose_num_candidates = cfg.pose_predict_num_candidates
num_views = cfg.num_views
plot_h = 4
plot_w = 6
num_views = int(min(num_views, plot_h * plot_w / 2))
if cfg.models_list:
model_names = parse_lines(cfg.models_list)
else:
model_names = dataset.file_names
num_models = len(model_names)
for k in range(num_models):
model_name = model_names[k]
sample = dataset.__getitem__(k)
images = sample['image']
masks = sample['mask']
if cfg.saved_camera:
cameras = sample['extrinsic']
cam_pos = sample['cam_pos']
if cfg.vis_depth_projs:
depths = sample['depth']
if cfg.variable_num_views:
num_views = sample['num_views']
print("{}/{} {}".format(k, num_models, model_name))
if pose_num_candidates == 1:
grid = np.empty((plot_h, plot_w), dtype=object)
else:
plot_w = pose_num_candidates + 1
if pose_student:
plot_w += 1
grid = np.empty((num_views, plot_w), dtype=object)
if save_pred:
all_pcs = np.zeros((num_views, pc_num_points, 3))
all_cameras = np.zeros((num_views, 4))
#all_voxels = np.zeros((num_views, vox_size, vox_size, vox_size))
#all_z_latent = np.zeros((num_views, cfg.fc_dim))
for view_idx in range(num_views):
input_image_np = images[[view_idx], :, :, :]
gt_mask_np = masks[[view_idx], :, :, :]
if cfg.saved_camera:
extr_mtr = cameras[view_idx, :, :]
cam_quaternion_np = quaternion_from_campos(cam_pos[view_idx, :])
cam_quaternion_np = np.expand_dims(cam_quaternion_np, axis=0)
else:
extr_mtr = np.zeros((4, 4))
code = 'images' if cfg.predict_pose else 'images_1'
input = {code: input_image_np,
'matrices': extr_mtr,
'camera_quaternion': cam_quaternion_np}
out = build_model(model, input, global_step)
input_image = out["inputs"]
cam_matrix = out["camera_extr_src"]
cam_quaternion = out["cam_quaternion"]
point_cloud = out["points_1"]
#gb = out["rgb_1"] if cfg.pc_rgb else None
#rojs = out["projs"]
#rojs_rgb = out["projs_rgb"]
#rojs_depth = out["projs_depth"]
cam_transform = out["cam_transform"]
#_latent = out["z_latent"]
#if cfg.pc_rgb:
# proj_tensor = projs_rgb
#elif cfg.vis_depth_projs:
# proj_tensor = projs_depth
#else:
# proj_tensor = projs
if pose_student:
camera_student_np = out["pose_student"]
predicted_camera = camera_student_np
else:
predicted_camera = cam_transf_np
#if cfg.vis_depth_projs:
# proj_np = normalise_depthmap(out["projs"])
# if depths is not None:
# depth_np = depths[view_idx, :, :, :]
# depth_np = normalise_depthmap(depth_np)
# else:
# depth_np = 1.0 - np.squeeze(gt_mask_np)
# if pose_student:
# proj_student_np = normalise_depthmap(proj_student_np)
#if save_voxels:
# if fast_conversion:
# voxels, _ = pointcloud2voxels3d_fast(cfg, input_pc, None)
# voxels = tf.expand_dims(voxels, axis=-1)
# voxels = smoothen_voxels3d(cfg, voxels, model.gauss_kernel())
# else:
# voxels = pointcloud2voxels(cfg, input_pc, model.gauss_sigma())
if cfg.predict_pose:
if cfg.save_rotated_points:
ref_rot = scipy.io.loadmat("{}/final_reference_rotation.mat".format(exp_dir))
ref_rot = ref_rot["rotation"]
pc_unrot = quaternion_rotate(input_pc, ref_quat)
point_cloud = pc_np_unrot
if cfg.pc_rgb:
gt_image = input_image_np
elif cfg.vis_depth_projs:
gt_image = depth_np
else:
gt_image = gt_mask_np
# if pose_num_candidates == 1:
# view_j = view_idx * 2 // plot_w
# view_i = view_idx * 2 % plot_w
# gt_image = np.squeeze(gt_image)
# grid[view_j, view_i] = mask4vis(cfg, gt_image, vis_size)
# curr_img = np.squeeze(out[projs])
# grid[view_j, view_i + 1] = mask4vis(cfg, curr_img, vis_size)
# if cfg.save_individual_images:
# curr_dir = os.path.join(save_dir, model_names[k])
# if not os.path.exists(curr_dir):
# os.makedirs(curr_dir)
# imageio.imwrite(os.path.join(curr_dir, '{}_{}.png'.format(view_idx, 'rgb_gt')),
# mask4vis(cfg, np.squeeze(input_image_np), vis_size))
# imageio.imwrite(os.path.join(curr_dir, '{}_{}.png'.format(view_idx, 'mask_pred')),
# mask4vis(cfg, np.squeeze(proj_np), vis_size))
# else:
# view_j = view_idx
# gt_image = np.squeeze(gt_image)
# grid[view_j, 0] = mask4vis(cfg, gt_image, vis_size)
# for kk in range(pose_num_candidates):
# curr_img = np.squeeze(out["projs"][kk, :, :, :].detach().cpu())
# grid[view_j, kk + 1] = mask4vis(cfg, curr_img, vis_size)
# if cfg.save_individual_images:
# curr_dir = os.path.join(save_dir, model_names[k])
# if not os.path.exists(curr_dir):
# os.makedirs(curr_dir)
# imageio.imwrite(os.path.join(curr_dir, '{}_{}_{}.png'.format(view_idx, kk, 'mask_pred')),
# mask4vis(cfg, np.squeeze(curr_img), vis_size))
# if cfg.save_individual_images:
# imageio.imwrite(os.path.join(curr_dir, '{}_{}.png'.format(view_idx, 'mask_gt')),
# mask4vis(cfg, np.squeeze(gt_mask_np), vis_size))
# if pose_student:
# grid[view_j, -1] = mask4vis(cfg, np.squeeze(proj_student_np.detach().cpu()), vis_size)
if save_pred:
#pc_np = pc_np.detach().cpu().numpy()
all_pcs[view_idx, :, :] = np.squeeze(point_cloud.detach().cpu())
#all_z_latent[view_idx] = z_latent.detach().cpu()
if cfg.predict_pose:
all_cameras[view_idx, :] = predicted_camera.detach().cpu()
# if save_voxels:
# # multiplying by two is necessary because
# # pc->voxel conversion expects points in [-1, 1] range
# pc_np_range = pc_np
# if not fast_conversion:
# pc_np_range *= 2.0
# voxels_np = sess.run(voxels, feed_dict={input_pc: pc_np_range})
# all_voxels[view_idx, :, :, :] = np.squeeze(voxels_np)
# vis_view = view_idx == 0 or cfg.vis_all_views
# if cfg.vis_voxels and vis_view:
# rgb_np = np.squeeze(rgb_np) if cfg.pc_rgb else None
# vis_pc(np.squeeze(pc_np), rgb=rgb_np)
#grid_merged = merge_grid(cfg, grid)
#imageio.imwrite("{}/{}_proj.png".format(save_dir, sample.file_names), grid_merged)
if save_pred:
if 0:
save_dict = {"points": all_pcs}
if cfg.predict_pose:
save_dict["camera_pose"] = all_cameras
scipy.io.savemat("{}/{}_pc.mat".format(save_pred_dir, model_names[k]),
mdict=save_dict)
else:
save_dict = {"points": all_pcs}
if cfg.predict_pose:
save_dict["camera_pose"] = all_cameras
with open("{}/{}_pc.pkl".format(save_pred_dir, model_names[k]), 'wb') as handle:
pickle.dump(save_dict, handle, protocol=pickle.HIGHEST_PROTOCOL)
def train():
cfg = app_config
setup_environment(cfg)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
train_dir = cfg.checkpoint_dir
mkdir_if_missing(train_dir)
split_name = "train"
dataset_folder = cfg.inp_dir
dataset = ShapeRecords(dataset_folder, cfg, split_name)
dataset_loader = torch.utils.data.DataLoader(dataset,
batch_size=cfg.batch_size,
shuffle=cfg.shuffle_dataset,
num_workers=4,
drop_last=True)
summary_writer = SummaryWriter(log_dir=train_dir, flush_secs=10)
ckpt_count = 1000
summary_count = 100
# loading pre existing model
# creating a new model
model = model_pc.ModelPointCloud(cfg)
model = model.to(device)
print(model.parameters)
log_dir = '../../dpc/run/model_run_data_lamp/'
mkdir_if_missing(log_dir)
learning_rate = 1e-4
optimizer = torch.optim.Adam(model.parameters(),
lr=learning_rate,
weight_decay=cfg.weight_decay)
# training steps
global_step = 38000
if global_step > 0:
checkpoint_path = os.path.join(log_dir,
'model.ckpt_{}.pth'.format(global_step))
print("Loading from path:", checkpoint_path)
checkpoint = torch.load(checkpoint_path)
global_step_val = checkpoint['global_step']
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
else:
global_step_val = global_step
model.train()
while global_step_val < cfg.max_number_of_steps:
step_loss = 0.0
for i, train_data in enumerate(dataset_loader, 0):
t9 = time.perf_counter()
for k in train_data.keys():
try:
train_data[k] = train_data[k].to(device)
except AttributeError:
pass
# get inputs by data processing
t0 = time.perf_counter()
inputs = model.preprocess(train_data, cfg.step_size)
t1 = time.perf_counter()
# zero the parameter gradients
optimizer.zero_grad()
t2 = time.perf_counter()
# dummy loss function
if global_step_val % summary_count == 0:
outputs = model(inputs,
global_step_val,
is_training=True,
run_projection=True,
summary_writer=summary_writer)
loss, min_loss = model.get_loss(inputs,
outputs,
summary_writer,
add_summary=True,
global_step=global_step_val)
summary_writer.add_image(
'prediction', outputs['projs'].detach().cpu().numpy()[
min_loss[0]].transpose(2, 0, 1), global_step_val)
summary_writer.add_image(
'actual',
inputs['masks'].detach().cpu().numpy()[0].transpose(
2, 0, 1), global_step_val)
#print(chamfer_distance( outputs['projs'].detach().cpu().numpy()[min_loss[0]].transpose(2, 0, 1), inputs['masks'].detach().cpu().numpy()[0].transpose(2, 0, 1))
else:
outputs = model(inputs,
global_step_val,
is_training=True,
run_projection=True)
loss, _ = model.get_loss(inputs, outputs, add_summary=False)
loss.backward()
optimizer.step()
del inputs
del outputs
t3 = time.perf_counter()
dt = t3 - t9
#print('Cuda {}'.format(t0-t9))
#print('Preprocess {}'.format(t1-t0))
#print('Forward {}'.format(t2-t1))
#print('Backward {}'.format(t3-t2))
step_loss += loss.item()
loss_avg = step_loss / (i + 1)
print(
f"step: {global_step_val}, loss= {loss.item():.5f}, loss_average = {loss_avg:.4f} ({dt:.3f} sec/step)"
)
if global_step_val % ckpt_count == 0: # save configuration
checkpoint_path = os.path.join(
log_dir, 'model.ckpt_{}.pth'.format(global_step_val))
print("PATH:", checkpoint_path)
torch.save(
{
'global_step': global_step_val,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss_avg
}, checkpoint_path)
global_step_val += 1
