def write_imgs_2(img_tuple, epoch, type_tuple=None, name_tuple=None, training=True):
print("Writing Images to ", output_dir)
if training:
cmap, uv, ab, bg, nor, dep, bg2, ori_gt,\
cmap_gt, uv_gt, ab_gt, bg_gt, nor_gt, dep_gt = img_tuple
else:
cmap, uv, ab, bg, nor, dep, bg2, ori_gt = img_tuple
print_img_with_reprocess(uv , "uv" , fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "uv.jpg"))
print_img_with_reprocess(ab , "ab" , fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "ab.jpg"))
print_img_with_reprocess(bg , "bg" , fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "bg.jpg"))
print_img_with_reprocess(bg2, "bg" , fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "bg2.jpg"))
#reprocess_np_auto(cmap, "")
print_img_with_reprocess(cmap, "exr" , fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "cmap.jpg")) #
print_img_with_reprocess(nor , "exr" , fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "nor.jpg")) #
print_img_with_reprocess(dep , "exr" , fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "dep.jpg")) #
print_img_with_reprocess(ori_gt, "ori" , fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "ori_gt.jpg"))
if training:
print_img_with_reprocess(uv_gt , "uv" , fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "uv_gt.jpg"))
print_img_with_reprocess(ab_gt , "ab" , fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "ab_gt.jpg"))
print_img_with_reprocess(bg_gt , "bg" , fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "bg_gt.jpg"))
print_img_with_reprocess(cmap_gt, "exr" , fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "cmap_gt.jpg")) #
print_img_with_reprocess(nor_gt , "exr" , fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "nor_gt.jpg")) #
print_img_with_reprocess(dep_gt , "exr" , fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "dep_gt.jpg")) #
print_img_with_reprocess(gt_clip(cmap_gt), "exr" , fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "cmap_gt2.jpg")) #
print_img_with_reprocess(gt_clip(nor_gt ), "exr" , fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "nor_gt2.jpg")) #
print_img_with_reprocess(gt_clip(dep_gt ), "exr" , fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "dep_gt2.jpg")) #
uv = reprocess_auto(uv, "uv")
bg2 = reprocess_auto(bg2, "bg")
ori_gt = reprocess_auto(ori_gt, "ori")
bw = uv2backward_trans_3(uv, bg2)
dewarp = bw_mapping_single_3(ori_gt, bw)
if training:
uv_gt = reprocess_auto(uv_gt, "uv")
bg_gt = reprocess_auto(bg_gt, "bg")
bw_gt = uv2backward_trans_3(uv_gt, bg_gt)
bw2 = uv2backward_trans_3(uv, bg_gt)
dewarp_gt = bw_mapping_single_3(ori_gt, bw_gt)
dewarp2 = bw_mapping_single_3(ori_gt, bw2)
print_img_auto(bw, "bw", fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "bw.jpg"))
print_img_auto(dewarp, "ori", fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "dewarp.jpg"))
if training:
print_img_auto(bw_gt,"bw", fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "bw_gt.jpg"))
print_img_auto(bw2, "bw", fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "bw2.jpg"))
print_img_auto(dewarp_gt,"ori", fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "dewarp_gt.jpg"))
print_img_auto(dewarp2, "ori", fname=tfilename(output_dir,"imgshow/epoch_{}".format(epoch), "dewarp2.jpg"))
def __getitem__(self, index):
npy_path = self.npy_list[index]
"""loading"""
# data = np.load(self.npy_dir + '/' + npy_name, allow_pickle=True)[()]
data = np.load(npy_path, allow_pickle=True)[()]
ori = data['ori']
ab = data['ab']
# bmap = data['bmap']
depth = data['depth']
normal = data['normal']
uv = data['uv']
cmap = data['cmap']
background = data['background']
if self.load_mod == "all":
return torch.from_numpy(ori), \
torch.from_numpy(ab), \
torch.from_numpy(depth), \
torch.from_numpy(normal), \
torch.from_numpy(cmap), \
torch.from_numpy(uv), \
torch.from_numpy(background)
# ori_1080 = data['ori_1080']
elif self.load_mod == "uvbw":
### ----------- Generate BW map ---------------
uv_2 = data_process.reprocess_auto(uv.transpose((1, 2, 0)), "uv")
mask = data_process.reprocess_auto(background.transpose((1, 2, 0)),
"background")
bw = uv2bw.uv2backward_trans_3(uv_2, mask)
bw = data_process.process_auto(bw, "bw")
bw = bw.transpose((2, 0, 1))
### --------------------------------------------
return torch.from_numpy(cmap), \
torch.from_numpy(uv), \
torch.from_numpy(bw), \
torch.from_numpy(background)
elif self.load_mod == "test_uvbw_mapping":
### ----------- Generate BW map ---------------
uv_2 = data_process.reprocess_auto(uv.transpose((1, 2, 0)), "uv")
mask = data_process.reprocess_auto(background.transpose((1, 2, 0)),
"background")
bw = uv2bw.uv2backward_trans_3(uv_2, mask)
bw = data_process.process_auto(bw, "bw")
bw = bw.transpose((2, 0, 1))
### --------------------------------------------
return torch.from_numpy(cmap), \
torch.from_numpy(uv), \
torch.from_numpy(bw), \
torch.from_numpy(background), \
torch.from_numpy(ori)
elif self.load_mod == "deform_cmap":
### ----------- Generate BW map ---------------
uv_2 = data_process.reprocess_auto(uv.transpose((1, 2, 0)), "uv")
mask = data_process.reprocess_auto(background.transpose((1, 2, 0)),
"background")
bw = uv2bw.uv2backward_trans_3(uv_2, mask)
bw2 = data_process.process_auto(bw, "bw")
bw2 = bw2.transpose((2, 0, 1))
deform = bw2deform.bw2deform(bw)
deform = data_process.process_auto(deform, "deform")
deform = deform.transpose((2, 0, 1))
### --------------------------------------------
return torch.from_numpy(cmap), \
torch.from_numpy(uv), \
torch.from_numpy(deform), \
torch.from_numpy(bw2), \
torch.from_numpy(ori), \
torch.from_numpy(background)
elif self.load_mod == "all_deform":
### ----------- Generate BW map ---------------
uv_2 = data_process.reprocess_auto(uv.transpose((1, 2, 0)), "uv")
mask = data_process.reprocess_auto(background.transpose((1, 2, 0)),
"background")
bw = uv2bw.uv2backward_trans_3(uv_2, mask)
bw2 = data_process.process_auto(bw, "bw")
bw2 = bw2.transpose((2, 0, 1))
deform = bw2deform.bw2deform(bw)
deform = data_process.process_auto(deform, "deform")
deform = deform.transpose((2, 0, 1))
# -----------------------------------------
return torch.from_numpy(ori), \
torch.from_numpy(ab), \
torch.from_numpy(depth), \
torch.from_numpy(normal), \
torch.from_numpy(cmap), \
torch.from_numpy(uv), \
torch.from_numpy(deform), \
torch.from_numpy(background)
def __getitem__(self, index):
image_name = self.image_name[index]
image_dir = self.image_dir
if True:
"""loading"""
ori = cv2.imread(image_dir + 'img/' + image_name)
ab = cv2.imread(image_dir + 'albedo/' + image_name)
ab = cv2.cvtColor(ab, cv2.COLOR_BGR2GRAY)
ab = ab[:, :, np.newaxis]
#depth = cv2.imread(self.image_dir + 'depth/' + image_name)
dep = cv2.imread(image_dir + 'depth/' + image_name[:-3] + 'exr',
cv2.IMREAD_UNCHANGED)
nor = cv2.imread(image_dir + 'shader_normal/' + image_name[:-3] +
'exr', cv2.IMREAD_UNCHANGED) #[-1,1]
cmap = cv2.imread(image_dir + '3dmap/' + image_name[:-3] + 'exr',
cv2.IMREAD_UNCHANGED) #[0,1]
uv = cv2.imread(image_dir + 'uv/' + image_name[:-3] + 'exr',
cv2.IMREAD_UNCHANGED) #[0,1]
uv = uv[:, :, 1:]
bg = cv2.threshold(dep[:, :, 0], 0, 1, cv2.THRESH_BINARY)[1]
#print("shape uv", uv.shape)
if self.load_mod == "original":
print("output original data")
return ori, ab, dep, nor, cmap, uv, bg
if self.load_mod == "nobw":
"""processing"""
ori, ab, dep, nor, cmap, uv, bg = augment(
(ori, ab, dep, nor, cmap, uv, bg), method=["crop"], imsize=256)
ori = process_auto(ori, "ori")
ab = process_auto(ab, "ab")
dep = process_auto(dep, "exr") #
nor = process_auto(nor, "exr") #
cmap = process_auto(cmap, "exr") #
uv = process_auto(uv, "uv")
bg = process_auto(bg, "bg")
return torch.from_numpy(ori.transpose((2,0,1))).float(), \
torch.from_numpy(ab.transpose((2,0,1))).float(), \
torch.from_numpy(dep.transpose((2,0,1))).float(), \
torch.from_numpy(nor.transpose((2,0,1))).float(), \
torch.from_numpy(cmap.transpose((2,0,1))).float(), \
torch.from_numpy(uv.transpose((2,0,1))).float(), \
torch.from_numpy(bg.transpose((2,0,1))).float(), \
if self.load_mod == "all":
"""processing"""
ori, ab, dep, nor, cmap, uv, bg = augment(
(ori, ab, dep, nor, cmap, uv, bg), method=["crop"], imsize=256)
bw = uv2bw.uv2backward_trans_3(uv, bg)
ori = process_auto(ori, "ori")
ab = process_auto(ab, "ab")
dep = process_auto(dep, "exr")
nor = process_auto(nor, "exr")
cmap = process_auto(cmap, "exr")
uv = process_auto(uv, "uv")
bg = process_auto(bg, "bg")
bw = process_auto(bw, "bw")
# ori,ab,dep,nor,cmap,uv,bg,bw
return torch.from_numpy(ori.transpose((2,0,1))).float(), \
torch.from_numpy(ab.transpose((2,0,1))).float(), \
torch.from_numpy(dep.transpose((2,0,1))).float(), \
torch.from_numpy(nor.transpose((2,0,1))).float(), \
torch.from_numpy(cmap.transpose((2,0,1))).float(), \
torch.from_numpy(uv.transpose((2,0,1))).float(), \
torch.from_numpy(bg.transpose((2,0,1))).float(), \
torch.from_numpy(bw.transpose((2,0,1))).float(), \
if self.load_mod == "new_ab":
ori, ab, dep, nor, cmap, uv, bg = augment(
(ori, ab, dep, nor, cmap, uv, bg), method=[], imsize=256)
#print("ab121", ab.shape)
ori = process_auto(ori, "ori")
ab = process_auto(ab, "ab")
dep = process_auto(dep, "exr")
nor = process_auto(nor, "exr")
cmap = process_auto(cmap, "exr")
uv = process_auto(uv, "uv")
bg = process_auto(bg, "bg")
ori2 = cv2.cvtColor(ori, cv2.COLOR_BGR2GRAY)[:, :, np.newaxis]
ab_diff = np.where(bg > 0, (ori2 - ab), 0)
ab_diff = process_auto(ab_diff, "ab")
return torch.from_numpy(ori.transpose((2,0,1))).float(), \
torch.from_numpy(ab_diff.transpose((2,0,1))).float(), \
torch.from_numpy(dep.transpose((2,0,1))).float(), \
torch.from_numpy(nor.transpose((2,0,1))).float(), \
torch.from_numpy(cmap.transpose((2,0,1))).float(), \
torch.from_numpy(uv.transpose((2,0,1))).float(), \
torch.from_numpy(bg.transpose((2,0,1))).float(), \
if self.load_mod == "extra_bg":
if random.random() > 0.66: #
ori, ab, dep, nor, cmap, uv, bg = get_composed_imgs(
(ori, ab, dep, nor, cmap, uv, bg))
ori, ab, dep, nor, cmap, uv, bg = augment(
(ori, ab, dep, nor, cmap, uv, bg),
method=["crop"],
imsize=256,
crop_rate=0.7)
else:
ori, ab, dep, nor, cmap, uv, bg = augment(
(ori, ab, dep, nor, cmap, uv, bg),
method=["crop"],
imsize=256,
crop_rate=0.9)
# import ipdb; ipdb.set_trace()
# return ori,ab,dep,nor,cmap,uv,bg
ori = process_auto(ori, "ori")
ab = process_auto(ab, "ab")
dep = process_auto(dep, "exr") #
nor = process_auto(nor, "exr") #
cmap = process_auto(cmap, "exr") #
uv = process_auto(uv, "uv")
bg = process_auto(bg, "bg")
# ab-diff
ori2 = cv2.cvtColor(ori, cv2.COLOR_BGR2GRAY)[:, :, np.newaxis]
ab_diff = np.where(bg > 0, (ori2 - ab), 0)
ab_diff = process_auto(ab_diff, "ab")
return torch.from_numpy(ori.transpose((2,0,1))).float(), \
torch.from_numpy(ab_diff.transpose((2,0,1))).float(), \
torch.from_numpy(dep.transpose((2,0,1))).float(), \
torch.from_numpy(nor.transpose((2,0,1))).float(), \
torch.from_numpy(cmap.transpose((2,0,1))).float(), \
torch.from_numpy(uv.transpose((2,0,1))).float(), \
torch.from_numpy(bg.transpose((2,0,1))).float(),
raise ValueError
def test_single(model, imgpath, writer):
test_name = generate_name()
img_ori = cv2.imread(imgpath)
parent, imgname = os.path.split(imgpath)
img = cv2.resize(img_ori, (256, 256))
_input = process_auto(img, "ori")
img_tensor = process_to_tensor(_input[np.newaxis, :, :, :])
img_tensor = img_tensor.cuda()
uv_map, cmap, nor_map, alb_map, dep_map, mask_map, \
_, _, _, _, _, deform_map = model(img_tensor)
alb_pred = alb_map[0, :, :, :]
uv_pred = uv_map[0, :, :, :]
mask_pred = mask_map[0, :, :, :]
mask_pred = torch.round(mask_pred)
cmap_pred = cmap[0, :, :, :]
dep_pred = dep_map[0, :, :, :]
nor_pred = nor_map[0, :, :, :]
deform_bw_map = deform2bw_tensor_batch(deform_map.detach().cpu())
deform_bw_pred = deform_bw_map[0, :, :, :]
uv_np = reprocess_auto(uv_pred, "uv")
mask_np = reprocess_auto(mask_pred, "background")
alb_np = reprocess_auto(alb_pred, "ab")
# cmap_np = reprocess_auto(cmap_pred, "cmap")
bw_np = uv2backward_trans_3(uv_np, mask_np)
dewarp_np = bw_mapping_single_3(img, bw_np)
bw_large = blur_bw_np(bw_np, img_ori)
alb_large, diff, ori_gray, img_gray = resize_albedo_np2(
img_ori, img, alb_np)
dewarp_ab_large = bw_mapping_single_3(alb_large, bw_large)
dewarp_large = bw_mapping_single_3(img_ori, bw_large)
print("shape: ", dewarp_ab_large.shape)
print_img_auto(ori_gray,
"ori",
fname=tfilename("output_test_single/",
"real-img-" + test_name,
imgname + "_ori_gray.jpg"))
print_img_auto(img_gray,
"ori",
fname=tfilename("output_test_single/",
"real-img-" + test_name,
imgname + "_img_gray.jpg"))
print_img_auto(alb_np,
"ab",
fname=tfilename("output_test_single/",
"real-img-" + test_name,
imgname + "_ab_np.jpg"))
print_img_auto(uv_np,
"uv",
fname=tfilename("output_test_single/",
"real-img-" + test_name,
imgname + "_uv_np.jpg"))
print_img_auto(bw_np,
"bw",
fname=tfilename("output_test_single/",
"real-img-" + test_name,
imgname + "_bw_np.jpg"))
print_img_auto(img_ori,
"ori",
fname=tfilename("output_test_single/",
"real-img-" + test_name,
imgname + "_imgori.jpg"))
print_img_auto(img,
"ori",
fname=tfilename("output_test_single/",
"real-img-" + test_name,
imgname + "_ori.jpg"))
print_img_auto(dewarp_np,
"ori",
fname=tfilename("output_test_single/",
"real-img-" + test_name,
imgname + "_dewarp_ori.jpg"))
print_img_auto(dewarp_ab_large,
"ab",
fname=tfilename("output_test_single/",
"real-img-" + test_name,
imgname + "_large_ab.jpg"))
print_img_auto(dewarp_large,
"ori",
fname=tfilename("output_test_single/",
"real-img-" + test_name,
imgname + "_large_ori.jpg"))
print_img_auto(alb_large,
"ab",
fname=tfilename("output_test_single/",
"real-img-" + test_name,
imgname + "_al_large.jpg"))
print_img_auto(diff,
"ab",
fname=tfilename("output_test_single/",
"real-img-" + test_name,
imgname + "_diff_large.jpg"))