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"))