def test(args):
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
loader = data_loader(data_path,
is_transform=True,
split="test",
img_size=(240, 320))
testloader = data.DataLoader(loader)
# Setup Model
model = Model_2b_depgd_GAP_MS(ResidualBlock, UpProj_Block, 1, num_class)
from collections import OrderedDict
state_dict = torch.load(args.model_path,
map_location=lambda storage, location: storage)
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
model.eval()
if torch.cuda.is_available():
model.cuda()
Aiou, Am_acc, acc_lbl, class_Aiou = accuracy(model, testloader)
print("PixAcc, mAcc, and mIoU are: %f, %f, %f" %
(acc_lbl, Am_acc, np.sum(class_Aiou[1:]) / float(num_class - 1)))
print("class Aiou:", class_Aiou[1:])
f_loss.write('Augmentation type: flip, centercrop\n\n')
f_loss.close()
f_iou = open(os.path.join(result_path, "log_acc.txt"), 'w')
f_iou.close()
f_iou = open(os.path.join(result_path, "log_val_acc.txt"), 'w')
f_iou.close()
# set up GPU
# we could do os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1, 2"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
# Data
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
tr_loader = data_loader(
data_path,
args.exp,
is_transform=True,
split='train',
img_size=(args.img_rows, args.img_cols),
augmentations=data_aug_tr,
)
te_loader = data_loader(data_path,
args.exp,
is_transform=True,
split='test',
img_size=(args.img_rows, args.img_cols),
augmentations=data_aug_te)
def train(args):
writer = SummaryWriter(comment=args.writer)
# data loader setting, train and evaluation
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
t_loader = data_loader(data_path, split='train', img_size=(args.img_rows, args.img_cols), img_norm=args.img_norm)
v_loader = data_loader(data_path, split='test', img_size=(args.img_rows, args.img_cols), img_norm=args.img_norm)
trainloader = data.DataLoader(t_loader, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers)
evalloader = data.DataLoader(v_loader, batch_size=args.batch_size, num_workers=args.num_workers)
print("Finish Loader Setup")
# Setup Model and load pretrained model
model_name = args.arch_RGB
# print(model_name)
model = get_model(model_name, True) # vgg_16
if args.pretrain: # True by default
if args.input == 'rgb': # only for rgb we have pretrain option
state = get_premodel(model, args.state_name)
model.load_state_dict(state)
model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count()))
elif args.input == 'd': # for d, load from result from...
print("Load training model: {}_{}_{}_{}_best.pkl".format(args.arch_RGB, args.dataset, args.loss, 1))
checkpoint = torch.load(pjoin(args.model_savepath_pretrain,
"{}_{}_{}_{}_best.pkl".format(args.arch_RGB, args.dataset, args.loss, 1)))
# model.load_state_dict(load_resume_state_dict(model, checkpoint['model_D_state']))
model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count()))
model.load_state_dict(checkpoint['model_D_state'])
else:
model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count()))
# model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count()))
# model_RGB = DataParallelWithCallback(model_RGB, device_ids=range(torch.cuda.device_count()))
model.cuda()
print("Finish model setup with model %s and state_dict %s" % (args.arch_RGB, args.state_name))
# optimizers and lr-decay setting
if args.pretrain: # True by default
optimizer_RGB = torch.optim.RMSprop(model.parameters(), lr=0.25 * args.l_rate)
scheduler_RGB = torch.optim.lr_scheduler.MultiStepLR(optimizer_RGB, milestones=[1, 2, 4, 8], gamma=0.5)
else:
optimizer_RGB = torch.optim.RMSprop(model.parameters(), lr=args.l_rate)
scheduler_RGB = torch.optim.lr_scheduler.MultiStepLR(optimizer_RGB, milestones=[1, 3, 5, 8, 11, 15], gamma=0.5)
# forward and backward
best_loss = 3
n_iter_t, n_iter_v = 0, 0
if args.dataset == 'matterport':
total_iter_t = 105432 / args.batch_size
elif args.dataset == 'scannet':
total_iter_t = 59743 / args.batch_size
else:
total_iter_t = 0
if not os.path.exists(args.model_savepath):
os.makedirs(args.model_savepath)
for epoch in range(args.n_epoch):
scheduler_RGB.step()
model.train()
for i, (images, labels, masks, valids, depthes, meshdepthes) in enumerate(trainloader):
n_iter_t += 1
images = Variable(images.contiguous().cuda())
labels = Variable(labels.contiguous().cuda())
masks = Variable(masks.contiguous().cuda())
optimizer_RGB.zero_grad()
if args.input == 'rgb':
outputs = model(images)
else:
depthes = Variable(depthes.contiguous().cuda())
if args.input == 'rgbd':
rgbd_input = torch.cat((images, depthes), dim=1)
outputs = model(rgbd_input)
elif args.input == 'd':
outputs = model(depthes)
loss, df = get_lossfun(args.loss, outputs, labels, masks)
if args.l1regular:
loss_rgl, df_rgl = get_lossfun('l1gra', outputs, labels, masks)
elif args.gradloss:
loss_grad, df_grad = get_lossfun('gradmap', outputs, labels, masks)
if args.l1regular:
outputs.backward(gradient=df, retain_graph=True)
outputs.backward(gradient=0.1 * df_rgl)
elif args.gradloss:
outputs.backward(gradient=df, retain_graph=True)
outputs.backward(gradient=0.5 * df_grad)
else:
outputs.backward(gradient=df)
optimizer_RGB.step()
if (i + 1) % 100 == 0:
if args.l1regular:
print("Epoch [%d/%d] Iter [%d/%d] Loss and RGL: %.4f, %.4f" % (
epoch + 1, args.n_epoch, i, total_iter_t, loss.data, loss_rgl.data))
elif args.gradloss:
print("Epoch [%d/%d] Iter [%d/%d] Loss and GradLoss: %.4f, %.4f" % (
epoch + 1, args.n_epoch, i, total_iter_t, loss.data, loss_grad.data))
else:
print("Epoch [%d/%d] Iter [%d/%d] Loss: %.4f" % (
epoch + 1, args.n_epoch, i, total_iter_t, loss.data))
if (i + 1) % 250 == 0:
writer.add_scalar('loss/trainloss', loss.data.item(), n_iter_t)
if args.l1regular:
writer.add_scalar('loss/trainloss_rgl', loss_rgl.data.item(), n_iter_t)
elif args.gradloss:
writer.add_scalar('loss/trainloss_grad', loss_grad.data.item(), n_iter_t)
writer.add_images('Image', images + 0.5, n_iter_t)
if args.input != 'rgb':
writer.add_images('Depth', np.repeat(
(depthes - torch.min(depthes)) / (torch.max(depthes) - torch.min(depthes)), 3, axis=1),
n_iter_t)
writer.add_images('Label', 0.5 * (labels.permute(0, 3, 1, 2) + 1), n_iter_t)
outputs_n = norm_tf(outputs)
writer.add_images('Output', outputs_n, n_iter_t)
model.eval()
mean_loss, sum_loss, sum_rgl, sum_grad = 0, 0, 0, 0
evalcount = 0
with torch.no_grad():
for i_val, (images_val, labels_val, masks_val, valids_val, depthes_val, meshdepthes_val) in tqdm(
enumerate(evalloader)):
n_iter_v += 1
images_val = Variable(images_val.contiguous().cuda())
labels_val = Variable(labels_val.contiguous().cuda())
masks_val = Variable(masks_val.contiguous().cuda())
if args.input == 'rgb':
outputs = model(images_val)
else:
depthes_val = Variable(depthes_val.contiguous().cuda())
if args.input == 'rgbd':
rgbd_input = torch.cat((images_val, depthes_val), dim=1)
outputs = model(rgbd_input)
elif args.input == 'd':
outputs = model(depthes_val)
loss, df = get_lossfun(args.loss, outputs, labels_val, masks_val, False) # valid_val not used infact
if args.l1regular:
loss_rgl, df_rgl = get_lossfun('l1gra', outputs, labels_val, masks_val, False)
elif args.gradloss:
loss_grad, df_grad = get_lossfun('gradmap', outputs, labels_val, masks_val, False)
if ((np.isnan(loss)) | (np.isinf(loss))):
sum_loss += 0
else:
sum_loss += loss
evalcount += 1
if args.l1regular:
sum_rgl += loss_rgl
elif args.gradloss:
sum_grad += loss_grad
if (i_val + 1) % 250 == 0:
# print("Epoch [%d/%d] Evaluation Loss: %.4f" % (epoch+1, args.n_epoch, loss))
writer.add_scalar('loss/evalloss', loss, n_iter_v)
writer.add_images('Eval Image', images_val + 0.5, n_iter_t)
if args.input != 'rgb':
writer.add_image('Depth', np.repeat(
(depthes_val - torch.min(depthes_val)) / (torch.max(depthes_val) - torch.min(depthes_val)),
3, axis=1), n_iter_t)
writer.add_images('Eval Label', 0.5 * (labels_val.permute(0, 3, 1, 2) + 1), n_iter_t)
outputs_n = norm_tf(outputs)
writer.add_images('Eval Output', outputs_n, n_iter_t)
mean_loss = sum_loss / evalcount
print("Epoch [%d/%d] Evaluation Mean Loss: %.4f" % (epoch + 1, args.n_epoch, mean_loss))
writer.add_scalar('loss/evalloss_mean', mean_loss, epoch)
writer.add_scalar('loss/evalloss_rgl_mean', sum_rgl / evalcount, epoch)
writer.add_scalar('loss/evalloss_grad_mean', sum_grad / evalcount, epoch)
if mean_loss < best_loss: # if (epoch+1)%20 == 0:
best_loss = mean_loss
state = {'epoch': epoch + 1,
'model_RGB_state': model.state_dict(),
'optimizer_RGB_state': optimizer_RGB.state_dict(), }
if args.pretrain:
if args.l1regular:
torch.save(state, pjoin(args.model_savepath,
"{}_{}_{}_{}_rgls_best.pkl".format(args.arch_RGB, args.dataset, args.loss,
args.model_num)))
elif args.gradloss:
torch.save(state, pjoin(args.model_savepath,
"{}_{}_{}_{}_grad_best.pkl".format(args.arch_RGB, args.dataset, args.loss,
args.model_num)))
else:
torch.save(state, pjoin(args.model_savepath,
"{}_{}_{}_{}_resume_RGB_best.pkl".format(args.arch_RGB, args.dataset,
args.loss, args.model_num)))
else:
torch.save(state, pjoin(args.model_savepath,
"{}_{}_{}_{}_resume_RGB_best.pkl".format(args.arch_RGB, args.dataset, args.loss,
args.model_num)))
print('Finish training for dataset %s trial %s' % (args.dataset, args.model_num))
# state = {'epoch': epoch+1,
# 'model_RGB_state': model_RGB.state_dict(),
# 'optimizer_RGB_state' : optimizer_RGB.state_dict(),}
# if args.pretrain:
# torch.save(state, pjoin(args.model_savepath, "{}_{}_{}_{}_RGB_final.pkl".format(args.arch_RGB, args.dataset, args.loss, args.model_num)))
# elif args.l1regular:
# torch.save(state, pjoin(args.model_savepath, "{}_{}_{}_{}_rgls_final.pkl".format(args.arch_RGB, args.dataset, args.loss, args.model_num)))
# else:
# torch.save(state, pjoin(args.model_savepath, "{}_{}_{}_{}_nopretrain_final.pkl".format(args.arch_RGB, args.dataset, args.loss, args.model_num)))
writer.export_scalars_to_json("./{}_{}_{}_{}.json".format(args.arch_RGB, args.dataset, args.loss, args.model_num))
writer.close()
def test(args):
# Setup Model
model_name_F = args.arch_F
model_F = get_model(model_name_F,True) # concat and output
model_F = torch.nn.DataParallel(model_F, device_ids=range(torch.cuda.device_count()))
if args.arch_map == 'map_conv' or args.arch_map == 'hybrid':
model_name_map = 'map_conv'
model_map = get_model(model_name_map,True) # concat and output
model_map = torch.nn.DataParallel(model_map, device_ids=range(torch.cuda.device_count()))
if args.model_full_name != '':
# Use the full name of model to load
print("Load training model: " + args.model_full_name)
checkpoint = torch.load(pjoin(args.model_savepath, args.model_full_name))
model_F.load_state_dict(checkpoint['model_F_state'])
model_map.load_state_dict(checkpoint["model_map_state"])
# Setup image
if args.imgset:
print("Test on dataset: {}".format(args.dataset))
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
v_loader = data_loader(data_path, split=args.test_split, img_size=(args.img_rows,args.img_cols), img_norm=args.img_norm,mode='seg')
evalloader = data.DataLoader(v_loader, batch_size=1)
print("Finish Loader Setup")
model_F.cuda()
model_F.eval()
if args.arch_map == 'map_conv' or args.arch_map == 'hybrid':
model_map.cuda()
model_map.eval()
sum_mean, sum_median, sum_small, sum_mid, sum_large, sum_num = [], [], [], [], [], []
sum_mean_b, sum_median_b, sum_small_b, sum_mid_b, sum_large_b, sum_num_b = [], [], [], [], [], []
sum_mean_s, sum_median_s, sum_small_s, sum_mid_s, sum_large_s, sum_num_s = [], [], [], [], [], []
sum_mean_c, sum_median_c, sum_small_c, sum_mid_c, sum_large_c, sum_num_c = [], [], [], [], [], []
evalcount = 0
with torch.no_grad():
for i_val, (images_val, labels_val, masks_val, valids_val, depthes_val, segment_val) in tqdm(enumerate(evalloader)):
# if i_val>10:
# break
images_val = Variable(images_val.contiguous().cuda())
labels_val = Variable(labels_val.contiguous().cuda())
masks_val = Variable(masks_val.contiguous().cuda())
valids_val = Variable(valids_val.contiguous().cuda())
depthes_val = Variable(depthes_val.contiguous().cuda())
segment_val = Variable(segment_val.contiguous().cuda())
# Bed:11 1191 494 786 1349
# Sofa:6 1313
# Chair:2 10 23 74 885 1184 1291 1338
segment_bed = torch.eq(segment_val,11)+torch.eq(segment_val,1191)+torch.eq(segment_val,494)+torch.eq(segment_val,786)+torch.eq(segment_val,1349)
segment_sofa = torch.eq(segment_val,6)+torch.eq(segment_val,1313)
segment_chair = torch.eq(segment_val,2)+torch.eq(segment_val,10)+torch.eq(segment_val,23)+torch.eq(segment_val,74)+torch.eq(segment_val,885)+torch.eq(segment_val,1184)+torch.eq(segment_val,1291)+torch.eq(segment_val,1338)
if segment_val.shape != masks_val.shape:
segment_bed = eval_mask_resize(segment_bed, args.img_rows, args.img_cols)
segment_sofa = eval_mask_resize(segment_sofa, args.img_rows, args.img_cols)
segment_chair = eval_mask_resize(segment_chair, args.img_rows, args.img_cols)
if args.arch_map == 'map_conv' or args.arch_map == 'hybrid':
outputs_valid = model_map(torch.cat((depthes_val, valids_val[:,np.newaxis,:,:]), dim=1))
outputs, outputs1, outputs2, outputs3,output_d = model_F(images_val, depthes_val, outputs_valid.squeeze(1))
else:
outputs, outputs1, outputs2, outputs3,output_d = model_F(images_val, depthes_val, valids_val)
outputs_n, pixelnum, mean_i, median_i, small_i, mid_i, large_i = eval_normal_pixel(outputs, labels_val, masks_val)
masks_bed_val = segment_bed.to(torch.float64)*masks_val
masks_sofa_val = segment_sofa.to(torch.float64)*masks_val
masks_chair_val = segment_chair.to(torch.float64)*masks_val
_, pixelnum_b, mean_i_b, median_i_b, small_i_b, mid_i_b, large_i_b = eval_normal_pixel(outputs, labels_val, masks_bed_val)
_, pixelnum_s, mean_i_s, median_i_s, small_i_s, mid_i_s, large_i_s = eval_normal_pixel(outputs, labels_val, masks_sofa_val)
_, pixelnum_c, mean_i_c, median_i_c, small_i_c, mid_i_c, large_i_c = eval_normal_pixel(outputs, labels_val, masks_chair_val)
# outputs_norm = np.squeeze(outputs_n.data.cpu().numpy(), axis=0)
# labels_val_norm = np.squeeze(labels_val.data.cpu().numpy(), axis=0)
# images_val = np.squeeze(images_val.data.cpu().numpy(), axis=0)
# images_val = images_val+0.5
# images_val = images_val.transpose(1, 2, 0)
# depthes_val = np.squeeze(depthes_val.data.cpu().numpy(), axis=0)
# depthes_val = np.transpose(depthes_val, [1,2,0])
# depthes_val = np.repeat(depthes_val, 3, axis = 2)
# masks_bed_val = np.squeeze(masks_bed_val.data.cpu().numpy(), axis=0)
# masks_sofa_val = np.squeeze(masks_sofa_val.data.cpu().numpy(), axis=0)
# masks_chair_val = np.squeeze(masks_chair_val.data.cpu().numpy(), axis=0)
# outputs_valid = np.squeeze(outputs_valid.data.cpu().numpy(), axis=0)
# outputs_valid = np.transpose(outputs_valid, [1,2,0])
# outputs_valid = np.repeat(outputs_valid, 3, axis = 2)
# if (i_val+1)%1 == 0:
# misc.imsave(pjoin(args.testset_out_path, "{}_fms_hybrid2.png".format(i_val+1)), outputs_norm)
# misc.imsave(pjoin(args.testset_out_path, "{}_fms1_l1.png".format(i_val+1)), outputs_norm1)
# misc.imsave(pjoin(args.testset_out_path, "{}_fms2_l1.png".format(i_val+1)), outputs_norm2)
# misc.imsave(pjoin(args.testset_out_path, "{}_fms3_l1.png".format(i_val+1)), outputs_norm3)
# misc.imsave(pjoin(args.testset_out_path, "{}_fmsd_l1.png".format(i_val+1)), outputs_normd)
# misc.imsave(pjoin(args.testset_out_path, "{}_d_l1_imgout.png".format(i_val+1)), outputs_norm)
# misc.imsave(pjoin(args.testset_out_path, "{}_gt.png".format(i_val+1)), labels_val_norm)
# misc.imsave(pjoin(args.testset_out_path, "{}_in.jpg".format(i_val+1)), images_val)
# # misc.imsave(pjoin(args.testset_out_path, "{}_depth.png".format(i_val+1)), depthes_val)
# misc.imsave(pjoin(args.testset_out_path, "{}_bed.png".format(i_val+1)), masks_bed_val)
# misc.imsave(pjoin(args.testset_out_path, "{}_sofa.png".format(i_val+1)), masks_sofa_val)
# misc.imsave(pjoin(args.testset_out_path, "{}_chair.png".format(i_val+1)), masks_chair_val)
# misc.imsave(pjoin(args.testset_out_path, "{}_ms_conf.png".format(i_val+1)), outputs_valid)
# if i_val == 0:
# outputs_mat = outputs_n.data.cpu().numpy()
# else:
# outputs_mat = np.concatenate((outputs_mat,outputs_n.data.cpu().numpy()), axis=0)
# accumulate the metrics in matrix
if ((np.isnan(mean_i))|(np.isinf(mean_i)) == False):
sum_mean.append(mean_i)
sum_median.append(median_i)
sum_small.append(small_i)
sum_mid.append(mid_i)
sum_large.append(large_i)
sum_num.append(pixelnum)
sum_mean_b.append(mean_i_b)
sum_median_b.append(median_i_b)
sum_small_b.append(small_i_b)
sum_mid_b.append(mid_i_b)
sum_large_b.append(large_i_b)
sum_num_b.append(pixelnum_b)
sum_mean_s.append(mean_i_s)
sum_median_s.append(median_i_s)
sum_small_s.append(small_i_s)
sum_mid_s.append(mid_i_s)
sum_large_s.append(large_i_s)
sum_num_s.append(pixelnum_s)
sum_mean_c.append(mean_i_c)
sum_median_c.append(median_i_c)
sum_small_c.append(small_i_c)
sum_mid_c.append(mid_i_c)
sum_large_c.append(large_i_c)
sum_num_c.append(pixelnum_c)
evalcount+=1
if (i_val+1) % 10 == 0:
print("Iteration %d Evaluation Loss: mean %.4f, median %.4f, 11.25 %.4f, 22.5 %.4f, 30 %.4f" % (i_val+1,
mean_i, median_i, small_i, mid_i, large_i))
# Summarize the result
eval_print(sum_mean, sum_median, sum_small, sum_mid, sum_large, sum_num, item='Pixel-Level')
eval_print(sum_mean_b, sum_median_b, sum_small_b, sum_mid_b, sum_large_b, sum_num_b, item='Bed-pixel')
eval_print(sum_mean_s, sum_median_s, sum_small_s, sum_mid_s, sum_large_s, sum_num_s, item='Sofa-pixel')
eval_print(sum_mean_c, sum_median_c, sum_small_c, sum_mid_c, sum_large_c, sum_num_c, item='Chair-pixel')
avg_mean = sum(sum_mean)/evalcount
sum_mean.append(avg_mean)
avg_median = sum(sum_median)/evalcount
sum_median.append(avg_median)
avg_small = sum(sum_small)/evalcount
sum_small.append(avg_small)
avg_mid = sum(sum_mid)/evalcount
sum_mid.append(avg_mid)
avg_large = sum(sum_large)/evalcount
sum_large.append(avg_large)
print("evalnum is %d, Evaluation Image-Level Mean Loss: mean %.4f, median %.4f, 11.25 %.4f, 22.5 %.4f, 30 %.4f" % (evalcount,
avg_mean, avg_median, avg_small, avg_mid, avg_large))
sum_matrix = np.transpose([sum_mean,sum_median,sum_small,sum_mid,sum_large])
if args.model_full_name != '':
sum_file = args.model_full_name[:-4] + '.csv'
np.savetxt(pjoin(args.model_savepath,sum_file), sum_matrix, fmt='%.6f', delimiter=',')
print("Saving to %s" % (sum_file))
def test(args):
# Setup Model
model_name = args.arch_RGB
model = get_model(model_name, True) # vgg_16
testset_out_default_path = "./result/"
testset_out_path = args.result_path
if args.imgset:
test_info = args.test_dataset + '_' + args.test_split
else:
if args.img_datalist == '':
# Single image
test_info = 'single'
else:
# Image list
test_info = 'batch'
if args.model_full_name != '':
# Use the full name of model to load
print("Load training model: " + args.model_full_name)
checkpoint = torch.load(
pjoin(args.model_savepath, args.model_full_name))
if testset_out_path == '':
testset_out_path = "{}{}_{}".format(testset_out_default_path,
args.model_full_name,
test_info)
model = torch.nn.DataParallel(model,
device_ids=range(
torch.cuda.device_count()))
model.load_state_dict(checkpoint['model_RGB_state'])
else:
# Pretrain model
print("Load pretrained model: {}".format(args.state_name))
state = get_premodel(model, args.state_name)
model = torch.nn.DataParallel(model,
device_ids=range(
torch.cuda.device_count()))
model.load_state_dict(state)
if testset_out_path == '':
testset_out_path = "{}pretrain_{}".format(testset_out_default_path,
test_info)
# Setup image
# Create output folder if needed
# if os.path.isdir(testset_out_path) == False:
# os.mkdir(testset_out_path)
if args.imgset:
print("Test on dataset: {}".format(args.test_dataset))
# Set up dataloader
data_loader = get_loader(args.test_dataset)
data_path = get_data_path(args.test_dataset)
v_loader = data_loader(data_path,
split=args.test_split,
img_size=(args.img_rows, args.img_cols),
img_norm=args.img_norm)
evalloader = data.DataLoader(v_loader, batch_size=1)
print("Finish Loader Setup")
model.cuda()
model.eval()
sum_mean, sum_median, sum_small, sum_mid, sum_large = 0, 0, 0, 0, 0
evalcount = 0
if args.numerical_result_path != '':
f = open(args.numerical_result_path, 'w')
with torch.no_grad():
# for i_val, (images_val, labels_val, masks_val, valids_val) in tqdm(enumerate(evalloader)):
for i_val, (images_val, labels_val, masks_val, valids_val,
depthes_val,
meshdepthes_val) in tqdm(enumerate(evalloader)):
images_val = Variable(images_val.contiguous().cuda())
labels_val = Variable(labels_val.contiguous().cuda())
masks_val = Variable(masks_val.contiguous().cuda())
if args.arch_RGB == 'ms':
outputs, outputs2, outputs3, outputs4, outputs5 = model(
images_val)
else:
outputs = model(images_val) # 1*ch*h*w
outputs_n, mean_i, median_i, small_i, mid_i, large_i = eval_normal(
outputs, labels_val, masks_val)
outputs_norm = np.squeeze(outputs_n.data.cpu().numpy(), axis=0)
labels_val_norm = np.squeeze(labels_val.data.cpu().numpy(),
axis=0)
images_val = np.squeeze(images_val.data.cpu().numpy(), axis=0)
images_val = images_val + 0.5
images_val = images_val.transpose(1, 2, 0)
outputs_norm = change_channel(outputs_norm)
# outputs_norm= temp_change_mlt_chanel(outputs_norm)
labels_val_norm = (labels_val_norm + 1) / 2 # scale to 0 1
# Change channel to have a better appearance for paper.
labels_val_norm = change_channel(labels_val_norm)
misc.imsave(
pjoin(testset_out_path, "{}_out.png".format(i_val + 1)),
outputs_norm)
misc.imsave(
pjoin(testset_out_path, "{}_gt.png".format(i_val + 1)),
labels_val_norm)
misc.imsave(
pjoin(testset_out_path, "{}_in.jpg".format(i_val + 1)),
images_val)
if ((np.isnan(mean_i)) | (np.isinf(mean_i))):
print('Error!')
sum_mean += 0
else:
sum_mean += mean_i
sum_median += median_i
sum_small += small_i
sum_mid += mid_i
sum_large += large_i
evalcount += 1
if (i_val + 1) % 1 == 0:
print(
"Iteration %d Evaluation Loss: mean %.4f, median %.4f, 11.25 %.4f, 22.5 %.4f, 30 %.4f"
% (i_val + 1, mean_i, median_i, small_i, mid_i,
large_i))
if (args.numerical_result_path != ''):
f.write(
"Iteration %d Evaluation Loss: mean %.4f, median %.4f, 11.25 %.4f, 22.5 %.4f, 30 %.4f\n"
% (i_val + 1, mean_i, median_i, small_i, mid_i,
large_i))
sum_mean = sum_mean / evalcount
sum_median = sum_median / evalcount
sum_small = sum_small / evalcount
sum_mid = sum_mid / evalcount
sum_large = sum_large / evalcount
if args.numerical_result_path != '':
f.write(
"evalnum is %d, Evaluation Mean Loss: mean %.4f, median %.4f, 11.25 %.4f, 22.5 %.4f, 30 %.4f"
% (evalcount, sum_mean, sum_median, sum_small, sum_mid,
sum_large))
f.close()
print(
"evalnum is %d, Evaluation Mean Loss: mean %.4f, median %.4f, 11.25 %.4f, 22.5 %.4f, 30 %.4f"
% (evalcount, sum_mean, sum_median, sum_small, sum_mid,
sum_large))
# end of dataset test
else:
if args.img_datalist == "":
# For single image, without GT
print("Read Input Image from : {}".format(args.img_path))
img = misc.imread(args.img_path)
if args.img_rot:
img = np.transpose(img, (1, 0, 2))
img = np.flipud(img)
orig_size = img.shape[:-1]
img = misc.imresize(
img,
(args.img_rows,
args.img_cols)) # Need resize the image to model inputsize
img = img.astype(np.float)
if args.img_norm:
img = (img - 128) / 255
# NHWC -> NCHW
img = img.transpose(2, 0, 1)
img = np.expand_dims(img, 0)
img = torch.from_numpy(img).float()
if torch.cuda.is_available():
model.cuda(0)
images = Variable(img.contiguous().cuda(0))
else:
images = Variable(img)
with torch.no_grad():
outputs = model(images)
outputs_norm = norm_imsave(outputs)
outputs_norm = np.squeeze(outputs_norm.data.cpu().numpy(), axis=0)
# Change channels
outputs_norm = change_channel(outputs_norm)
misc.imsave(args.out_path, outputs_norm)
print("Complete")
# end of test on single image
else:
# For image list without GT
data_list = get_dataList(args.img_datalist)
for img_path in data_list:
print("Read Input Image from : {}".format(img_path))
img = misc.imread(pjoin(args.img_dataroot, img_path))
height, width, channels = img.shape
output_filename = img_path.split('/')[-1]
if args.img_rot:
img = np.transpose(img, (1, 0, 2))
img = np.flipud(img)
orig_size = img.shape[:-1]
img = misc.imresize(img, (
args.img_rows,
args.img_cols)) # Need resize the image to model inputsize
img = img.astype(np.float)
if args.img_norm:
img = (img - 128) / 255
# NHWC -> NCHW
img = img.transpose(2, 0, 1)
img = np.expand_dims(img, 0)
img = torch.from_numpy(img).float()
if torch.cuda.is_available():
model.cuda(0)
images = Variable(img.contiguous().cuda(0))
else:
images = Variable(img)
with torch.no_grad():
outputs = model(images)
outputs_norm = norm_imsave(outputs)
outputs_norm = np.squeeze(outputs_norm.data.cpu().numpy(),
axis=0)
# Change channels
outputs_norm = change_channel(outputs_norm)
# Resize to the original size, if needed
# outputs_norm = misc.imresize(outputs_norm, (height, width))
# Save the result
misc.imsave(pjoin(args.out_path, output_filename),
outputs_norm)
# Save to mat file, if needed
# outputs_mat = outputs.tolist()
# mat_filename = output_filename.replace(output_filename.split('.')[-1], 'mat')
# sio.savemat(pjoin(testset_out_path, mat_filename), {'normal': outputs_mat});
print("Complete")
pass
def test(args):
# Setup Model
# Setup the fusion model (RGB+Depth)
model_name_F = args.arch_F
model_F = get_model(model_name_F, True) # concat and output
model_F = torch.nn.DataParallel(model_F,
device_ids=range(
torch.cuda.device_count()))
# Setup the map model
if args.arch_map == 'map_conv':
model_name_map = args.arch_map
model_map = get_model(model_name_map, True) # concat and output
model_map = torch.nn.DataParallel(model_map,
device_ids=range(
torch.cuda.device_count()))
if args.model_full_name != '':
# Use the full name of model to load
print("Load training model: " + args.model_full_name)
checkpoint = torch.load(
pjoin(args.model_savepath, args.model_full_name))
model_F.load_state_dict(checkpoint['model_F_state'])
model_map.load_state_dict(checkpoint["model_map_state"])
# Setup image
if args.imgset:
print("Test on dataset: {}".format(args.dataset))
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
v_loader = data_loader(data_path,
split=args.test_split,
img_size=(args.img_rows, args.img_cols),
img_norm=args.img_norm)
evalloader = data.DataLoader(v_loader, batch_size=1)
print("Finish Loader Setup")
model_F.cuda()
model_F.eval()
if args.arch_map == 'map_conv':
model_map.cuda()
model_map.eval()
sum_mean, sum_median, sum_small, sum_mid, sum_large, sum_num = [], [], [], [], [], []
evalcount = 0
with torch.no_grad():
for i_val, (images_val, labels_val, masks_val, valids_val,
depthes_val,
meshdepthes_val) in tqdm(enumerate(evalloader)):
images_val = Variable(images_val.contiguous().cuda())
labels_val = Variable(labels_val.contiguous().cuda())
masks_val = Variable(masks_val.contiguous().cuda())
valids_val = Variable(valids_val.contiguous().cuda())
depthes_val = Variable(depthes_val.contiguous().cuda())
if args.arch_map == 'map_conv':
outputs_valid = model_map(
torch.cat(
(depthes_val, valids_val[:, np.newaxis, :, :]),
dim=1))
outputs, outputs1, outputs2, outputs3, output_d = model_F(
images_val, depthes_val, outputs_valid.squeeze(1))
else:
outputs, outputs1, outputs2, outputs3, output_d = model_F(
images_val, depthes_val, valids_val)
outputs_n, pixelnum, mean_i, median_i, small_i, mid_i, large_i = eval_normal_pixel(
outputs, labels_val, masks_val)
outputs_norm = np.squeeze(outputs_n.data.cpu().numpy(), axis=0)
labels_val_norm = np.squeeze(labels_val.data.cpu().numpy(),
axis=0)
images_val = np.squeeze(images_val.data.cpu().numpy(), axis=0)
images_val = images_val + 0.5
images_val = images_val.transpose(1, 2, 0)
depthes_val = np.squeeze(depthes_val.data.cpu().numpy(),
axis=0)
depthes_val = np.transpose(depthes_val, [1, 2, 0])
depthes_val = np.repeat(depthes_val, 3, axis=2)
outputs_norm = change_channel(outputs_norm)
labels_val_norm = (labels_val_norm + 1) / 2
labels_val_norm = change_channel(labels_val_norm)
# if (i_val+1)%10 == 0:
misc.imsave(
pjoin(args.testset_out_path,
"{}_MS_hyb.png".format(i_val + 1)), outputs_norm)
misc.imsave(
pjoin(args.testset_out_path,
"{}_gt.png".format(i_val + 1)), labels_val_norm)
misc.imsave(
pjoin(args.testset_out_path,
"{}_in.jpg".format(i_val + 1)), images_val)
misc.imsave(
pjoin(args.testset_out_path,
"{}_depth.png".format(i_val + 1)), depthes_val)
# accumulate the metrics in matrix
if ((np.isnan(mean_i)) | (np.isinf(mean_i)) == False):
sum_mean.append(mean_i)
sum_median.append(median_i)
sum_small.append(small_i)
sum_mid.append(mid_i)
sum_large.append(large_i)
sum_num.append(pixelnum)
evalcount += 1
if (i_val + 1) % 10 == 0:
print(
"Iteration %d Evaluation Loss: mean %.4f, median %.4f, 11.25 %.4f, 22.5 %.4f, 30 %.4f"
% (i_val + 1, mean_i, median_i, small_i, mid_i,
large_i))
# Summarize the result
eval_print(sum_mean,
sum_median,
sum_small,
sum_mid,
sum_large,
sum_num,
item='Pixel-Level')
avg_mean = sum(sum_mean) / evalcount
sum_mean.append(avg_mean)
avg_median = sum(sum_median) / evalcount
sum_median.append(avg_median)
avg_small = sum(sum_small) / evalcount
sum_small.append(avg_small)
avg_mid = sum(sum_mid) / evalcount
sum_mid.append(avg_mid)
avg_large = sum(sum_large) / evalcount
sum_large.append(avg_large)
print(
"evalnum is %d, Evaluation Image-Level Mean Loss: mean %.4f, median %.4f, 11.25 %.4f, 22.5 %.4f, 30 %.4f"
% (evalcount, avg_mean, avg_median, avg_small, avg_mid,
avg_large))
sum_matrix = np.transpose(
[sum_mean, sum_median, sum_small, sum_mid, sum_large])
if args.model_full_name != '':
sum_file = args.model_full_name[:-4] + '.csv'
np.savetxt(pjoin(args.model_savepath, sum_file),
sum_matrix,
fmt='%.6f',
delimiter=',')
print("Saving to %s" % (sum_file))
# end of dataset test
else:
if os.path.isdir(args.out_path) == False:
os.mkdir(args.out_path)
print("Read Input Image from : {}".format(args.img_path))
for i in os.listdir(args.img_path):
if not i.endswith('.jpg'):
continue
print i
input_f = args.img_path + i
depth_f = args.depth_path + i[:-4] + '.png'
output_f = args.out_path + i[:-4] + '_rgbd.png'
img = misc.imread(input_f)
orig_size = img.shape[:-1]
if args.img_rot:
img = np.transpose(img, (1, 0, 2))
img = np.flipud(img)
img = misc.imresize(img, (
args.img_cols,
args.img_rows)) # Need resize the image to model inputsize
else:
img = misc.imresize(img, (
args.img_rows,
args.img_cols)) # Need resize the image to model inputsize
img = img.astype(np.float)
if args.img_norm:
img = (img - 128) / 255
# NHWC -> NCHW
img = img.transpose(2, 0, 1)
img = np.expand_dims(img, 0)
img = torch.from_numpy(img).float()
if args.img_rot:
depth = png_reader_32bit(depth_f,
(args.img_rows, args.img_cols))
depth = np.transpose(depth, (1, 0))
depth = np.flipud(depth)
# valid = png_reader_uint8(mask_f, (args.img_rows,args.img_cols))
# valid = np.transpose(valid, (1,0))
# valid = np.flipud(valid)
else:
depth = png_reader_32bit(depth_f,
(args.img_rows, args.img_cols))
# valid = png_reader_uint8(mask_f, (args.img_rows,args.img_cols))
depth = depth.astype(float)
# Please change to the scale so that scaled_depth=1 corresponding to real 10m depth
# matterpot depth=depth/40000 scannet depth=depth/10000
depth = depth / (args.d_scale)
if depth.ndim == 3: # to dim 2
depth = depth[:, :, 0]
# if valid.ndim == 3: #to dim 2
# valid = valid[:,:,0]
# valid = 1-depth
# valid[valid>1] = 1
valid = (depth > 0.0001).astype(float)
# valid = depth.astype(float)
depth = depth[np.newaxis, :, :]
depth = np.expand_dims(depth, 0)
valid = np.expand_dims(valid, 0)
depth = torch.from_numpy(depth).float()
valid = torch.from_numpy(valid).float()
if torch.cuda.is_available():
model_F.cuda()
model_F.eval()
if args.arch_map == 'map_conv':
model_map.cuda()
model_map.eval()
images = Variable(img.contiguous().cuda())
depth = Variable(depth.contiguous().cuda())
valid = Variable(valid.contiguous().cuda())
else:
images = Variable(img)
depth = Variable(depth)
valid = Variable(valid)
with torch.no_grad():
if args.arch_map == 'map_conv':
outputs_valid = model_map(
torch.cat((depth, valid[:, np.newaxis, :, :]), dim=1))
outputs, outputs1, outputs2, outputs3, output_d = model_F(
images, depth, outputs_valid.squeeze(1))
else:
outputs, outputs1, outputs2, outputs3, output_d = model_F(
images, depth, outputs_valid)
outputs_norm = norm_imsave(outputs)
outputs_norm = np.squeeze(outputs_norm.data.cpu().numpy(), axis=0)
# outputs_norm = misc.imresize(outputs_norm, orig_size)
outputs_norm = change_channel(outputs_norm)
misc.imsave(output_f, outputs_norm)
print("Complete")
def train(args):
writer = SummaryWriter(comment=args.writer)
# data loader setting, train and evaluatfconvion
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
t_loader = data_loader(data_path,
split='train',
img_size=(args.img_rows, args.img_cols),
img_norm=args.img_norm)
v_loader = data_loader(data_path,
split='test',
img_size=(args.img_rows, args.img_cols),
img_norm=args.img_norm)
trainloader = data.DataLoader(t_loader,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers)
evalloader = data.DataLoader(v_loader,
batch_size=args.batch_size,
num_workers=args.num_workers)
print("Finish Loader Setup")
# Setup Model and load pretrained model
model_name_F = args.arch_F
model_F = get_model(model_name_F, True) # concat and output
model_F = torch.nn.DataParallel(model_F,
device_ids=range(
torch.cuda.device_count()))
if args.resume: # use previously trained model
# Two types of resume training
checkpoint = torch.load(args.resume_model_path)
if checkpoint.has_key('model_RGB_state'):
# Resume with a RGB model.
# Load the RGB model and change it as encoder in the fusion network
# Only the first three conv block will be used.
state = get_fconv_premodel(model_F, checkpoint['model_RGB_state'])
else:
# Resume a RGBD model
state = checkpoint["model_F_state"]
model_F.load_state_dict(state)
model_F.cuda()
print("Finish model setup")
# Setup model for confidence map
if args.arch_map == 'map_conv':
model_name_map = args.arch_map
model_map = get_model(model_name_map, True) # concat and output
model_map = torch.nn.DataParallel(model_map,
device_ids=range(
torch.cuda.device_count()))
if args.resume and checkpoint.has_key("model_map_state"):
model_map.load_state_dict(checkpoint["model_map_state"])
model_map.cuda()
print("Finish model_map setup")
# optimizers and lr-decay setting
if args.resume:
optimizer_F = torch.optim.RMSprop(model_F.parameters(),
lr=0.1 * args.l_rate)
scheduler_F = torch.optim.lr_scheduler.MultiStepLR(
optimizer_F, milestones=[2, 4, 6, 9, 12], gamma=0.5)
if args.arch_map == 'map_conv':
optimizer_map = torch.optim.RMSprop(model_map.parameters(),
lr=0.1 * args.l_rate)
# scheduler_map = torch.optim.lr_scheduler.MultiStepLR(optimizer_map, milestones=[1, 3, 5, 7, 9, 11, 13], gamma=0.5)#second trial
scheduler_map = torch.optim.lr_scheduler.MultiStepLR(
optimizer_map, milestones=[2, 4, 6, 9, 12], gamma=0.5)
else:
optimizer_F = torch.optim.RMSprop(model_F.parameters(), lr=args.l_rate)
scheduler_F = torch.optim.lr_scheduler.MultiStepLR(
optimizer_F, milestones=[1, 3, 5, 8, 11, 15], gamma=0.5)
best_loss = 1
n_iter_t, n_iter_v = 0, 0
if args.dataset == 'matterport':
total_iter_t = 105432 / args.batch_size
elif args.dataset == 'scannet':
total_iter_t = 59743 / args.batch_size
else:
total_iter_t = 0
if not os.path.exists(args.model_savepath):
os.makedirs(args.model_savepath)
# forward and backward
for epoch in range(args.n_epoch):
scheduler_F.step()
model_F.train()
if args.arch_map == 'map_conv':
scheduler_map.step()
model_map.train()
for i, (images, labels, masks, valids, depthes,
meshdepthes) in enumerate(trainloader):
n_iter_t += 1
images = Variable(images.contiguous().cuda())
labels = Variable(labels.contiguous().cuda())
masks = Variable(masks.contiguous().cuda())
valids = Variable(valids.contiguous().cuda())
depthes = Variable(depthes.contiguous().cuda())
optimizer_F.zero_grad()
if args.arch_map == 'map_conv':
optimizer_map.zero_grad()
if args.arch_map == 'map_conv':
outputs_valid = model_map(
torch.cat((depthes, valids[:, np.newaxis, :, :]),
dim=1)) # 1c output
outputs, outputs1, outputs2, outputs3, output_d = model_F(
images, depthes, outputs_valid.squeeze(1))
else:
outputs, outputs1, outputs2, outputs3, output_d = model_F(
images, depthes, valids)
# outputs, outputs1, outputs2, outputs3, output_d = model_F(images, depthes, valids)
loss, df = get_lossfun(args.loss, outputs, labels, masks)
if args.hybrid_loss:
loss1, df1 = get_lossfun(args.loss, outputs1, labels, masks)
loss2, df2 = get_lossfun('cosine', outputs2, labels, masks)
loss3, df3 = get_lossfun('cosine', outputs3, labels, masks)
outputs.backward(gradient=df, retain_graph=True)
if args.hybrid_loss:
outputs1.backward(gradient=0.5 * df1, retain_graph=True)
outputs2.backward(gradient=0.25 * df2, retain_graph=True)
outputs3.backward(gradient=0.125 * df3)
optimizer_F.step()
if args.arch_map == 'map_conv':
optimizer_map.step()
if (i + 1) % 100 == 0:
if args.hybrid_loss:
print(
"Epoch [%d/%d] Iter [%d/%d] Loss: %.4f,%.4f,%.4f,%.4f "
% (epoch + 1, args.n_epoch, i, total_iter_t, loss.data,
loss1.data, loss2.data, loss3.data))
else:
print(
"Epoch [%d/%d] Iter [%d/%d] Loss: %.4f" %
(epoch + 1, args.n_epoch, i, total_iter_t, loss.data))
if (i + 1) % 250 == 0:
writer.add_scalar('loss/trainloss', loss.data.item(), n_iter_t)
writer.add_images('Image', images + 0.5, n_iter_t)
writer.add_images('Label',
0.5 * (labels.permute(0, 3, 1, 2) + 1),
n_iter_t)
writer.add_images(
'Depth',
np.repeat((depthes - torch.min(depthes)) /
(torch.max(depthes) - torch.min(depthes)),
3,
axis=1), n_iter_t)
# writer.add_image('Mesh_Depth', np.repeat((meshdepthes-torch.min(depthes))/(torch.max(depthes)-torch.min(depthes)), 3, axis = 1), n_iter_t)
outputs_n = norm_tf(outputs)
if (args.hybrid_loss):
outputs_n1 = norm_tf(outputs1)
outputs_n2 = norm_tf(outputs2)
outputs_n3 = norm_tf(outputs3)
output_nd = norm_tf(output_d)
writer.add_images('Output', outputs_n, n_iter_t)
if (args.hybrid_loss):
writer.add_images('Output1', outputs_n1, n_iter_t)
writer.add_images('Output2', outputs_n2, n_iter_t)
writer.add_images('Output3', outputs_n3, n_iter_t)
writer.add_images('Output_depth', output_nd, n_iter_t)
if args.arch_map == 'map_conv':
outputs_valid_1 = (outputs_valid - torch.min(outputs_valid)
) / (torch.max(outputs_valid) -
torch.min(outputs_valid))
writer.add_images('Output_Mask',
outputs_valid_1.repeat([1, 3, 1, 1]),
n_iter_t)
model_F.eval()
if args.arch_map == 'map_conv':
model_map.eval()
mean_loss, sum_loss = 0, 0
evalcount = 0
with torch.no_grad():
for i_val, (images_val, labels_val, masks_val, valids_val,
depthes_val,
meshdepthes_val) in tqdm(enumerate(evalloader)):
n_iter_v += 1
images_val = Variable(images_val.contiguous().cuda())
labels_val = Variable(labels_val.contiguous().cuda())
masks_val = Variable(masks_val.contiguous().cuda())
valids_val = Variable(valids_val.contiguous().cuda())
depthes_val = Variable(depthes_val.contiguous().cuda())
if args.arch_map == 'map_conv':
outputs_valid = model_map(
torch.cat(
(depthes_val, valids_val[:, np.newaxis, :, :]),
dim=1))
outputs, outputs1, outputs2, outputs3, output_d = model_F(
images_val, depthes_val, outputs_valid.squeeze(1))
else:
outputs, outputs1, outputs2, outputs3, output_d = model_F(
images_val, depthes_val, valids_val)
# outputs, outputs1, outputs2, outputs3,output_d = model_F(images_val, depthes_val, valids_val)
loss, df = get_lossfun(args.loss, outputs, labels_val,
masks_val,
False) # valid_val not used infact
if ((np.isnan(loss)) | (np.isinf(loss))):
sum_loss += 0
else:
sum_loss += loss
evalcount += 1
if (i_val + 1) % 250 == 0:
# print("Epoch [%d/%d] Evaluation Loss: %.4f" % (epoch+1, args.n_epoch, loss))
writer.add_scalar('loss/evalloss', loss, n_iter_v)
writer.add_images('Eval Image', images_val + 0.5, n_iter_t)
writer.add_images(
'Eval Label',
0.5 * (labels_val.permute(0, 3, 1, 2) + 1), n_iter_t)
writer.add_images(
'Eval Depth',
np.repeat(
(depthes_val - torch.min(depthes_val)) /
(torch.max(depthes_val) - torch.min(depthes_val)),
3,
axis=1), n_iter_t)
# writer.add_image('Eval Mesh Depth', np.repeat((meshdepthes_val-torch.min(depthes_val))/(torch.max(meshdepthes_val)-torch.min(depthes_val)), 3, axis = 1), n_iter_t)
outputs_n = norm_tf(outputs)
output_nd = norm_tf(output_d)
writer.add_images('Eval Output', outputs_n, n_iter_t)
writer.add_images('Eval Output_depth', output_nd, n_iter_t)
if args.arch_map == 'map_conv':
outputs_valid_1 = (outputs_valid -
torch.min(outputs_valid)) / (
torch.max(outputs_valid) -
torch.min(outputs_valid))
writer.add_images('Eval Output_Mask',
outputs_valid_1.repeat([1, 3, 1, 1]),
n_iter_t)
# mean_loss = sum_loss/(evalloader.__len__())
mean_loss = sum_loss / evalcount
print("Epoch [%d/%d] Evaluation Mean Loss: %.4f" %
(epoch + 1, args.n_epoch, mean_loss))
writer.add_scalar('loss/evalloss_mean', mean_loss, epoch)
if mean_loss < best_loss: # if (epoch+1)%20 == 0:
best_loss = mean_loss
if args.hybrid_loss:
other_info = 'hybrid'
else:
other_info = 'nohybrid'
if args.resume:
other_info += '_resume'
if args.arch_map == 'map_conv':
state = {
'epoch': epoch + 1,
'model_F_state': model_F.state_dict(),
'optimizer_F_state': optimizer_F.state_dict(),
'model_map_state': model_map.state_dict(),
'optimizer_map_state': optimizer_map.state_dict(),
}
torch.save(
state,
pjoin(
args.model_savepath, "{}_{}_{}_{}_{}_best.pkl".format(
args.arch_F, args.dataset, args.loss,
args.model_num, other_info)))
# For retrain purpose
# torch.save(state, pjoin(args.model_savepath,
# "{}_{}_{}_{}_{}_{}_best.pkl".format(args.arch_F, args.dataset, args.loss,
# 'e' + str(epoch), args.model_num,
# other_info)))
else:
state = {
'epoch': epoch + 1,
'model_F_state': model_F.state_dict(),
'optimizer_F_state': optimizer_F.state_dict(),
}
torch.save(
state,
pjoin(
args.model_savepath, "{}_{}_{}_{}_{}_best.pkl".format(
args.arch_F, args.dataset, args.loss,
args.model_num, other_info)))
# state = {'epoch': epoch+1,
# 'model_F_state': model_F.state_dict(),
# 'optimizer_F_state' : optimizer_F.state_dict(),}
# torch.save(state, pjoin(args.model_savepath, "{}_{}_{}_{}_best.pkl".format(args.arch_F, args.dataset, args.loss, args.model_num)))
print('Finish training for dataset %s trial %s' %
(args.dataset, args.model_num))
writer.export_scalars_to_json("./{}_{}_{}_{}.json".format(
args.arch_F, args.dataset, args.loss, args.model_num))
writer.close()
def train(args):
writer = SummaryWriter(comment=args.writer)
# data loader setting, train and evaluat fconvion
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
t_loader = data_loader(data_path, split='train', img_size=(args.img_H, args.img_W), img_norm=args.img_norm)
v_loader = data_loader(data_path, split='test', img_size=(args.img_H, args.img_W), img_norm=args.img_norm)
trainloader = data.DataLoader(t_loader, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers)
evalloader = data.DataLoader(v_loader, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers)
print("Finish Loader Setup")
# Setup Model and load pretrained model
model_normal = args.arch_N
model_N = get_model(model_normal, True) # concat and output
model_N = torch.nn.DataParallel(model_N, device_ids=range(torch.cuda.device_count()))
model_N.cuda()
state = torch.load(args.PATH)
# model_N.load_state_dict(state)
print("Finish model setup")
# Setup depthmodel
depth_model_name = args.arch_D
DepthModel = get_model(depth_model_name, True) # concat and output
DepthModel = torch.nn.DataParallel(DepthModel, device_ids=range(torch.cuda.device_count()))
DepthModel.cuda()
# DepthModel.load_state_dict(torch.load(args.PATH))
print("Finish DepthModel setup")
# optimizers and lr-decay setting
optimizer_N = torch.optim.RMSprop(model_N.parameters(), lr=args.l_rate)
scheduler_N = torch.optim.lr_scheduler.MultiStepLR(optimizer_N, milestones=[10, 20, 40, 60], gamma=0.2)
# optimizer_P = torch.optim.RMSprop(model_point.parameters(), lr=args.l_rate)
# scheduler_P = torch.optim.lr_scheduler.MultiStepLR(optimizer_P, milestones=[10, 20, 30, 40], gamma=0.5)
optimizer_D = torch.optim.RMSprop(DepthModel.parameters(), lr=0.1 * args.l_rate)
scheduler_D = torch.optim.lr_scheduler.MultiStepLR(optimizer_D, milestones=[10, 20, 40, 60], gamma=0.2)#second trial
best_loss = 1
n_iter_t, n_iter_v = 0, 0
cos = nn.CosineSimilarity(dim=1, eps=0)
total_iter_t = len(t_loader)
loss_func = ModelLoss(fx=args.fx, fy=args.fy, img_size=(args.img_H, args.img_W))
if not os.path.exists(args.model_savepath):
os.makedirs(args.model_savepath)
# forward and backward
for epoch in range(args.n_epoch):
model_N.train()
DepthModel.train()
# training
for i, (images, depths, normal) in enumerate(trainloader):
n_iter_t += 1
images = Variable(images.contiguous().cuda())
normal = Variable(normal.contiguous().cuda())
depths = Variable(depths.contiguous().cuda())
optimizer_N.zero_grad()
optimizer_D.zero_grad()
pca_normal, pointnet_feature= DepthModel(depths) #
coarse_normal, refined_normal = model_N(images, pca_normal, pointnet_feature)
# result_img = merge_into_row_with_gt(images, depths, normal, coarse_normal, refined_normal)
# save_image(result_img, args.model_savepath+'/00result'+str(i)+'.png')
# loss_depth = loss_func.criterion(output_depth, depths)
# loss_depth = torch.log(torch.abs(output_depth - depthes) + 0.5).mean()
# loss_depth = total_loss(output_depth, depthes)
# loss_pca_normal = torch.min(torch.sqrt(((normal.reshape(-1, 3) - pca_normal.reshape(-1, 3)) ** 2).sum(-1)),
# torch.sqrt(((normal.reshape(-1, 3) + pca_normal.reshape(-1, 3)) ** 2).sum(-1))).mean()
loss_cN = torch.abs(1 - cos(coarse_normal, normal)).mean()
loss_rN = torch.abs(1 - cos(refined_normal, normal)).mean()
loss = loss_rN + loss_cN
loss.backward()
optimizer_N.step()
scheduler_N.step()
optimizer_D.step()
scheduler_D.step()
if (i+1) % 100== 0:
# print("Train: Epoch [%d/%d], Iter [%d/%d], depth loss: %.4f" % (
# epoch + 1, args.n_epoch, i, total_iter_t, loss.data.item()))
print("Train: Epoch [%d/%d], Iter [%d/%d], Loss: %.4f, "
"coarse normal: %.4f, refine normal: %.4f" % (
epoch + 1, args.n_epoch, i+1, total_iter_t, loss.data.item(),
loss_cN.data.item(), loss_rN.data.item()))
if (i + 1) % 2000 == 0:
writer.add_scalar('loss/trainloss', loss.data.item(), n_iter_t)
writer.add_images('RGB', images.detach().cpu().numpy() + 0.5, n_iter_t)
writer.add_images('Depth_GT',
np.repeat(((depths - torch.min(depths)) / (torch.max(depths) - torch.min(depths))).detach().cpu().numpy(),
3, axis=1), n_iter_t)
writer.add_images('Normal_GT', 0.5 * (normal.detach().cpu().numpy() + 1), n_iter_t)
# writer.add_images('Pred_Depth',
# np.repeat(((output_depth - torch.min(output_depth)) / (torch.max(output_depth) - torch.min(output_depth))).detach().cpu().numpy(),
# 3, axis=1), n_iter_t)
writer.add_images('Coarse_Normal', 0.5 * (coarse_normal.detach().cpu().numpy() + 1), n_iter_t)
writer.add_images('Refined_Normal', 0.5 * (refined_normal.detach().cpu().numpy() + 1), n_iter_t)
# output_nd = norm_tf(output_depth)
# outputs_cn = norm_tf(coarse_normal)
# output_rd = norm_tf(refined_normal)
model_N.eval()
DepthModel.eval()
mean_loss, sum_loss = 0, 0
evalcount = 0
# errorSum = {'MSE': 0, 'RMSE': 0, 'ABS_REL': 0, 'LG10': 0,
# 'MAE': 0, 'DELTA1': 0, 'DELTA2': 0, 'DELTA3': 0,
# 'mean': 0, 'median': 0, '11.25': 0, '22.5': 0, '30': 0}
# evaluation
with torch.no_grad():
for i_val, (images_val, depths_val, normals_val) in tqdm(enumerate(evalloader)):
n_iter_v += 1
images_val = Variable(images_val.contiguous().cuda())
normals_val = Variable(normals_val.contiguous().cuda())
depths_val = Variable(depths_val.contiguous().cuda())
pca_normal, pointnet_feature = DepthModel(depths_val)
coarse_normal_val, refined_normal_val = model_N(images_val, pca_normal, pointnet_feature)
# loss_depth = torch.log(torch.abs(output_depth_val - depthes_val) + 0.5).mean()
# loss_depth = loss_func.criterion(output_depth_val, depths_val)
# loss_pca_normal = torch.min(torch.sqrt(((normals_val - pca_normal_val) ** 2).sum(-1)),
# torch.sqrt(((normals_val + pca_normal_val) ** 2).sum(-1))).mean()
loss_cN = torch.abs(1 - cos(coarse_normal_val, normals_val)).mean()
loss_rN = torch.abs(1 - cos(refined_normal_val, normals_val)).mean()
loss = loss_cN + loss_rN
if np.isnan(loss.detach().cpu().numpy()) | np.isinf(loss.detach().cpu().numpy()):
sum_loss += 0
else:
sum_loss += loss.data.item()
evalcount += 1
if (i_val+1) % 100==0:
# errors
errors = evaluateError(refined_normal_val, normals_val)
# errorSum = addErrors(errorSum, errors, args.batch_size)
# averageError = averageErrors(errorSum, args.batch_size)
print('metrics:')
print(errors)
if (i_val+1) % 1000 == 0:
# print("Epoch [%d/%d] Evaluation Loss: %.4f" % (epoch+1, args.n_epoch, loss))
writer.add_scalar('loss/evalloss', loss.data.item(), n_iter_v)
writer.add_images('Eval_RGB', images_val + 0.5, n_iter_v)
writer.add_images('Eval_Depth_GT',
np.repeat(
((depths_val - torch.min(depths_val)) / (torch.max(depths_val) - torch.min(depths_val))).detach().cpu().numpy(),
3, axis=1), n_iter_v)
writer.add_images('Eval_Normal_GT', 0.5 * (normals_val.detach().cpu().numpy() + 1), n_iter_v)
# writer.add_images('Eval_Pred_Depth',
# np.repeat(((output_depth_val - torch.min(output_depth_val)) / (
# torch.max(output_depth_val) - torch.min(output_depth_val))).detach().cpu().numpy(),
# 3, axis=1), n_iter_v)
writer.add_images('Eval_Coarse_Normal', 0.5 * (coarse_normal_val.detach().cpu().numpy() + 1), n_iter_v)
writer.add_images('Eval_Refined_Normal', 0.5 * (refined_normal_val.detach().cpu().numpy() + 1), n_iter_v)
mean_loss = sum_loss / evalcount
print("Epoch [%d/%d] Evaluation Mean Loss: %.4f" % (epoch + 1, args.n_epoch, mean_loss))
writer.add_scalar('loss/evalloss_mean', mean_loss, epoch)
# if mean_loss < best_loss:
# best_loss = mean_loss
# state = {'epoch': epoch + 1,
# 'model_N_state': model_N.state_dict(),
# 'optimizer_N_state': optimizer_N.state_dict(),
# 'DepthModel_state': DepthModel.state_dict(),
# 'optimizer_depth_state': optimizer_D.state_dict(), }
# torch.save(state, pjoin(args.model_savepath,
# "{}_{}_{}_best.pkl".format(args.arch_N, args.dataset, args.loss)))
print('Finish training for dataset %s trial %s' % (args.dataset, args.model_num))
writer.export_scalars_to_json("./{}_{}_{}.json".format(args.arch_N, args.dataset, args.loss))
writer.close()
# gt: lb, h*w*3
# mask: gt!=0,h*w
# valid: rawdepth!=0, h*w
# rawdepth: depth with hole, 1*h*w
# meshdepth: depth with hole, 1*h*w
return im, torch.ones(
(3, self.img_size[0],
self.img_size[1])), torch.ones(self.img_size), valid, depth, depth
# Leave code for debugging purposes
if __name__ == '__main__':
# Config your local data path
from loader import get_data_path
local_path = get_data_path('sevenscenes')
bs = 5
dst = sevenScenesLoader(root=local_path)
testloader = data.DataLoader(dst, batch_size=bs)
for i, data in enumerate(testloader):
imgs, labels, masks, valids, depths, meshdepths = data
imgs = imgs.numpy()
imgs = np.transpose(imgs, [0, 2, 3, 1])
imgs = imgs + 0.5
labels = labels.numpy()
labels = 0.5 * (labels + 1)
masks = masks.numpy()
masks = np.repeat(masks[:, :, :, np.newaxis], 3, axis=3)