def draw_concat(Gs,Zs,reals,NoiseAmp,in_s,mode,m_noise,m_image,opt): G_z = in_s if len(Gs) > 0: if mode == 'rand': count = 0 pad_noise = int(((opt.ker_size-1)*opt.num_layer)/2) if opt.mode == 'animation_train': pad_noise = 0 for G,Z_opt,real_curr,real_next,noise_amp in zip(Gs,Zs,reals,reals[1:],NoiseAmp): if count == 0: z = functions.generate_noise([1, Z_opt.shape[2] - 2 * pad_noise, Z_opt.shape[3] - 2 * pad_noise], device=opt.device) z = z.expand(1, 3, z.shape[2], z.shape[3]) else: z = functions.generate_noise([opt.nc_z,Z_opt.shape[2] - 2 * pad_noise, Z_opt.shape[3] - 2 * pad_noise], device=opt.device) z = m_noise(z) G_z = G_z[:,:,0:real_curr.shape[2],0:real_curr.shape[3]] G_z = m_image(G_z) z_in = noise_amp*z+G_z G_z = G(z_in.detach(),G_z) G_z = imresize(G_z,1/opt.scale_factor,opt) G_z = G_z[:,:,0:real_next.shape[2],0:real_next.shape[3]] count += 1 if mode == 'rec': count = 0 for G,Z_opt,real_curr,real_next,noise_amp in zip(Gs,Zs,reals,reals[1:],NoiseAmp): G_z = G_z[:, :, 0:real_curr.shape[2], 0:real_curr.shape[3]] G_z = m_image(G_z) z_in = noise_amp*Z_opt+G_z G_z = G(z_in.detach(),G_z) G_z = imresize(G_z,1/opt.scale_factor,opt) G_z = G_z[:,:,0:real_next.shape[2],0:real_next.shape[3]] #if count != (len(Gs)-1): # G_z = m_image(G_z) count += 1 return G_z
def draw_concat(Gs, Zs, reals, NoiseAmp, in_s, mode, m_noise, m_image, opt, noise_mode: NoiseMode): G_z = in_s if len(Gs) > 0: if mode == 'rand': count = 0 pad_noise = int(((opt.ker_size-1)*opt.num_layer)/2) for G,(Z_opt1, Z_opt2),real_curr,real_next,(noise_amp1, noise_amp2) in zip(Gs,Zs,reals,reals[1:],NoiseAmp): noise_amp = None if noise_mode == NoiseMode.Z1: z1 = _create_noise_for_draw_concat(opt, count, pad_noise, m_noise, Z_opt1, noise_mode) z2 = torch.zeros(z1.shape, device=opt.device) noise_amp = noise_amp1 elif noise_mode == NoiseMode.Z2: z2 = _create_noise_for_draw_concat(opt, count, pad_noise, m_noise, Z_opt2, noise_mode) z1 = torch.zeros(z2.shape, device=opt.device) noise_amp = noise_amp2 elif noise_mode == NoiseMode.MIXED: z1 = _create_noise_for_draw_concat(opt, count, pad_noise, m_noise, Z_opt1, noise_mode) z2 = _create_noise_for_draw_concat(opt, count, pad_noise, m_noise, Z_opt2, noise_mode) else: raise NotImplementedError z = functions.merge_noise_vectors(z1, z2, opt.noise_vectors_merge_method) G_z = G_z[:,:,0:real_curr.shape[2],0:real_curr.shape[3]] G_z = m_image(G_z) z_in = noise_amp*z+G_z G_z = G(z_in.detach(), G_z)[0] G_z = imresize(G_z,1/opt.scale_factor,opt) G_z = G_z[:,:,0:real_next.shape[2],0:real_next.shape[3]] count += 1 if mode == 'rec': count = 0 for G,(Z_opt1, Z_opt2),real_curr,real_next,(noise_amp1, noise_amp2) in zip(Gs,Zs,reals,reals[1:],NoiseAmp): G_z = G_z[:, :, 0:real_curr.shape[2], 0:real_curr.shape[3]] G_z = m_image(G_z) noise_amp = None if noise_mode == NoiseMode.Z1: Z_opt2_zeros = torch.zeros(Z_opt2.shape, device=opt.device) Z_opt = functions.merge_noise_vectors(Z_opt1, Z_opt2_zeros, opt.noise_vectors_merge_method) noise_amp = noise_amp1 elif noise_mode == NoiseMode.Z2: Z_opt1_zeros = torch.zeros(Z_opt1.shape, device=opt.device) Z_opt = functions.merge_noise_vectors(Z_opt1_zeros, Z_opt2, opt.noise_vectors_merge_method) noise_amp = noise_amp2 elif noise_mode == NoiseMode.MIXED: Z_opt = functions.merge_noise_vectors(Z_opt1, Z_opt2, opt.noise_vectors_merge_method) else: raise NotImplementedError z_in = noise_amp*Z_opt+G_z G_z = G(z_in.detach(),G_z)[0] G_z = imresize(G_z,1/opt.scale_factor,opt) G_z = G_z[:,:,0:real_next.shape[2],0:real_next.shape[3]] #if count != (len(Gs)-1): # G_z = m_image(G_z) count += 1 return G_z
def draw_concat(Gs,Zs,reals,NoiseAmp,in_s,mode,m_noise,m_image,opt): G_z = in_s if len(Gs) > 0: if mode == 'rand': count = 0 pad_noise = int(((opt.ker_size-1)*opt.num_layer)/2) if opt.mode == 'animation_train': pad_noise = 0 for G,Z_opt,real_curr,real_next,noise_amp in zip(Gs,Zs,reals,reals[1:],NoiseAmp): if count == 0: z = functions.generate_noise([1, Z_opt.shape[1] - 2 * pad_noise, Z_opt.shape[2] - 2 * pad_noise]) z = tf.broadcast_to(z, [1, z.shape[1], z.shape[2], 3]) else: z = functions.generate_noise([opt.nc_z,Z_opt.shape[1] - 2 * pad_noise, Z_opt.shape[2] - 2 * pad_noise]) z = m_noise(z) G_z = G_z[:,0:real_curr.shape[1],0:real_curr.shape[2],:] #PY: NCWH, TF:NWHC G_z = m_image(G_z) z_in = noise_amp*z+G_z G_z = G(z_in,G_z, training=True) G_z = imresize(G_z,1/opt.scale_factor,opt) G_z = G_z[:,0:real_next.shape[1],0:real_next.shape[2], :] count += 1 if mode == 'rec': count = 0 for G,Z_opt,real_curr,real_next,noise_amp in zip(Gs,Zs,reals,reals[1:],NoiseAmp): G_z = G_z[:, 0:real_curr.shape[1], 0:real_curr.shape[2], :] G_z = m_image(G_z) z_in = noise_amp*Z_opt+G_z G_z = G(z_in,G_z, training=True) G_z = imresize(G_z,1/opt.scale_factor,opt) G_z = G_z[:,0:real_next.shape[1],0:real_next.shape[2], :] count += 1 return G_z
def draw_concat(Gs, Zs, reals, NoiseAmp, in_s, mode, m_noise, m_image, opt): G_z = in_s # if it's not the first scale, else do nothign if len(Gs) > 0: # if in random mode if mode == 'rand': count = 0 pad_noise = int(((opt.ker_size - 1) * opt.num_layer) / 2) #from each scale for G, Z_opt, real_curr, real_next, noise_amp in zip( Gs, Zs, reals, reals[1:], NoiseAmp): # for the first loop if count == 0: #generate the noise z = functions.generate_noise([ 1, Z_opt.shape[2] - 2 * pad_noise, Z_opt.shape[3] - 2 * pad_noise ], device=opt.device) #broadcast it to correct shape z = z.expand(1, 3, z.shape[2], z.shape[3]) else: #direct generate the noise z = functions.generate_noise([ opt.nc_z, Z_opt.shape[2] - 2 * pad_noise, Z_opt.shape[3] - 2 * pad_noise ], device=opt.device) #padding the noise z = m_noise(z) #------------------------------------------------------------ #generate a shape of current real image's [width,height] from G_z(in_s) G_z = G_z[:, :, 0:real_curr.shape[2], 0:real_curr.shape[3]] #padding it with images G_z = m_image(G_z) #amplify the generated noise, then add with the G_z z_in = noise_amp * z + G_z #generate a new output from generator G_z = G(z_in.detach(), G_z) #resize the graph with 1/opt.scale_factor G_z = imresize(G_z, 1 / opt.scale_factor, opt) #generate a shape of current real image's [width,height] from G_z(in_s) G_z = G_z[:, :, 0:real_next.shape[2], 0:real_next.shape[3]] count += 1 if mode == 'rec': count = 0 #from each scale for G, Z_opt, real_curr, real_next, noise_amp in zip( Gs, Zs, reals, reals[1:], NoiseAmp): # do same thing except G_z = G_z[:, :, 0:real_curr.shape[2], 0:real_curr.shape[3]] G_z = m_image(G_z) z_in = noise_amp * Z_opt + G_z # for here we use Z_opt instead of generated noise G_z = G(z_in.detach(), G_z) G_z = imresize(G_z, 1 / opt.scale_factor, opt) G_z = G_z[:, :, 0:real_next.shape[2], 0:real_next.shape[3]] count += 1 return G_z
def draw_concat(Gs, Zs, reals, NoiseAmp, in_s, mode, m_noise, m_image, opt): """ Generate through all higher level Gs """ G_z = in_s # G_z is the current image output if len( Gs ) > 0: # skipped for the initial pyr level, since there is no previous G to generate if mode == 'rand': # using random noise map Z_opt count = 0 pad_noise = int(((opt.ker_size - 1) * opt.num_layer) / 2) if opt.mode == 'animation_train': pad_noise = 0 for G, Z_opt, real_curr, real_next, noise_amp in zip( Gs, Zs, reals, reals[1:], NoiseAmp): if count == 0: # Z_opt is not really used, except its size. z = functions.generate_noise([ 1, Z_opt.shape[2] - 2 * pad_noise, Z_opt.shape[3] - 2 * pad_noise ], device=opt.device) z = z.expand(1, 3, z.shape[2], z.shape[3]) # same value along color channel else: z = functions.generate_noise( [ opt.nc_z, Z_opt.shape[2] - 2 * pad_noise, Z_opt.shape[3] - 2 * pad_noise ], device=opt.device) # noise including color z = m_noise(z) G_z = G_z[:, :, 0:real_curr.shape[2], 0:real_curr.shape[3]] G_z = m_image(G_z) z_in = noise_amp * z + G_z G_z = G(z_in.detach(), G_z) G_z = imresize(G_z, 1 / opt.scale_factor, opt) # upsample it to current level G_z = G_z[:, :, 0:real_next.shape[2], 0:real_next.shape[3]] count += 1 if mode == 'rec': # using reconstruction vectors Z_opt count = 0 for G, Z_opt, real_curr, real_next, noise_amp in zip( Gs, Zs, reals, reals[1:], NoiseAmp): G_z = G_z[:, :, 0:real_curr.shape[2], 0:real_curr.shape[ 3]] # make sure the size is the same as real pyr G_z = m_image(G_z) z_in = noise_amp * Z_opt + G_z # use the loaded noise amplitude G_z = G(z_in.detach(), G_z) # THis is the iteration equation for G_z G_z = imresize(G_z, 1 / opt.scale_factor, opt) # upsample it to current level G_z = G_z[:, :, 0:real_next.shape[2], 0:real_next.shape[ 3]] # make sure the size is the same as real pyr #if count != (len(Gs)-1): # G_z = m_image(G_z) count += 1 return G_z
def train(opt,Gs,Zs,reals1,reals2,NoiseAmp): real1_, real2_ = functions.read_image(opt) in_s = 0 scale_num = 0 real1 = imresize(real1_,opt.scale1,opt) real2 = imresize(real2_, opt.scale1, opt) reals1 = functions.creat_reals_pyramid(real1,reals1,opt) reals2 = functions.creat_reals_pyramid(real2, reals2, opt) nfc_prev = 0 while scale_num<opt.stop_scale+1: opt.nfc = min(opt.nfc_init * pow(2, math.floor(scale_num / 4)), 128) opt.min_nfc = min(opt.min_nfc_init * pow(2, math.floor(scale_num / 4)), 128) opt.out_ = functions.generate_dir2save(opt) opt.outf = '%s/%d' % (opt.out_,scale_num) try: os.makedirs(opt.outf) except OSError: pass plt.imsave('%s/real1_scale.png' % (opt.outf), functions.convert_image_np(reals1[scale_num]), vmin=0, vmax=1) plt.imsave('%s/real2_scale.png' % (opt.outf), functions.convert_image_np(reals2[scale_num]), vmin=0, vmax=1) D_curr,G_curr = init_models(opt) if (nfc_prev==opt.nfc): G_curr.load_state_dict(torch.load('%s/%d/netG.pth' % (opt.out_,scale_num-1))) D_curr.load_state_dict(torch.load('%s/%d/netD.pth' % (opt.out_,scale_num-1))) z_curr,in_s,G_curr = train_single_scale(D_curr,G_curr,reals1,reals2,Gs,Zs,in_s,NoiseAmp,opt) G_curr = functions.reset_grads(G_curr,False) G_curr.eval() D_curr = functions.reset_grads(D_curr,False) D_curr.eval() Gs.append(G_curr) Zs.append(z_curr) NoiseAmp.append(opt.noise_amp) torch.save(Zs, '%s/Zs.pth' % (opt.out_)) torch.save(Gs, '%s/Gs.pth' % (opt.out_)) torch.save(reals1, '%s/reals1.pth' % (opt.out_)) torch.save(reals2, '%s/reals2.pth' % (opt.out_)) torch.save(NoiseAmp, '%s/NoiseAmp.pth' % (opt.out_)) scale_num+=1 nfc_prev = opt.nfc del D_curr,G_curr return
def SinGAN_generate(Gs,Zs,reals,NoiseAmp,opt,in_s=None,scale_v=1,scale_h=1,n=0,gen_start_scale=0,num_samples=50): #if torch.is_tensor(in_s) == False: if in_s is None: in_s = torch.full(reals[0].shape, 0, device=opt.device) images_cur = [] for G,Z_opt,noise_amp in zip(Gs,Zs,NoiseAmp): pad1 = ((opt.ker_size-1)*opt.num_layer)/2 m = nn.ZeroPad2d(int(pad1)) nzx = (Z_opt.shape[2]-pad1*2)*scale_v nzy = (Z_opt.shape[3]-pad1*2)*scale_h images_prev = images_cur images_cur = [] for i in range(0,num_samples,1): if n == 0: z_curr = functions.generate_noise([1,nzx,nzy], device=opt.device) z_curr = z_curr.expand(1,3,z_curr.shape[2],z_curr.shape[3]) z_curr = m(z_curr) else: z_curr = functions.generate_noise([opt.nc_z,nzx,nzy], device=opt.device) z_curr = m(z_curr) if images_prev == []: I_prev = m(in_s) #I_prev = m(I_prev) #I_prev = I_prev[:,:,0:z_curr.shape[2],0:z_curr.shape[3]] #I_prev = functions.upsampling(I_prev,z_curr.shape[2],z_curr.shape[3]) else: I_prev = images_prev[i] I_prev = imresize(I_prev,1/opt.scale_factor, opt) if opt.mode != "SR": I_prev = I_prev[:, :, 0:round(scale_v * reals[n].shape[2]), 0:round(scale_h * reals[n].shape[3])] I_prev = m(I_prev) I_prev = I_prev[:,:,0:z_curr.shape[2],0:z_curr.shape[3]] I_prev = functions.upsampling(I_prev,z_curr.shape[2],z_curr.shape[3]) else: I_prev = m(I_prev) if n < gen_start_scale: z_curr = Z_opt z_in = noise_amp*(z_curr)+I_prev I_curr = G(z_in.detach(),I_prev) if opt.mode == 'train': dir2save = '%s/RandomSamples/%s/gen_start_scale=%d' % (opt.out, opt.input_name[:-4], gen_start_scale) else: dir2save = functions.generate_dir2save(opt) try: os.makedirs(dir2save) except OSError: pass if (opt.mode != "harmonization") & (opt.mode != "editing") & (opt.mode != "SR") & (opt.mode != "paint2image"): plt.imsave('%s/%d.png' % (dir2save, i), functions.convert_image_np(I_curr.detach()), vmin=0,vmax=1) #plt.imsave('%s/%d_%d.png' % (dir2save,i,n),functions.convert_image_np(I_curr.detach()), vmin=0, vmax=1) #plt.imsave('%s/in_s.png' % (dir2save), functions.convert_image_np(in_s), vmin=0,vmax=1) images_cur.append(I_curr) n+=1 return I_curr.detach()
def adjust_scales2image(real_, opt): # print(real_.shape) # [1, 3, 300, 300] # print(real_.shape[2]) # 300 # print(real_.shape[3]) # 300 opt.num_scales = int((math.log(math.pow(opt.min_size / (real_.shape[2]), 1), opt.scale_factor_init))) + 1 # print(opt.num_scales) # scale: 9 scale2stop = int( math.log(min([opt.max_size, max([real_.shape[2], real_.shape[3]])]) / max([real_.shape[2], real_.shape[3]]), opt.scale_factor_init)) # print('scale2stop', scale2stop) # scale2stop: 0 opt.stop_scale = opt.num_scales - scale2stop # print('stop_scale', opt.stop_scale) # out: 9 opt.scale1 = min(opt.max_size / max([real_.shape[2], real_.shape[3]]), 1) # min(250/max([real_.shape[0],real_.shape[1]]),1) # print('scale1', opt.scale1) # out: 1.0 real = imresize(real_, opt.scale1, opt) # print('imresize_real', real.shape) # [1, 3, 300, 300] opt.scale_factor = math.pow(opt.min_size / (real.shape[2]), 1 / (opt.stop_scale)) # print(opt.scale_factor) # 0.7587 # opt.scale_factor = math.pow(opt.min_size/(min(real_.shape[0],real_.shape[1])),1/(opt.stop_scale)) scale2stop = int( math.log(min([opt.max_size, max([real_.shape[2], real_.shape[3]])]) / max([real_.shape[2], real_.shape[3]]), opt.scale_factor_init)) # print(scale2stop) # 0 opt.stop_scale = opt.num_scales - scale2stop # print(opt.stop_scale) # 9 return real
def preprocess_content_image(opt, reals,scale): real = functions.read_image(opt) functions.adjust_scales2image(real, opt) ref = functions.read_image_dir('%s/%s' % (opt.ref_dir, opt.ref_name), opt) if ref.shape[3] != real.shape[3]: ref = imresize_to_shape(ref, [real.shape[2], real.shape[3]], opt) ref = ref[:, :, :real.shape[2], :real.shape[3]] N = len(reals) - 1 n = scale in_s = imresize(ref, pow(opt.scale_factor, (N - n + 1)), opt) in_s = in_s[:, :, :reals[n - 1].shape[2], :reals[n - 1].shape[3]] in_s = imresize(in_s, 1 / opt.scale_factor, opt) in_s = in_s[:, :, :reals[n].shape[2], :reals[n].shape[3]] return in_s
def adjust_scales2image(real_, opt): # opt.num_scales = int((math.log(math.pow(opt.min_size / (real_.shape[2]), 1), opt.scale_factor_init))) + 1 opt.num_scales = math.ceil((math.log( math.pow(opt.min_size / (min(real_.shape[2], real_.shape[3])), 1), opt.scale_factor_init))) + 1 scale2stop = math.ceil( math.log( min([opt.max_size, max([real_.shape[2], real_.shape[3]])]) / max([real_.shape[2], real_.shape[3]]), opt.scale_factor_init)) opt.stop_scale = opt.num_scales - scale2stop opt.scale1 = min(opt.max_size / max([real_.shape[2], real_.shape[3]]), 1) # min(250/max([real_.shape[0],real_.shape[1]]),1) real = imresize(real_, opt.scale1, opt) # opt.scale_factor = math.pow(opt.min_size / (real.shape[2]), 1 / (opt.stop_scale)) opt.scale_factor = math.pow( opt.min_size / (min(real.shape[2], real.shape[3])), 1 / (opt.stop_scale)) scale2stop = math.ceil( math.log( min([opt.max_size, max([real_.shape[2], real_.shape[3]])]) / max([real_.shape[2], real_.shape[3]]), opt.scale_factor_init)) opt.stop_scale = opt.num_scales - scale2stop return real
def adjust_scales2image(real_, opt): #opt.num_scales = int((math.log(math.pow(opt.min_size / (real_.shape[2]), 1), opt.scale_factor_init))) + 1 opt.num_scales = math.ceil((math.log( math.pow(opt.min_size / (min(real_.shape[2], real_.shape[3])), 1), opt.scale_factor_init))) + 1 scale2stop = math.ceil( math.log( min([opt.max_size, max([real_.shape[2], real_.shape[3]])]) / max([real_.shape[2], real_.shape[3]]), opt.scale_factor_init)) opt.stop_scale = opt.num_scales - scale2stop opt.scale1 = min(opt.max_size / max([real_.shape[2], real_.shape[3]]), 1) # min(250/max([real_.shape[0],real_.shape[1]]),1) #taking spare pixels for mask in case needed if opt.mode != 'random_samples': opt.mask_coords[::2] = np.floor( np.array(opt.mask_coords[::2]).astype(np.float) * opt.scale1) opt.mask_coords[1::2] = np.ceil( np.array(opt.mask_coords[1::2]).astype(np.float) * opt.scale1) real = imresize(real_, opt.scale1, opt) #opt.scale_factor = math.pow(opt.min_size / (real.shape[2]), 1 / (opt.stop_scale)) opt.scale_factor = math.pow( opt.min_size / (min(real.shape[2], real.shape[3])), 1 / (opt.stop_scale)) scale2stop = math.ceil( math.log( min([opt.max_size, max([real_.shape[2], real_.shape[3]])]) / max([real_.shape[2], real_.shape[3]]), opt.scale_factor_init)) opt.stop_scale = opt.num_scales - scale2stop return real
def adjust_scales2image_SR(real_, opt): opt.min_size = 18 opt.num_scales = (int((math.log( opt.min_size / min(real_.shape[2], real_.shape[3]), opt.scale_factor_init, ))) + 1) scale2stop = int( math.log( min(opt.max_size, max(real_.shape[2], real_.shape[3])) / max(real_.shape[0], real_.shape[3]), opt.scale_factor_init, )) opt.stop_scale = opt.num_scales - scale2stop opt.scale1 = min(opt.max_size / max([real_.shape[2], real_.shape[3]]), 1) # min(250/max([real_.shape[0],real_.shape[1]]),1) real = imresize(real_, opt.scale1, opt) # opt.scale_factor = math.pow(opt.min_size / (real.shape[2]), 1 / (opt.stop_scale)) opt.scale_factor = math.pow( opt.min_size / (min(real.shape[2], real.shape[3])), 1 / (opt.stop_scale)) scale2stop = int( math.log( min(opt.max_size, max(real_.shape[2], real_.shape[3])) / max(real_.shape[0], real_.shape[3]), opt.scale_factor_init, )) opt.stop_scale = opt.num_scales - scale2stop return real
def creat_reals_pyramid(real, reals, opt): real = real[:, 0:3, :, :] for i in range(0, opt.stop_scale + 1, 1): scale = math.pow(opt.scale_factor, opt.stop_scale - i) curr_real = imresize(real, scale, opt) reals.append(curr_real) return reals
def create_reals_pyramid(real, reals, opt): real = real[:, 0:3, :, :] for i in range(0, opt.stop_scale + 1, 1): scale = math.pow(opt.scale_factor, opt.stop_scale - i) #resize the real image to correct scale curr_real = imresize(real, scale, opt) reals.append(curr_real) return reals # the image list
def creat_reals_pyramid(real, reals, opt): if opt.input_type == 'image': real = real[:, 0:3, :, :] for i in range(0, opt.stop_scale + 1, 1): scale = math.pow(opt.scale_factor, opt.stop_scale - i) curr_real = imresize(real, scale, opt) # print('@ creat_reals_pyramid: curr_real.shape = ', curr_real.shape) reals.append(curr_real) return reals
def create_pyramid(im,pyr_list,opt, mode=None): for i in range(0,opt.stop_scale+1,1): scale = math.pow(opt.scale_factor,opt.stop_scale-i) if mode == "mask": curr_im = imresize_mask(im,scale,opt) else: curr_im = imresize(im,scale,opt) pyr_list.append(curr_im.to(opt.device)) return pyr_list
def creat_reals_pyramid(real, reals, opt): if real.shape[1] == 4: real = real[:, 0:4, :, :] # added by vajira else: real = real[:, 0:3, :, :] for i in range(0, opt.stop_scale + 1, 1): scale = math.pow(opt.scale_factor, opt.stop_scale - i) curr_real = imresize(real, scale, opt) reals.append(curr_real) return reals
def train(opt, Gs, Zs, reals, NoiseAmp): real_ = functions.read_image(opt) in_s = 0 scale_num = 0 real = imresize(real_, opt.scale1, opt) reals = functions.creat_reals_pyramid(real, reals, opt) nfc_prev = 0 netD_optimizer = tf.keras.optimizers.Adam(learning_rate=opt.lr_d, beta_1=opt.beta1, beta_2=0.999) netG_optimizer = tf.keras.optimizers.Adam(learning_rate=opt.lr_g, beta_1=opt.beta1, beta_2=0.999) while scale_num < opt.stop_scale + 1: opt.nfc = min(opt.nfc_init * pow(2, math.floor(scale_num / 4)), 128) opt.min_nfc = min(opt.min_nfc_init * pow(2, math.floor(scale_num / 4)), 128) opt.out_ = functions.generate_dir2save(opt) opt.outf = '%s/%d' % (opt.out_, scale_num) try: os.makedirs(opt.outf) except OSError: pass plt.imsave('%s/real_scale.png' % (opt.outf), functions.convert_image_np(reals[scale_num]), vmin=0, vmax=1) D_curr, G_curr = init_models(opt) if nfc_prev == opt.nfc: D_curr.load_weights('%s/%d/netD' % (opt.out_, scale_num - 1)) G_curr.load_weights('%s/%d/netG' % (opt.out_, scale_num - 1)) z_curr, in_s, G_curr = train_single_scale(D_curr, G_curr, reals, Gs, Zs, in_s, NoiseAmp, opt, scale_num, netG_optimizer, netD_optimizer) Gs.append(G_curr) Zs.append(z_curr) NoiseAmp.append(opt.noise_amp) with open('%s/Zs.pkl' % (opt.out_), 'wb') as f: pickle.dump(Zs, f) with open('%s/reals.pkl' % (opt.out_), 'wb') as f: pickle.dump(reals, f) with open('%s/NoiseAmp.pkl' % (opt.out_), 'wb') as f: pickle.dump(NoiseAmp, f) scale_num += 1 nfc_prev = opt.nfc del D_curr, G_curr return None
def cache_input_output(Gs, Zs, NoiseAmp, reals, scale_in=None, scale_out=None): """ cache time-series input at scale i, and time-series output at scale j. both i and j start at 0 index :output: cache_dict -- dict, { 'input': list of time-series np.array, 'output': list of time-series np.array } """ cache_dict = defaultdict(list) # by default cache first scale input and final scale output scale_in = 0 if scale_in is None else scale_in scale_out = len(Gs) - 1 if scale_out is None else scale_out # create layer for boarder padding pad_image = int(((ker_size - 1) * num_layer) / 2) m_image = nn.ZeroPad2d(int(pad_image)) in_s = torch.full(Zs[0].shape, 0, device=device) frames_curr = [] # out loop is scale iteration for scale_n, (G, Z_opt, noise_amp, real) in enumerate(zip(Gs, Zs, NoiseAmp, reals)): frames_prev = frames_curr frames_curr = [] z_prev1, z_prev2 = compute_z_prev(scale_n, Z_opt, device) # inner loop is time iteration for t in range(0, 100, 1): z_diff = compute_z_diff(scale_n, Z_opt, z_prev1, z_prev2, beta, device) z_curr = compute_z_curr(Z_opt, z_prev1, z_diff, alpha) z_prev2 = z_prev1 z_prev1 = z_curr # overwrite z_curr if init at higher scale if scale_n < scale_start: z_curr = Z_opt if frames_prev == []: I_prev = in_s else: I_prev = frames_prev[t] I_prev = imresize(I_prev, 1 / scale_factor, opt) # edit I_prev = I_prev[:, :, 0:real.shape[2], 0:real.shape[3]] I_prev = m_image(I_prev) z_in = noise_amp * z_curr + I_prev I_curr = G(z_in.detach(), I_prev) frames_curr.append(I_curr) # cache results if scale_n == scale_in: z_in = tensor_to_np(z_in) cache_dict['input'].append(z_in) if scale_n == scale_out: I_curr = tensor_to_np(I_curr) cache_dict['output'].append(I_curr) return cache_dict
def create_masks_pyramid(real, masks, opt): real = real[:, 0:3, :, :] for i in range(0, opt.stop_scale + 1, 1): scale = math.pow(opt.scale_factor, opt.stop_scale - i) curr_real = imresize(real, scale, opt) curr_mask = torch.ones_like(curr_real) curr_coords = [ int(np.ceil(coord * scale)) for coord in opt.mask_coords ] curr_mask[:, :, curr_coords[0]:curr_coords[1], curr_coords[2]:curr_coords[3]] = 0 masks.append(curr_mask) return masks
def SinGAN_SR(opt, Gs, Zs, reals, NoiseAmp): mode = opt.mode in_scale, iter_num = functions.calc_init_scale(opt) opt.scale_factor = 1 / in_scale opt.scale_factor_init = 1 / in_scale opt.mode = 'SR_train' #opt.alpha = 100 opt.stop_scale = 0 dir2trained_model = functions.generate_dir2save(opt) if (os.path.exists(dir2trained_model)): #print('Trained model does not exist, training SinGAN for SR') Gs, Zs, reals, NoiseAmp = functions.load_trained_pyramid(opt) opt.mode = mode else: SR_train(opt, Gs, Zs, reals, NoiseAmp) opt.mode = mode print('%f' % pow(in_scale, iter_num)) Zs_sr = [] reals_sr = [] NoiseAmp_sr = [] Gs_sr = [] real = reals[-1] #read_image(opt) for j in range(1, iter_num + 1, 1): real_ = imresize(real, pow(1 / opt.scale_factor, j), opt) real_ = real_[:, :, 0:int(pow(1 / opt.scale_factor, j) * real.shape[2]), 0:int(pow(1 / opt.scale_factor, j) * real.shape[3])] reals_sr.append(real_) Gs_sr.append(Gs[-1]) NoiseAmp_sr.append(NoiseAmp[-1]) z_opt = torch.full(real_.shape, 0, device=opt.device) m = nn.ZeroPad2d(5) z_opt = m(z_opt) Zs_sr.append(z_opt) out = SinGAN_generate(Gs_sr, Zs_sr, reals_sr, NoiseAmp_sr, opt, in_s=reals_sr[0], num_samples=1) dir2save = functions.generate_dir2save(opt) plt.imsave('%s.png' % (dir2save), functions.convert_image_np(out.detach()), vmin=0, vmax=1) return
def adjust_scales2image(real_,opt): #opt.num_scales = int((math.log(math.pow(opt.min_size / (real_.shape[2]), 1), opt.scale_factor_init))) + 1 # num_scales: how many levels of pyramids opt.num_scales = int((math.log(math.pow(opt.min_size / (min(real_.shape[2], real_.shape[3])), 1), opt.scale_factor_init))) # scale2stop: for the largest patch size, what ratio wrt the image shape in terms of scaler_factor_init # 1:0; 1/2:1; etc scale2stop = math.ceil(math.log(min([opt.max_size, max([real_.shape[2], real_.shape[3]])]) / max([real_.shape[2], real_.shape[3]]),opt.scale_factor_init)) # stop_scale: level to stop, since reach the maximum size. opt.stop_scale = opt.num_scales - scale2stop # scale1: for the largest patch size, what ratio wrt the image shape opt.scale1 = min(opt.max_size / max([real_.shape[2], real_.shape[3]]),1) # min(250/max([real_.shape[0],real_.shape[1]]),1) real = imresize(real_, opt.scale1, opt) # scale_factor: evenly divide the scale_factor_init opt.scale_factor = opt.scale_factor_init # scale2stop = math.ceil(math.log(min([opt.max_size, max([real_.shape[2], real_.shape[3]])]) / max([real_.shape[2], real_.shape[3]]),opt.scale_factor_init)) # opt.stop_scale = opt.num_scales - scale2stop return real
def adjust_scales2image_SR(real_, opt): # 定义最小尺寸为18 opt.min_size = 18 # scales数目 opt.num_scales = int( (math.log(math.pow(opt.min_size / (min([real_.shape[2], real_.shape[3]])), 1), opt.scale_factor_init))) + 1 print('aaaaaaaaaaaa', opt.num_scales) scale2stop = int( math.log(min([opt.max_size, max([real_.shape[2], real_.shape[3]])]) / max([real_.shape[2], real_.shape[3]]), opt.scale_factor_init)) opt.stop_scale = opt.num_scales - scale2stop opt.scale1 = min(opt.max_size / max([real_.shape[2], real_.shape[3]]), 1) # min(250/max([real_.shape[0],real_.shape[1]]),1) real = imresize(real_, opt.scale1, opt) opt.scale_factor = math.pow(opt.min_size / (min([real_.shape[2], real_.shape[3]])), 1 / (opt.stop_scale)) # opt.scale_factor = math.pow(opt.min_size/(min(real_.shape[0],real_.shape[1])),1/(opt.stop_scale)) scale2stop = int( math.log(min([opt.max_size, max([real_.shape[2], real_.shape[3]])]) / max([real_.shape[2], real_.shape[3]]), opt.scale_factor_init)) opt.stop_scale = opt.num_scales - scale2stop return real
def test_pyramid(images): parser = get_arguments() parser.add_argument('--input_dir', help='input image dir', default='Input/Images') #parser.add_argument('--input_name', help='input image name', required=True) parser.add_argument('--mode', help='task to be done', default='train') opt = parser.parse_args("") opt.input_name = 'blank' opt = functions.post_config(opt) real = functions.np2torch(images[0], opt) functions.adjust_scales2image(real, opt) all_reals = [] for image in images: reals = [] real_ = functions.np2torch(image, opt) real = imresize(real_, opt.scale1, opt) reals = functions.creat_reals_pyramid(real, reals, opt) all_reals.append(reals) return np.array(all_reals).T
def train(opt, Gs, Zs, reals, NoiseAmp): real_ = functions.read_image(opt) in_s = 0 scale_num = 0 real = imresize(real_, opt.scale1, opt) reals = functions.creat_reals_pyramid(real, reals, opt) nfc_prev = 0 memory = [] ##storing memory time = [] ##storing time while scale_num < opt.stop_scale + 1: opt.nfc = min(opt.nfc_init * pow(2, math.floor(scale_num / 4)), 128) opt.min_nfc = min(opt.min_nfc_init * pow(2, math.floor(scale_num / 4)), 128) opt.out_ = functions.generate_dir2save(opt) opt.outf = '%s/%d' % (opt.out_, scale_num) try: os.makedirs(opt.outf) except OSError: pass #plt.imsave('%s/in.png' % (opt.out_), functions.convert_image_np(real), vmin=0, vmax=1) #plt.imsave('%s/original.png' % (opt.out_), functions.convert_image_np(real_), vmin=0, vmax=1) plt.imsave('%s/real_scale.png' % (opt.outf), functions.convert_image_np(reals[scale_num]), vmin=0, vmax=1) D_curr, G_curr = init_models(opt) if (nfc_prev == opt.nfc): G_curr.load_state_dict( torch.load('%s/%d/netG.pth' % (opt.out_, scale_num - 1))) D_curr.load_state_dict( torch.load('%s/%d/netD.pth' % (opt.out_, scale_num - 1))) start = datetime.datetime.now() z_curr, in_s, G_curr, mbs, percent = train_single_scale( D_curr, G_curr, reals, Gs, Zs, in_s, NoiseAmp, opt) memory.append([mbs, percent]) end = datetime.datetime.now() elapsed = end - start time.append(elapsed) print(f'time: {elapsed}') G_curr = functions.reset_grads(G_curr, False) G_curr.eval() D_curr = functions.reset_grads(D_curr, False) D_curr.eval() Gs.append(G_curr) Zs.append(z_curr) NoiseAmp.append(opt.noise_amp) torch.save(Zs, '%s/Zs.pth' % (opt.out_)) torch.save(Gs, '%s/Gs.pth' % (opt.out_)) torch.save(reals, '%s/reals.pth' % (opt.out_)) torch.save(NoiseAmp, '%s/NoiseAmp.pth' % (opt.out_)) torch.save(full_memory, '%s/full_memory.pth' % (opt.out_)) torch.save(full_time, '%s/full_time.pth' % (opt.out_)) scale_num += 1 nfc_prev = opt.nfc del D_curr, G_curr #torch.save(full_memory, '%s/full_memory.pk' % (opt.out_)) #torch.save(full_time, '%s/full_time.pk' % (opt.out_)) print(memory) print(time) print(full_memory) print(full_time) #pk.dump(full_memory, open('Full_memory', 'wb')) return
(opt.input_dir, opt.ref_name[:-4], opt.ref_name[-4:]), opt) mask = functions.read_image_dir( '%s/%s_mask%s' % (opt.ref_dir, opt.ref_name[:-4], opt.ref_name[-4:]), opt) if ref.shape[3] != real.shape[3]: mask = imresize_to_shape(mask, [real.shape[2], real.shape[3]], opt) mask = mask[:, :, :real.shape[2], :real.shape[3]] ref = imresize_to_shape(ref, [real.shape[2], real.shape[3]], opt) ref = ref[:, :, :real.shape[2], :real.shape[3]] mask = functions.dilate_mask(mask, opt) N = len(reals) - 1 n = opt.inpainting_start_scale in_s = imresize(ref, pow(opt.scale_factor, (N - n + 1)), opt) in_s = in_s[:, :, :reals[n - 1].shape[2], :reals[n - 1].shape[3]] in_s = imresize(in_s, 1 / opt.scale_factor, opt) in_s = in_s[:, :, :reals[n].shape[2], :reals[n].shape[3]] out = SinGAN_generate(Gs[n:], Zs[n:], reals, NoiseAmp[n:], opt, in_s, n=n, num_samples=1) out = (1 - mask) * real + mask * out plt.imsave('%s/start_scale=%d.png' % (dir2save, opt.inpainting_start_scale), functions.convert_image_np(out.detach()),
except OSError: pass real = functions.read_image(opt) real = functions.adjust_scales2image(real, opt) Gs, Zs, reals, NoiseAmp = functions.load_trained_pyramid(opt) if (opt.editing_start_scale < 1) | (opt.editing_start_scale > (len(Gs) - 1)): print("injection scale should be between 1 and %d" % (len(Gs) - 1)) else: ref = functions.read_image_dir( '%s/%s' % (opt.ref_dir, opt.ref_name), opt) mask = functions.read_image_dir( '%s/%s_mask%s' % (opt.ref_dir, opt.ref_name[:-4], opt.ref_name[-4:]), opt) if ref.shape[3] != real.shape[3]: mask = imresize(mask, real.shape[3] / ref.shape[3], opt) mask = mask[:, :, :real.shape[2], :real.shape[3]] ref = imresize(ref, real.shape[3] / ref.shape[3], opt) ref = ref[:, :, :real.shape[2], :real.shape[3]] mask = functions.dilate_mask(mask, opt) N = len(reals) - 1 n = opt.editing_start_scale in_s = imresize(ref, pow(opt.scale_factor, (N - n + 1)), opt) in_s = in_s[:, :, :reals[n - 1].shape[2], :reals[n - 1].shape[3]] in_s = imresize(in_s, 1 / opt.scale_factor, opt) in_s = in_s[:, :, :reals[n].shape[2], :reals[n].shape[3]] out = SinGAN_generate(Gs[n:], Zs[n:], reals, NoiseAmp[n:],
def generate_gif(Gs,Zs,reals,NoiseAmp,opt,alpha=0.1,beta=0.9,start_scale=2,fps=10): in_s = torch.full(Zs[0].shape, 0, device=opt.device) images_cur = [] count = 0 for G,Z_opt,noise_amp,real in zip(Gs,Zs,NoiseAmp,reals): pad_image = int(((opt.ker_size - 1) * opt.num_layer) / 2) nzx = Z_opt.shape[2] nzy = Z_opt.shape[3] #pad_noise = 0 #m_noise = nn.ZeroPad2d(int(pad_noise)) m_image = nn.ZeroPad2d(int(pad_image)) images_prev = images_cur images_cur = [] if count == 0: z_rand = functions.generate_noise([1,nzx,nzy], device=opt.device) z_rand = z_rand.expand(1,3,Z_opt.shape[2],Z_opt.shape[3]) z_prev1 = 0.95*Z_opt +0.05*z_rand z_prev2 = Z_opt else: z_prev1 = 0.95*Z_opt +0.05*functions.generate_noise([opt.nc_z,nzx,nzy], device=opt.device) z_prev2 = Z_opt for i in range(0,100,1): if count == 0: z_rand = functions.generate_noise([1,nzx,nzy], device=opt.device) z_rand = z_rand.expand(1,3,Z_opt.shape[2],Z_opt.shape[3]) diff_curr = beta*(z_prev1-z_prev2)+(1-beta)*z_rand else: diff_curr = beta*(z_prev1-z_prev2)+(1-beta)*(functions.generate_noise([opt.nc_z,nzx,nzy], device=opt.device)) z_curr = alpha*Z_opt+(1-alpha)*(z_prev1+diff_curr) z_prev2 = z_prev1 z_prev1 = z_curr if images_prev == []: I_prev = in_s else: I_prev = images_prev[i] I_prev = imresize(I_prev, 1 / opt.scale_factor, opt) I_prev = I_prev[:, :, 0:real.shape[2], 0:real.shape[3]] I_prev = m_image(I_prev) if count < start_scale: z_curr = Z_opt z_in = noise_amp*z_curr+I_prev I_curr = G(z_in.detach(),I_prev) if (count == len(Gs)-1): I_curr = functions.denorm(I_curr).detach() I_curr = I_curr[0,:,:,:].cpu().numpy() I_curr = I_curr.transpose(1, 2, 0)*255 I_curr = I_curr.astype(np.uint8) images_cur.append(I_curr) count += 1 dir2save = functions.generate_dir2save(opt) try: os.makedirs('%s/start_scale=%d' % (dir2save,start_scale) ) except OSError: pass imageio.mimsave('%s/start_scale=%d/alpha=%f_beta=%f.gif' % (dir2save,start_scale,alpha,beta),images_cur,fps=fps) del images_cur
real = functions.read_image(opt) opt.min_size = 18 real = functions.adjust_scales2image_SR(real, opt) train(opt, Gs, Zs, reals, NoiseAmp) opt.mode = mode print('%f' % pow(in_scale, iter_num)) Zs_sr = [] reals_sr = [] NoiseAmp_sr = [] Gs_sr = [] real = reals[-1] # read_image(opt) real_ = real opt.scale_factor = 1 / in_scale opt.scale_factor_init = 1 / in_scale for j in range(1, iter_num + 1, 1): real_ = imresize(real_, pow(1 / opt.scale_factor, 1), opt) reals_sr.append(real_) Gs_sr.append(Gs[-1]) NoiseAmp_sr.append(NoiseAmp[-1]) z_opt = torch.full(real_.shape, 0, dtype=torch.float32, device=opt.device) m = nn.ZeroPad2d(5) z_opt = m(z_opt) Zs_sr.append(z_opt) out = SinGAN_generate(Gs_sr, Zs_sr, reals_sr, NoiseAmp_sr, opt,
def train(opt, Gs, Zs, reals, NoiseAmp): real_ = functions.read_image(opt) in_s = 0 scale_num = 0 real = imresize(real_, opt.scale1, opt) reals = functions.creat_reals_pyramid(real, reals, opt) nfc_prev = 0 while scale_num < opt.stop_scale + 1: opt.nfc = min(opt.nfc_init * pow(2, math.floor(scale_num / 4)), 128) opt.min_nfc = min(opt.min_nfc_init * pow(2, math.floor(scale_num / 4)), 128) if opt.fast_training: if (scale_num > 0) & (scale_num % 4 == 0): opt.niter = opt.niter // 2 ''' if (scale_num == opt.stop_scale): opt.nfc = 128 opt.min_nfc = 128 ''' opt.out_ = functions.generate_dir2save(opt) opt.outf = '%s/%d' % (opt.out_, scale_num) try: os.makedirs(opt.outf) except OSError: pass #plt.imsave('%s/in.png' % (opt.out_), functions.convert_image_np(real), vmin=0, vmax=1) #plt.imsave('%s/original.png' % (opt.out_), functions.convert_image_np(real_), vmin=0, vmax=1) plt.imsave('%s/real_scale.png' % (opt.outf), functions.convert_image_np(reals[scale_num]), vmin=0, vmax=1) D_curr, G_curr = init_models(opt) if (nfc_prev == opt.nfc): G_curr.load_state_dict( torch.load('%s/%d/netG.pth' % (opt.out_, scale_num - 1))) D_curr.load_state_dict( torch.load('%s/%d/netD.pth' % (opt.out_, scale_num - 1))) z_curr, in_s, G_curr = train_single_scale(D_curr, G_curr, reals, Gs, Zs, in_s, NoiseAmp, opt) G_curr = functions.reset_grads(G_curr, False) G_curr.eval() D_curr = functions.reset_grads(D_curr, False) D_curr.eval() Gs.append(G_curr) Zs.append(z_curr) NoiseAmp.append(opt.noise_amp) torch.save(Zs, '%s/Zs.pth' % (opt.out_)) torch.save(Gs, '%s/Gs.pth' % (opt.out_)) torch.save(reals, '%s/reals.pth' % (opt.out_)) torch.save(NoiseAmp, '%s/NoiseAmp.pth' % (opt.out_)) scale_num += 1 nfc_prev = opt.nfc del D_curr, G_curr return