def main(): parser = argparse.ArgumentParser(description='Find latent representation of reference images using perceptual loss') parser.add_argument('--batch_size', default=1, help='Batch size for generator and perceptual model', type=int) parser.add_argument('--resolution',default=1024,type=int) parser.add_argument('--weight_file',default="weight_files/pytorch/karras2019stylegan-ffhq-1024x1024.pt",type=str) parser.add_argument('--latent_file',default="latent_W/0.npy") args=parser.parse_args() g_all = nn.Sequential(OrderedDict([ ('g_mapping', G_mapping()), #('truncation', Truncation(avg_latent)), ('g_synthesis', G_synthesis(resolution=args.resolution)) ])) g_all.load_state_dict(torch.load(args.weight_file, map_location=device)) g_all.eval() g_all.to(device) g_mapping,g_synthesis=g_all[0],g_all[1] boundary_name=["stylegan_ffhq_gender_w_boundary.npy","stylegan_ffhq_age_w_boundary.npy","stylegan_ffhq_pose_w_boundary.npy","stylegan_ffhq_eyeglasses_w_boundary.npy","stylegan_ffhq_smile_w_boundary.npy"] semantic=["gender","age","pose","eye_glass","smile"] for i in range(5): latents_0=np.load(args.latent_file) latents_0=torch.tensor(latents_0).to(device)#.unsqueeze(0) boundary=np.load("boundaries/"+boundary_name[i]) make_morph(boundary,i,latents_0,g_synthesis,semantic)
def main(): parser = argparse.ArgumentParser( description= 'Find latent representation of reference images using perceptual loss') parser.add_argument('--batch_size', default=1, help='Batch size for generator and perceptual model', type=int) parser.add_argument('--resolution', default=1024, type=int) parser.add_argument( '--weight_file', default="weight_files/pytorch/karras2019stylegan-ffhq-1024x1024.pt", type=str) parser.add_argument('--latent_file1', default="latent_W/0.npy") parser.add_argument('--latent_file2', default="latent_W/sample.npy") args = parser.parse_args() g_all = nn.Sequential( OrderedDict([ ('g_mapping', G_mapping()), #('truncation', Truncation(avg_latent)), ('g_synthesis', G_synthesis(resolution=args.resolution)) ])) g_all.load_state_dict(torch.load(args.weight_file, map_location=device)) g_all.eval() g_all.to(device) g_mapping, g_synthesis = g_all[0], g_all[1] latents_0 = np.load(args.latent_file1) latents_1 = np.load(args.latent_file2) latents_0 = torch.tensor(latents_0).to(device) latents_1 = torch.tensor(latents_1).to(device) for i in range(100): alpha = (1 / 100) * i latents = alpha * latents_0 + (1 - alpha) * latents_1 synth_img = g_synthesis(latents) synth_img = (synth_img + 1.0) / 2.0 save_image(synth_img.clamp(0, 1), "morph_result/encode1/{}.png".format(i))
def image_crossover_face(BASE_DIR, RAW_DIR, rand_uuid, process_selection, gender): ALIGNED_IMAGE_DIR = f'{BASE_DIR}aligned/' os.mkdir(ALIGNED_IMAGE_DIR) TARGET_IMAGE_DIR = f'{BASE_DIR}target/' os.mkdir(TARGET_IMAGE_DIR) if process_selection == 0 and gender == 'female': TARGET_SOURCE_DIR = '../image_2_style_gan/source/target/female/' elif process_selection == 0 and gender == 'male': TARGET_SOURCE_DIR = '../image_2_style_gan/source/target/male/' else: TARGET_SOURCE_DIR = f'{BASE_DIR}raw_target/aligned/' FINAL_IMAGE_DIR = f'{BASE_DIR}final/' os.mkdir(FINAL_IMAGE_DIR) MASK_DIR = f'{BASE_DIR}mask/' os.mkdir(MASK_DIR) model_resolution = 1024 ITERATION = 150 aligned_image_names = align_images(RAW_DIR, ALIGNED_IMAGE_DIR) try: if not aligned_image_names: print('\nFailed Face Alignment. Abort process.') shutil.rmtree(BASE_DIR) # UUID 디렉터리 삭제 sys.exit(1) aligned_image_name = ALIGNED_IMAGE_DIR + os.listdir( ALIGNED_IMAGE_DIR)[0] mask_maker(aligned_image_name, MASK_DIR) ingredient_name = ALIGNED_IMAGE_DIR + os.listdir(ALIGNED_IMAGE_DIR)[0] random_target_image_index = random.randint( 0, len(os.listdir(TARGET_SOURCE_DIR)) - 1) target_name = TARGET_SOURCE_DIR + os.listdir( TARGET_SOURCE_DIR)[random_target_image_index] target_number = os.listdir( TARGET_SOURCE_DIR)[random_target_image_index] print(f"target_number: {target_number}") except IndexError as e: print( "\nMissing file(s).\nCheck if all of source images prepared properly and try again." ) print(f"Aligned_image_names function: {e}") shutil.rmtree(BASE_DIR) # UUID 디렉터리 삭제 sys.exit(1) try: mask_name = MASK_DIR + os.listdir(MASK_DIR)[0] except Exception as e: shutil.copyfile('../image_2_style_gan/source/ref_mask/ref_mask.png', '{}ref_mask.png'.format(MASK_DIR)) mask_name = MASK_DIR + os.listdir(MASK_DIR)[0] final_name = FINAL_IMAGE_DIR + str(rand_uuid) + '.png' g_all = nn.Sequential( OrderedDict([ ('g_mapping', G_mapping()), # ('truncation', Truncation(avg_latent)), ('g_synthesis', G_synthesis(resolution=model_resolution)) ])) g_all.load_state_dict( torch.load( f"../image_2_style_gan/torch_weight_files/karras2019stylegan-ffhq-{model_resolution}x{model_resolution}.pt", map_location=device)) g_all.eval() g_all.to(device) g_mapping, g_synthesis = g_all[0], g_all[1] img_0 = image_reader_color(target_name) #(1,3,1024,1024) -1~1 img_0 = img_0.to(device) img_1 = image_reader_color(ingredient_name) img_1 = img_1.to(device) #(1,3,1024,1024) save_image(img_0, '../image_2_style_gan/source/trghistadjs.png') save_image(img_1, '../image_2_style_gan/source/orghistadjs.png') blur_mask0 = image_reader_color(mask_name).to(device) blur_mask0 = blur_mask0[:, 0, :, :].unsqueeze(0) blur_mask1 = blur_mask0.clone() blur_mask1 = 1 - blur_mask0 MSE_Loss = nn.MSELoss(reduction="mean") upsample2d = torch.nn.Upsample(scale_factor=0.5, mode='bilinear', align_corners=False) img_p0 = img_0.clone() #resize for perceptual net img_p0 = upsample2d(img_p0) img_p0 = upsample2d(img_p0) #(1,3,256,256) img_p1 = img_1.clone() img_p1 = upsample2d(img_p1) img_p1 = upsample2d(img_p1) #(1,3,256,256) perceptual_net = VGG16_for_Perceptual(n_layers=[2, 4, 14, 21]).to( device) #conv1_1,conv1_2,conv2_2,conv3_3 dlatent = torch.zeros((1, 18, 512), requires_grad=True, device=device) optimizer = optim.Adam({dlatent}, lr=0.01, betas=(0.9, 0.999), eps=1e-8) loss_list = [] print( "Start ---------------------------------------------------------------------------------------------" ) for i in range( ITERATION): # [img_0 : Target IMG] / [img_1 : Ingredient IMG] img_noise = random_pixel_image(min_float=0.2, max_float=0.8).to(device) optimizer.zero_grad() synth_img = g_synthesis(dlatent) synth_img = (synth_img * 0.8 + img_noise * 0.2) / 2 loss_wl0 = caluclate_loss(synth_img, img_0, perceptual_net, img_p0, blur_mask0, MSE_Loss, upsample2d) loss_wl1 = caluclate_loss(synth_img, img_1, perceptual_net, img_p1, blur_mask1, MSE_Loss, upsample2d) loss = loss_wl0 + loss_wl1 loss.backward() optimizer.step() loss_np = loss.detach().cpu().numpy() loss_0 = loss_wl0.detach().cpu().numpy() loss_1 = loss_wl1.detach().cpu().numpy() loss_list.append(loss_np) if i % 10 == 0: print("iter{}: loss --{}, loss0 --{}, loss1 --{}".format( i, loss_np, loss_0, loss_1)) elif i == (ITERATION - 1): save_image(synth_img.clamp(0, 1), final_name) origin_name = '{}{}_origin.png'.format(FINAL_IMAGE_DIR, str(rand_uuid)) os.replace(ingredient_name, origin_name) os.remove(mask_name) # 마스크 파일 삭제 print( "Complete -------------------------------------------------------------------------------------" ) return origin_name, final_name
def main(): parser = argparse.ArgumentParser( description= 'Find latent representation of reference images using perceptual loss') parser.add_argument('--batch_size', default=1, help='Batch size for generator and perceptual model', type=int) parser.add_argument('--resolution', default=1024, type=int) parser.add_argument('--src_im1', default="source_image/sample.png") parser.add_argument('--src_im2', default="source_image/0.png") parser.add_argument( '--weight_file', default="weight_files/pytorch/karras2019stylegan-ffhq-1024x1024.pt", type=str) parser.add_argument('--iteration', default=1000, type=int) args = parser.parse_args() g_all = nn.Sequential( OrderedDict([ ('g_mapping', G_mapping()), #('truncation', Truncation(avg_latent)), ('g_synthesis', G_synthesis(resolution=args.resolution)) ])) g_all.load_state_dict(torch.load(args.weight_file, map_location=device)) g_all.eval() g_all.to(device) g_mapping, g_synthesis = g_all[0], g_all[1] img_0 = image_reader(args.src_im1) #(1,3,1024,1024) -1~1 img_0 = img_0.to(device) img_1 = image_reader(args.src_im2) img_1 = img_1.to(device) #(1,3,1024,1024) MSE_Loss = nn.MSELoss(reduction="mean") upsample2d = torch.nn.Upsample(scale_factor=0.5, mode='bilinear') img_p0 = img_0.clone() #resize for perceptual net img_p0 = upsample2d(img_p0) img_p0 = upsample2d(img_p0) #(1,3,256,256) img_p1 = img_1.clone() img_p1 = upsample2d(img_p1) img_p1 = upsample2d(img_p1) #(1,3,256,256) perceptual_net = VGG16_for_Perceptual(n_layers=[2, 4, 14, 21]).to( device) #conv1_1,conv1_2,conv2_2,conv3_3 dlatent_a = torch.zeros((1, 18, 512), requires_grad=True, device=device) #appearace latent s1 dlatent_e = torch.zeros((1, 18, 512), requires_grad=True, device=device) # expression latent s2 optimizer = optim.Adam({dlatent_a, dlatent_e}, lr=0.01, betas=(0.9, 0.999), eps=1e-8) alpha = torch.zeros((1, 18, 512)).to(device) alpha[:, 3:5, :] = 1 print("Start") loss_list = [] for i in range(args.iteration): optimizer.zero_grad() synth_img_a = g_synthesis(dlatent_a) synth_img_a = (synth_img_a + 1.0) / 2.0 synth_img_e = g_synthesis(dlatent_e) synth_img_e = (synth_img_e + 1.0) / 2.0 loss_1 = caluclate_contentloss(synth_img_a, perceptual_net, img_p1, MSE_Loss, upsample2d) loss_1.backward() optimizer.step() loss_2 = caluclate_styleloss(synth_img_e, img_p0, perceptual_net, upsample2d) loss_2.backward() optimizer.step() loss_1 = loss_1.detach().cpu().numpy() loss_2 = loss_2.detach().cpu().numpy() dlatent1 = dlatent_a * alpha + dlatent_e * (1 - alpha) dlatent2 = dlatent_a * (1 - alpha) + dlatent_e * alpha synth_img1 = g_synthesis(dlatent1) synth_img1 = (synth_img1 + 1.0) / 2.0 synth_img2 = g_synthesis(dlatent2) synth_img2 = (synth_img2 + 1.0) / 2.0 if i % 10 == 0: print("iter{}: loss0 --{}, loss1 --{}".format( i, loss_1, loss_2)) save_image(synth_img_a.clamp(0, 1), "save_image/exchange/a/{}_a.png".format(i)) save_image(synth_img_e.clamp(0, 1), "save_image/exchange/e/{}_e.png".format(i)) save_image( synth_img1.clamp(0, 1), "save_image/exchange/result1/{}_exchange1.png".format(i)) save_image( synth_img2.clamp(0, 1), "save_image/exchange/result2/{}_exchange2.png".format(i)) np.save("latent_W/exchange1.npy", dlatent1.detach().cpu().numpy()) np.save("latent_W/exchange2.npy", dlatent2.detach().cpu().numpy())
def image_crossover_eyes(BASE_DIR, RAW_DIR, rand_uuid, process_selection, gender): ALIGNED_IMAGE_DIR = f'{BASE_DIR}aligned/' os.mkdir(ALIGNED_IMAGE_DIR) TARGET_IMAGE_DIR = f'{BASE_DIR}target/' os.mkdir(TARGET_IMAGE_DIR) ITERATION = 150 BATCH_SIZE = 1 # Not in usage. if process_selection == 0 and gender == 'female': TARGET_SOURCE_DIR = '../image_2_style_gan/source/target/female/' elif process_selection == 0 and gender == 'male': TARGET_SOURCE_DIR = '../image_2_style_gan/source/target/male/' else: TARGET_SOURCE_DIR = f'{BASE_DIR}raw_target/aligned/' FINAL_IMAGE_DIR = f'{BASE_DIR}final/' os.mkdir(FINAL_IMAGE_DIR) MASK_DIR = f'{BASE_DIR}mask/' os.mkdir(MASK_DIR) model_resolution = 1024 aligned_image_names = align_images(RAW_DIR, ALIGNED_IMAGE_DIR) try: if not aligned_image_names: print('\nFailed Face Alignment. Abort process.') shutil.rmtree(BASE_DIR) sys.exit(1) aligned_image_name = ALIGNED_IMAGE_DIR + \ os.listdir(ALIGNED_IMAGE_DIR)[0] mask_name, face_mask_name, eyes_mask_name, brows_mask_name, lids_mask_name = precision_facial_masks( aligned_image_name, MASK_DIR) ingredient_name = ALIGNED_IMAGE_DIR + os.listdir(ALIGNED_IMAGE_DIR)[0] random_target_image_index = random.randint( 0, len(os.listdir(TARGET_SOURCE_DIR))-1) target_name = TARGET_SOURCE_DIR + \ os.listdir(TARGET_SOURCE_DIR)[random_target_image_index] except IndexError as e: print("Missing file(s).\nCheck if all of source images prepared properly and try again.") print(f"Error of [ Alignment and Mask Creation ] part : {e}") shutil.rmtree(BASE_DIR) sys.exit(1) final_name = FINAL_IMAGE_DIR + str(rand_uuid) + '.png' g_all = nn.Sequential(OrderedDict([('g_mapping', G_mapping()), # ('truncation', Truncation(avg_latent)), ('g_synthesis', G_synthesis(resolution=model_resolution))])) g_all.load_state_dict(torch.load( f"../image_2_style_gan/torch_weight_files/karras2019stylegan-ffhq-{model_resolution}x{model_resolution}.pt", map_location=device)) g_all.eval() g_all.to(device) g_mapping, g_synthesis = g_all[0], g_all[1] img_0 = image_reader_color(target_name) # (1,3,1024,1024) -1~1 img_0 = img_0.to(device) img_1 = image_reader_color(ingredient_name) img_1 = img_1.to(device) # (1,3,1024,1024) blur_mask0_1 = image_reader_gray(mask_name).to(device) blur_mask0_2 = image_reader_gray(eyes_mask_name).to(device) blur_mask0_3 = image_reader_gray(lids_mask_name).to(device) # blur_mask0_4 = image_reader_gray(face_mask_name).to(device) # blur_mask0_5 = image_reader_gray(brows_mask_name).to(device) blur_mask1 = 1-blur_mask0_1 blur_mask_eyes = blur_mask0_1-blur_mask0_2 blur_mask_lids = blur_mask0_1-torch.clamp(blur_mask0_3-blur_mask0_2, 0, 1) img_0 = target_channel_manipulator(img_0, img_1, blur_mask0_1) img_1 = eyes_channel_matching(img_1, blur_mask0_2) save_image(img_0, '../image_2_style_gan/source/trghistadjs.png') save_image(img_1, '../image_2_style_gan/source/orghistadjs.png') save_image(blur_mask0_1, '../image_2_style_gan/source/mask.png') MSE_Loss = nn.MSELoss(reduction="mean") upsample2d = torch.nn.Upsample(scale_factor=0.5, mode='nearest') img_p0 = img_0.clone() # resize for perceptual net img_p0 = upsample2d(img_p0) img_p0 = upsample2d(img_p0) # (1,3,256,256) img_p1 = img_1.clone() img_p1 = upsample2d(img_p1) img_p1 = upsample2d(img_p1) # (1,3,256,256) # img_noise = random_pixel_image(min_float=0.3, max_float=0.7).to(device) perceptual_net = VGG16_for_Perceptual(n_layers=[2, 4, 14, 21]).to(device) # conv1_1,conv1_2,conv2_2,conv3_3 dlatent = torch.zeros((1, 18, 512), requires_grad=True, device=device) # requires_grad를 'True'로 두어 오차 역전파 과정 중 해당 Tensor에 대한 변화도를 계산하도록 한다. optimizer = optim.Adam({dlatent}, lr=0.01, betas=(0.9, 0.999), eps=1e-8) loss_list = [] print("Start ---------------------------------------------------------------------------------------------") # [img_0 : Target IMG] / [img_1 : Ingredient IMG] for i in range(ITERATION): img_noise = random_pixel_image(min_float=0.3, max_float=0.7).to(device) optimizer.zero_grad() # 매 Iter 시마다 가중치들의 변화도가 누적되지 않도록 그 변화도를 초기화시켜 준다. synth_img = g_synthesis(dlatent) synth_img = (synth_img*0.75 + img_noise*0.25) / 2 # 랜덤 노이즈인 이미지로 시작하고, 이미 특정 DATA Set으로 학습된 모델을 사용한다. # 그렇기 때문에 오차 비교 부분을 없앨 경우, 내장된 가중치를 가지고 특정한 한 인물의 이미지만을 계속 만들어낸다. # 순전파 과정을 수행한다. loss_wl0 = caluclate_loss(synth_img, img_0, perceptual_net, img_p0, blur_mask_eyes, MSE_Loss, upsample2d) # 'loss_wl0'은 변화해 갈 synth_imge와 img_0(Target Image)에 각각 'blur_mask_eyes'마스크를 적용한 Image들 간의 오차를 계산해 낸다. loss_wl1 = caluclate_loss(synth_img, img_1, perceptual_net, img_p1, blur_mask_lids, MSE_Loss, upsample2d) # 'loss_wl1'은 변화해 갈 synth_imge와 img_1(Ingredient Image)에 각각 'blur_mask_lids'마스크를 적용한 Image들 간의 오차를 계산해 낸다. loss = loss_wl0 + loss_wl1 # 최종 loss는 loss_wl0, loss_wl1 두 loss 모두를 더한 값이다. loss.backward() # 오차 역전파 과정을 진행한다. optimizer.step() # 가중치의 변화도를 .step()함수를 호출해 갱신한다. loss_np = loss.detach().cpu().numpy() loss_0 = loss_wl0.detach().cpu().numpy() loss_1 = loss_wl1.detach().cpu().numpy() # 위 순전파 과정에서 구해진 해당 Iter의 Loss값들을 연산 기록으로부터 분리한 후 별도의 변수에 따로 저장한다. loss_list.append(loss_np) if i % 10 == 0: print("iter{}: loss --{}, loss0 --{}, loss1 --{}".format(i, loss_np, loss_0, loss_1)) elif i == (ITERATION - 1): save_image(img_1*blur_mask1 + synth_img*blur_mask0_1, final_name) origin_name = '{}{}_origin.png'.format(FINAL_IMAGE_DIR, str(rand_uuid)) os.replace(ingredient_name, origin_name) print("Complete -------------------------------------------------------------------------------------") return origin_name, final_name
def image_crossover(BASE_DIR, RAW_DIR, rand_uuid, client_img_name, process_selection, gender): ALIGNED_IMAGE_DIR = f'{BASE_DIR}aligned/' os.mkdir(ALIGNED_IMAGE_DIR) TARGET_IMAGE_DIR = f'{BASE_DIR}target/' os.mkdir(TARGET_IMAGE_DIR) if process_selection == 0 and gender == 'female': TARGET_SOURCE_DIR = '../image_2_style_gan/source/target/female/' elif process_selection == 0 and gender == 'male': TARGET_SOURCE_DIR = '../image_2_style_gan/source/target/male/' else: TARGET_SOURCE_DIR = f'{BASE_DIR}raw_target/aligned/' FINAL_IMAGE_DIR = f'{BASE_DIR}final/' os.mkdir(FINAL_IMAGE_DIR) MASK_DIR = f'{BASE_DIR}mask/' os.mkdir(MASK_DIR) model_resolution = 1024 parser = argparse.ArgumentParser( description= 'Find latent representation of reference images using perceptual loss') parser.add_argument('--batch_size', default=5, help='Batch size for generator and perceptual model', type=int) parser.add_argument('--resolution', default=model_resolution, type=int) parser.add_argument('--src_im1', default=TARGET_IMAGE_DIR) parser.add_argument('--src_im2', default=ALIGNED_IMAGE_DIR) parser.add_argument('--mask', default=MASK_DIR) parser.add_argument( '--weight_file', default= "../image_2_style_gan/weight_files/pytorch/karras2019stylegan-ffhq-1024x1024.pt", type=str) parser.add_argument('--iteration', default=150, type=int) args = parser.parse_args() if os.path.isdir(os.path.dirname(args.weight_file)) is not True: os.makedirs(os.path.dirname(args.weight_file), exist_ok=True) aligned_image_names = align_images(RAW_DIR, args.src_im2) try: if not aligned_image_names: print( '\nNo raw-image detected. Process proceeds without alignment.') shutil.rmtree(BASE_DIR) # UUID 디렉터리 삭제 sys.exit(1) aligned_image_name = [args.src_im2 + os.listdir(args.src_im2)[0]] mask_maker(aligned_image_name, args.mask) ingredient_name = args.src_im2 + os.listdir(args.src_im2)[0] random_target_image_index = random.randint( 0, len(os.listdir(TARGET_SOURCE_DIR)) - 1) # random_target_image_index = 3 target_name = TARGET_SOURCE_DIR + os.listdir( TARGET_SOURCE_DIR)[random_target_image_index] print(os.listdir(TARGET_SOURCE_DIR) [random_target_image_index]) # Image file 이름 확인 except IndexError as e: print( "\nMissing file(s).\nCheck if all of source images prepared properly and try again." ) print(f"Aligned_image_names function: {e}") shutil.rmtree(BASE_DIR) # UUID 디렉터리 삭제 sys.exit(1) try: mask_name = args.mask + os.listdir(args.mask)[0] except Exception as e: shutil.copyfile('../image_2_style_gan/source/ref_mask/ref_mask.png', '{}ref_mask.png'.format(args.mask)) mask_name = args.mask + os.listdir(args.mask)[0] # file_names = [] final_name = FINAL_IMAGE_DIR + str(rand_uuid) + '.png' g_all = nn.Sequential( OrderedDict([ ('g_mapping', G_mapping()), # ('truncation', Truncation(avg_latent)), ('g_synthesis', G_synthesis(resolution=args.resolution)) ])) g_all.load_state_dict(torch.load(args.weight_file, map_location=device)) g_all.eval() g_all.to(device) g_mapping, g_synthesis = g_all[0], g_all[1] img_0 = image_reader(target_name) #(1,3,1024,1024) -1~1 img_0 = img_0.to(device) img_1 = image_reader(ingredient_name) img_1 = img_1.to(device) #(1,3,1024,1024) blur_mask0 = image_reader(mask_name).to(device) blur_mask0 = blur_mask0[:, 0, :, :].unsqueeze(0) blur_mask1 = blur_mask0.clone() blur_mask1 = 1 - blur_mask1 MSE_Loss = nn.MSELoss(reduction="mean") upsample2d = torch.nn.Upsample(scale_factor=0.5, mode='bilinear', align_corners=True) img_p0 = img_0.clone() #resize for perceptual net img_p0 = upsample2d(img_p0) img_p0 = upsample2d(img_p0) #(1,3,256,256) img_p1 = img_1.clone() img_p1 = upsample2d(img_p1) img_p1 = upsample2d(img_p1) #(1,3,256,256) perceptual_net = VGG16_for_Perceptual(n_layers=[2, 4, 14, 21]).to( device) #conv1_1,conv1_2,conv2_2,conv3_3 dlatent = torch.zeros((1, 18, 512), requires_grad=True, device=device) optimizer = optim.RMSprop({dlatent}, lr=0.01, alpha=0.99, eps=1e-08, weight_decay=0, momentum=0, centered=False) loss_list = [] print( "Start ---------------------------------------------------------------------------------------" ) for i in range(args.iteration): optimizer.zero_grad() synth_img = g_synthesis(dlatent) synth_img = (synth_img + 1.0) / 2.0 loss_wl0 = caluclate_loss(synth_img, img_0, perceptual_net, img_p0, blur_mask0, MSE_Loss, upsample2d) loss_wl1 = caluclate_loss(synth_img, img_1, perceptual_net, img_p1, blur_mask1, MSE_Loss, upsample2d) loss = loss_wl0 + loss_wl1 loss.backward() optimizer.step() loss_np = loss.detach().cpu().numpy() loss_0 = loss_wl0.detach().cpu().numpy() loss_1 = loss_wl1.detach().cpu().numpy() loss_list.append(loss_np) if i % 10 == 0: print("iter{}: loss -- {}, loss0 --{}, loss1 --{}".format( i, loss_np, loss_0, loss_1)) # file_name = "{}_{}_{}.png".format(MEDIUM_IMAGE_DIR, client_ip, i) # save_image(synth_img.clamp(0, 1), file_name) # if i > 10: # file_names.append(file_name) # np.save(r"../image_2_style_gan/latent_W/crossover_{}.npy".format(client_ip), dlatent.detach().cpu().numpy()) elif i == (args.iteration - 1): save_image(synth_img.clamp(0, 1), final_name) origin_name = '{}{}_origin.png'.format(FINAL_IMAGE_DIR, str(rand_uuid)) os.replace(ingredient_name, origin_name) os.remove(mask_name) # 마스크 파일 삭제 print( "Complete ---------------------------------------------------------------------------------------------" ) return origin_name, final_name
from collections import OrderedDict import torch import torch.nn as nn device = 'cuda:0' if torch.cuda.is_available() else 'cpu' from image_2_style_gan.stylegan_layers import G_mapping,G_synthesis,D_basic resolution=1024 g_all = nn.Sequential(OrderedDict([ ('g_mapping', G_mapping()), #('truncation', Truncation(avg_latent)), ('g_synthesis', G_synthesis(resolution=resolution)) ])) d_basic = D_basic(resolution=resolution) a=True tensorflow_dir= "weight_files/tensorflow/" pytorch_dir= "weight_files/pytorch/" weight_name="karras2019stylegan-ffhq-{}x{}".format(resolution, resolution) if a: # this can be run to get the weights, but you need the reference implementation and weights import pickle, torch, collections
def main(): parser = argparse.ArgumentParser( description= 'Find latent representation of reference images using perceptual loss') parser.add_argument('--batch_size', default=1, help='Batch size for generator and perceptual model', type=int) parser.add_argument('--resolution', default=1024, type=int) parser.add_argument('--src_im', default="sample.png") parser.add_argument('--src_dir', default="source_image/") parser.add_argument( '--weight_file', default="weight_files/pytorch/karras2019stylegan-ffhq-1024x1024.pt", type=str) parser.add_argument('--iteration', default=1000, type=int) args = parser.parse_args() g_all = nn.Sequential( OrderedDict([ ('g_mapping', G_mapping()), #('truncation', Truncation(avg_latent)), ('g_synthesis', G_synthesis(resolution=args.resolution)) ])) g_all.load_state_dict(torch.load(args.weight_file, map_location=device)) g_all.eval() g_all.to(device) g_mapping, g_synthesis = g_all[0], g_all[1] name = args.src_im.split(".")[0] img = image_reader(args.src_dir + args.src_im) #(1,3,1024,1024) -1~1 img = img.to(device) MSE_Loss = nn.MSELoss(reduction="mean") img_p = img.clone() #Perceptual loss 用画像 upsample2d = torch.nn.Upsample(scale_factor=256 / args.resolution, mode='bilinear') #VGG入力のため(256,256)にリサイズ img_p = upsample2d(img_p) perceptual_net = VGG16_for_Perceptual(n_layers=[2, 4, 14, 21]).to(device) dlatent = torch.zeros((1, 18, 512), requires_grad=True, device=device) optimizer = optim.Adam({dlatent}, lr=0.01, betas=(0.9, 0.999), eps=1e-8) print("Start") loss_list = [] for i in range(args.iteration): optimizer.zero_grad() synth_img = g_synthesis(dlatent) synth_img = (synth_img + 1.0) / 2.0 mse_loss, perceptual_loss = caluclate_loss(synth_img, img, perceptual_net, img_p, MSE_Loss, upsample2d) loss = mse_loss + perceptual_loss loss.backward() optimizer.step() loss_np = loss.detach().cpu().numpy() loss_p = perceptual_loss.detach().cpu().numpy() loss_m = mse_loss.detach().cpu().numpy() loss_list.append(loss_np) if i % 10 == 0: print( "iter{}: loss -- {}, mse_loss --{}, percep_loss --{}".format( i, loss_np, loss_m, loss_p)) save_image(synth_img.clamp(0, 1), "save_image/encode1/{}.png".format(i)) #np.save("loss_list.npy",loss_list) np.save("latent_W/{}.npy".format(name), dlatent.detach().cpu().numpy())
def image_crossover(client_ip, time_flag): # 외부로부터 접속한 Client의 IP와 Image 처리를 개시할 당시의 시간 정보를 받아온다. # 이 둘은 파일명의 중복 방지와 Client간 구분을 위해 활용될 것이다. MEDIUM_IMAGE_DIR = r'../image_2_style_gan/save_image/crossover/' # Image처리 과정 도중 생성되는 중간 산물들을 임시 저장할 경로를 설정한다. FINAL_IMAGE_DIR = r'static/images/' # 처리 완료된 Image를 파일로 저장할 경로를 설정한다. if os.path.isdir(FINAL_IMAGE_DIR) is not True: os.makedirs(FINAL_IMAGE_DIR, exist_ok=True) # 경로에 해당하는 폴더가 존재하지 않는 것이 하나라도 있을 경우, 자동으로 해당 경로 전체에 해당하는 모든 폴더를 생성한다. # 이미 존재하는 폴더에 대해서는 오류 발생을 무시하고 과정을 진행한다. if os.path.isdir(MEDIUM_IMAGE_DIR) is not True: os.makedirs(MEDIUM_IMAGE_DIR, exist_ok=True) # 경로에 해당하는 폴더가 존재하지 않는 것이 하나라도 있을 경우, 자동으로 해당 경로 전체에 해당하는 모든 폴더를 생성한다. # 이미 존재하는 폴더에 대해서는 오류 발생을 무시하고 과정을 진행한다. # Command Prompt 등을 통해 입력받는 경우를 위해 아래와 같이 'ArgumentParser'를 통해 Argument의 형태로 Parameter 값들을 입력받을 수 있다. # 하지만, 현재 사용자로부터 직접 입력받도록 돼있는 항목은 원본 Image 파일 뿐이므로, 나머지는 'default'값을 설정하는 형태로, # 지금은 사실상 메소드의 내부에서 변수들을 초기 값을 주며 선언하는 것과 별 다를 바 없는 구조이다. parser = argparse.ArgumentParser( description= 'Find latent representation of reference images using perceptual loss') parser.add_argument('--batch_size', default=6, help='Batch size for generator and perceptual model', type=int) # 입력받은 Image를 Model에 적용시킬 때 수행할 작업의 Batch Size를 설정한다. parser.add_argument('--resolution', default=1024, type=int) # 결과물이 출력될 해상도를 설정한다. parser.add_argument('--src_im1', default=r"../image_2_style_gan/source_image/target/") # 목표 Image 즉, 원본 Image에서 바뀔 부분의 기준이 되는 Image 파일의 경로를 설정한다. parser.add_argument( '--src_im2', default=r"../image_2_style_gan/source_image/ingredient/") # 변형시킬 원본이 될 Image의 경로를 설정한다. parser.add_argument('--mask', default=r"../image_2_style_gan/source_image/mask/") # 원본에서 바꿀 부분만을 지정할 수 있도록 할 Masking 파일의 경로를 설정한다. parser.add_argument( '--weight_file', default= r"../image_2_style_gan/weight_files/pytorch/karras2019stylegan-ffhq-1024x1024.pt", type=str) # Model이 이미 학습한 가중치들을 지닌 파일의 경로를 설정한다. parser.add_argument('--iteration', default=150, type=int) # 처리를 반복할 횟수를 설정한다. args = parser.parse_args( ) # 'parser' 객체에 적재한 Argument들을 Parsing해와 변수에 담아 둔다. # if client_img_name == '': # raw_image_names = os.listdir(r'../image_2_style_gan/img/') # else: # raw_image_names = client_img_name # raw_image_names = r'../image_2_style_gan/img/' aligned_image_names = align_images(args.src_im2) # 정렬이 완료된 Image를 저장할 경로를 Parameter로 넘겨주며, 메소드를 작동시키고 정렬된 파일(들)의 이름을 반환받는다. try: # Raw-Image(얼굴이 정렬되지 않은 Image)가 존재하지 않는 경우, 안내 메시지를 띄우고 그대로 진행시키는 부분이다. if not aligned_image_names: print( '\nNo raw-image detected. Process proceeds without alignment.') aligned_image_names = [args.src_im2 + os.listdir(args.src_im2)[0]] # 메소드로부터 파일명 목록을 넘겨받지 못한(=빈 목록을 넘겨받은) 대신, 이미 경로 안에 존재하는 재료를 사용하도록 한다. # 다만, 현재 API에 연동된 상태에선 무조건 Raw-Image를 넘겨받아 진행하게 돼있기 때문에, # 해당 오류가 발생하는 경우엔 아예 정상적인 진행이 불가할 가능성이 크다. mask_maker( aligned_image_names, args.mask) # 얼굴을 정렬한 Image를 획득하면, 이를 가지고 Masking Image를 생성한다. ingredient_name = args.src_im2 + os.listdir( args.src_im2)[0] # 원본 Image의 경로 + 파일명을 결합한다. target_name = args.src_im1 + os.listdir( args.src_im1)[0] # 목표 Image의 경로 + 파일명을 결합한다. except IndexError as e: # 세 가지의 필수 요소(원본 Image, Masking Image, 대상 Image) 중 하나라도 빠진 것이 있으면 메시지로 경고하고 Process를 중단시킨다. print( "\nMissing file(s).\nCheck if all of source images prepared properly and try again." ) sys.exit(1) try: mask_name = args.mask + os.listdir( args.mask)[0] # Masking Image의 경로 + 파일명을 결합한다. # 만약 여기서 오류가 발생하면 이는 해당 내용이 존재하지 않는다(= Masking Image를 만들어내지 못했다)는 의미이다. except Exception as e: # Masking Image를 제작하지 못한 경우, 기존의 범용 Masking Image를 복사해 가져와 사용하게 된다. shutil.copyfile( r'../image_2_style_gan/source_image/ref_mask/ref_mask.png', '{}ref_mask.png'.format(args.mask)) mask_name = args.mask + os.listdir(args.mask)[0] FINAL_IMAGE_DIR = FINAL_IMAGE_DIR + client_ip + time_flag + '/' # 결과물 Image를 저장할 경로를 Client의 IP와 작업 개시 시간을 이용해 조합한다. if os.path.isdir(FINAL_IMAGE_DIR) is not True: os.mkdir(FINAL_IMAGE_DIR) # 해당 경로 이름으로 실제 폴더를 생성한다. # file_names = [] final_name = FINAL_IMAGE_DIR + time_flag + '.png' # 결과물 Image의 파일명을 조합한다. # ============ 아래는 학습된 Model을 가지고 Network을 통과시켜 Image를 얻어내는 부분이다. ================ g_all = nn.Sequential( OrderedDict([ ('g_mapping', G_mapping()), #('truncation', Truncation(avg_latent)), ('g_synthesis', G_synthesis(resolution=args.resolution)) ])) g_all.load_state_dict(torch.load(args.weight_file, map_location=device)) g_all.eval() g_all.to(device) g_mapping, g_synthesis = g_all[0], g_all[1] img_0 = image_reader(target_name) #(1,3,1024,1024) -1~1 img_0 = img_0.to(device) img_1 = image_reader(ingredient_name) img_1 = img_1.to(device) #(1,3,1024,1024) blur_mask0 = image_reader(mask_name).to(device) blur_mask0 = blur_mask0[:, 0, :, :].unsqueeze(0) blur_mask1 = blur_mask0.clone() blur_mask1 = 1 - blur_mask1 MSE_Loss = nn.MSELoss(reduction="mean") upsample2d = torch.nn.Upsample(scale_factor=0.5, mode='bilinear') img_p0 = img_0.clone() #resize for perceptual net img_p0 = upsample2d(img_p0) img_p0 = upsample2d(img_p0) #(1,3,256,256) img_p1 = img_1.clone() img_p1 = upsample2d(img_p1) img_p1 = upsample2d(img_p1) #(1,3,256,256) perceptual_net = VGG16_for_Perceptual(n_layers=[2, 4, 14, 21]).to( device) #conv1_1,conv1_2,conv2_2,conv3_3 dlatent = torch.zeros((1, 18, 512), requires_grad=True, device=device) optimizer = optim.Adam({dlatent}, lr=0.01, betas=(0.9, 0.999), eps=1e-8) loss_list = [] print( "Start ---------------------------------------------------------------------------------------" ) for i in range(args.iteration): optimizer.zero_grad() synth_img = g_synthesis(dlatent) synth_img = (synth_img + 1.0) / 2.0 loss_wl0 = caluclate_loss(synth_img, img_0, perceptual_net, img_p0, blur_mask0, MSE_Loss, upsample2d) loss_wl1 = caluclate_loss(synth_img, img_1, perceptual_net, img_p1, blur_mask1, MSE_Loss, upsample2d) loss = loss_wl0 + loss_wl1 loss.backward() optimizer.step() loss_np = loss.detach().cpu().numpy() loss_0 = loss_wl0.detach().cpu().numpy() loss_1 = loss_wl1.detach().cpu().numpy() loss_list.append(loss_np) if i % 10 == 0: # Iteration 도중 적당한 단위마다 산출된 오차를 출력한다. print("iter{}: loss -- {}, loss0 --{}, loss1 --{}".format( i, loss_np, loss_0, loss_1)) # file_name = "{}_{}_{}.png".format(MEDIUM_IMAGE_DIR, client_ip, i) # save_image(synth_img.clamp(0, 1), file_name) # if i > 10: # file_names.append(file_name) # np.save(r"../image_2_style_gan/latent_W/crossover_{}.npy".format(client_ip), dlatent.detach().cpu().numpy()) elif i == (args.iteration - 1): save_image(synth_img.clamp(0, 1), final_name) # 과정이 종료됐을 때의 결과물을 지정한 파일명으로 저장한다. # gif_buffer = [] # durations = [] # gif_buffer.append(Image.open(ingredient_name)) # durations.append(3.00) # # for file in file_names: # gif_buffer.append(Image.open(file)) # durations.append(0.04) # # gif_buffer.append((Image.open(final_name))) # durations.append(3.00) # imageio.mimsave('{}{}.gif'.format(FINAL_IMAGE_DIR, time_flag), gif_buffer, duration=durations) # del gif_buffer # for file in os.listdir(r'../image_2_style_gan/save_image/crossover/'): # dir = r'../image_2_style_gan/save_image/crossover/' + file # os.remove(dir) origin_name = '{}{}_origin.png'.format(FINAL_IMAGE_DIR, time_flag) # 원본 이미지도 비교에 이용하기 위해 가져올 것이므로, 그 파일명을 설정한다. os.replace(ingredient_name, origin_name) # 원본 파일을 재명명하며 가져온다. # os.remove(ingredient_name) # os.replace(mask_name, '{}Used_{}_mask.png'.format(FINAL_IMAGE_DIR, time_flag)) os.remove( mask_name ) # Masking Image는 일회용이므로 삭제한다.(Reference Mask인 경우에는 원본이 다른 폴더에서 존재하기 때문에 그것은 계속 남아있다.) # shutil.copyfile(target_name, '{}Used_{}_target.png'.format(FINAL_IMAGE_DIR, time_flag)) print( "Complete ---------------------------------------------------------------------------------------------" ) return origin_name, final_name # 작업이 완료되면 원본 파일과 결과물 파일의 경로와 이름을 호출측에 반환한다.
def image_crossover(rand_uuid, client_img_name): MEDIUM_IMAGE_DIR = '../image_2_style_gan/images/medium/' if os.path.isdir(MEDIUM_IMAGE_DIR) is not True: os.makedirs(MEDIUM_IMAGE_DIR, exist_ok=True) FINAL_IMAGE_DIR = '../image_2_style_gan/images/final/' if os.path.isdir(FINAL_IMAGE_DIR) is not True: os.makedirs(FINAL_IMAGE_DIR, exist_ok=True) # MEDIUM_IMAGE_DIR_ELEM = MEDIUM_IMAGE_DIR.split("/") # DIR = MEDIUM_IMAGE_DIR_ELEM[0] # for ELEM in MEDIUM_IMAGE_DIR_ELEM: # if ELEM != '' and os.path.isdir(DIR) is not True: # DIR += ELEM # os.mkdir(DIR) # DIR += '/' parser = argparse.ArgumentParser(description='Find latent representation of reference images using perceptual loss') parser.add_argument('--batch_size', default=6, help='Batch size for generator and perceptual model', type=int) parser.add_argument('--resolution', default=1024, type=int) parser.add_argument('--src_im1', default="../image_2_style_gan/source_image/target/") parser.add_argument('--src_im2', default="../image_2_style_gan/images/medium/") parser.add_argument('--mask', default="../image_2_style_gan/images/mask/") parser.add_argument('--weight_file', default="../image_2_style_gan/weight_files/pytorch/karras2019stylegan-ffhq-1024x1024.pt", type=str) parser.add_argument('--iteration', default=150, type=int) args = parser.parse_args() # if client_img_name == '': # raw_image_names = os.listdir(r'../image_2_style_gan/img/') # else: # raw_image_names = client_img_name # raw_image_names = r'../image_2_style_gan/img/' aligned_image_names = align_images(args.src_im2) try: if not aligned_image_names: print('\nNo raw-image detected. Process proceeds without alignment.') aligned_image_names = [args.src_im2 + os.listdir(args.src_im2)[0]] mask_maker(aligned_image_names, args.mask) ingredient_name = args.src_im2 + os.listdir(args.src_im2)[0] target_name = args.src_im1 + os.listdir(args.src_im1)[0] except IndexError as e: print("\nMissing file(s).\nCheck if all of source images prepared properly and try again.") print(f"Aligned_image_names function: {e}") os.remove(client_img_name) sys.exit(1) try: mask_name = args.mask + os.listdir(args.mask)[0] except Exception as e: shutil.copyfile('../image_2_style_gan/source_image/ref_mask/ref_mask.png', '{}ref_mask.png'.format(args.mask)) mask_name = args.mask + os.listdir(args.mask)[0] FINAL_IMAGE_DIR = FINAL_IMAGE_DIR + str(rand_uuid) + '/' if os.path.isdir(FINAL_IMAGE_DIR) is not True: os.mkdir(FINAL_IMAGE_DIR) # file_names = [] final_name = FINAL_IMAGE_DIR + str(rand_uuid) + '.png' g_all = nn.Sequential(OrderedDict([ ('g_mapping', G_mapping()), #('truncation', Truncation(avg_latent)), ('g_synthesis', G_synthesis(resolution=args.resolution)) ])) g_all.load_state_dict(torch.load(args.weight_file, map_location=device)) g_all.eval() g_all.to(device) g_mapping,g_synthesis=g_all[0],g_all[1] img_0=image_reader(target_name) #(1,3,1024,1024) -1~1 img_0=img_0.to(device) img_1=image_reader(ingredient_name) img_1=img_1.to(device) #(1,3,1024,1024) blur_mask0=image_reader(mask_name).to(device) blur_mask0=blur_mask0[:,0,:,:].unsqueeze(0) blur_mask1=blur_mask0.clone() blur_mask1=1-blur_mask1 MSE_Loss=nn.MSELoss(reduction="mean") upsample2d=torch.nn.Upsample(scale_factor=0.5, mode='bilinear') img_p0=img_0.clone() #resize for perceptual net img_p0=upsample2d(img_p0) img_p0=upsample2d(img_p0) #(1,3,256,256) img_p1=img_1.clone() img_p1=upsample2d(img_p1) img_p1=upsample2d(img_p1) #(1,3,256,256) perceptual_net=VGG16_for_Perceptual(n_layers=[2,4,14,21]).to(device) #conv1_1,conv1_2,conv2_2,conv3_3 dlatent=torch.zeros((1,18,512),requires_grad=True,device=device) optimizer=optim.Adam({dlatent},lr=0.01,betas=(0.9,0.999),eps=1e-8) loss_list = [] print("Start ---------------------------------------------------------------------------------------") for i in range(args.iteration): optimizer.zero_grad() synth_img=g_synthesis(dlatent) synth_img = (synth_img + 1.0) / 2.0 loss_wl0=caluclate_loss(synth_img,img_0,perceptual_net,img_p0,blur_mask0,MSE_Loss,upsample2d) loss_wl1=caluclate_loss(synth_img,img_1,perceptual_net,img_p1,blur_mask1,MSE_Loss,upsample2d) loss=loss_wl0+loss_wl1 loss.backward() optimizer.step() loss_np=loss.detach().cpu().numpy() loss_0=loss_wl0.detach().cpu().numpy() loss_1=loss_wl1.detach().cpu().numpy() loss_list.append(loss_np) if i % 10 == 0: print("iter{}: loss -- {}, loss0 --{}, loss1 --{}".format(i,loss_np,loss_0,loss_1)) # file_name = "{}_{}_{}.png".format(MEDIUM_IMAGE_DIR, client_ip, i) # save_image(synth_img.clamp(0, 1), file_name) # if i > 10: # file_names.append(file_name) # np.save(r"../image_2_style_gan/latent_W/crossover_{}.npy".format(client_ip), dlatent.detach().cpu().numpy()) elif i == (args.iteration - 1): save_image(synth_img.clamp(0, 1), final_name) # gif_buffer = [] # durations = [] # gif_buffer.append(Image.open(ingredient_name)) # durations.append(3.00) # # for file in file_names: # gif_buffer.append(Image.open(file)) # durations.append(0.04) # # gif_buffer.append((Image.open(final_name))) # durations.append(3.00) # imageio.mimsave('{}{}.gif'.format(FINAL_IMAGE_DIR, time_flag), gif_buffer, duration=durations) # del gif_buffer # for file in os.listdir(r'../image_2_style_gan/save_image/crossover/'): # dir = r'../image_2_style_gan/save_image/crossover/' + file # os.remove(dir) origin_name = '{}{}_origin.png'.format(FINAL_IMAGE_DIR, str(rand_uuid)) os.replace(ingredient_name, origin_name) # os.remove(ingredient_name) # os.replace(mask_name, '{}Used_{}_mask.png'.format(FINAL_IMAGE_DIR, time_flag)) os.remove(mask_name) # shutil.copyfile(target_name, '{}Used_{}_target.png'.format(FINAL_IMAGE_DIR, time_flag)) print("Complete ---------------------------------------------------------------------------------------------") return origin_name, final_name