def generate(model_name,
anchor_image=None,
direction=None,
transfer=None,
noise_solutions=None,
factor=0.25,
base=None,
insert_limit=4):
#direction = 'L, R, T, B'
parser = get_arguments()
parser.add_argument('--input_dir',
help='input image dir',
default='Input/Images')
parser.add_argument('--mode',
help='random_samples | random_samples_arbitrary_sizes',
default='random_samples')
# for random_samples:
parser.add_argument('--gen_start_scale',
type=int,
help='generation start scale',
default=0)
opt = parser.parse_args("")
opt.input_name = model_name
if model_name == 'islands2_basis_2.jpg': #HARDCODED
opt.scale_factor = 0.6
opt = functions.post_config(opt)
Gs = []
Zs = []
reals = []
NoiseAmp = []
real = functions.read_image(opt)
#opt.input_name = anchor #CHANGE TO ANCHOR HERE
anchor = functions.read_image(opt)
functions.adjust_scales2image(real, opt)
Gs, Zs, reals, NoiseAmp = functions.load_trained_pyramid(opt)
in_s = functions.generate_in2coarsest(reals, 1, 1, opt)
array = SinGAN_anchor_generate(Gs,
Zs,
reals,
NoiseAmp,
opt,
gen_start_scale=opt.gen_start_scale,
anchor_image=anchor_image,
direction=direction,
transfer=transfer,
noise_solutions=noise_solutions,
factor=factor,
base=base,
insert_limit=insert_limit)
return array
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
print(
'random samples for image %s, start scale=%d, already exist' %
(opt.input_name, opt.gen_start_scale))
elif opt.mode == 'random_samples_arbitrary_sizes':
print(
'random samples for image %s at size: scale_h=%f, scale_v=%f, already exist'
% (opt.input_name, opt.scale_h, opt.scale_v))
else:
try:
os.makedirs(dir2save)
except OSError:
pass
if opt.mode == 'random_samples':
real = functions.read_image(opt)
functions.adjust_scales2image(real, opt)
Gs, Zs, reals, NoiseAmp = functions.load_trained_pyramid(opt)
# 生成最粗糙的图像
in_s = functions.generate_in2coarsest(reals, 1, 1, opt)
SinGAN_generate(Gs,
Zs,
reals,
NoiseAmp,
opt,
gen_start_scale=opt.gen_start_scale)
# elif opt.mode == 'random_samples_arbitrary_sizes':
# real = functions.read_image(opt)
# functions.adjust_scales2image(real, opt)
# Gs, Zs, reals, NoiseAmp = functions.load_trained_pyramid(opt)
# # 传入scale_v, scale_h 生成任意尺寸的图像 scale 1:1.5
def test_generate(model_name,
anchor_image=None,
direction=None,
transfer=None,
noise_solutions=None,
factor=0.25,
base=None,
insert_limit=4):
#direction = 'L, R, T, B'
parser = get_arguments()
parser.add_argument('--input_dir',
help='input image dir',
default='Input/Images')
parser.add_argument('--mode',
help='random_samples | random_samples_arbitrary_sizes',
default='random_samples')
# for random_samples:
parser.add_argument('--gen_start_scale',
type=int,
help='generation start scale',
default=0)
opt = parser.parse_args("")
opt.input_name = model_name
opt = functions.post_config(opt)
Gs = []
Zs = []
reals = []
NoiseAmp = []
opt.input_name = 'island_basis_0.jpg' #grabbing image that exists...
real = functions.read_image(opt)
#opt.input_name = anchor #CHANGE TO ANCHOR HERE
#anchor = functions.read_image(opt)
functions.adjust_scales2image(real, opt)
opt.input_name = 'test1.jpg' #grabbing model that we want
Gs, Zs, reals, NoiseAmp = functions.load_trained_pyramid(opt)
#dummy stuff for dimensions
reals = []
real_ = real
real = imresize(real_, opt.scale1, opt)
reals = functions.creat_reals_pyramid(real, reals, opt)
in_s = functions.generate_in2coarsest(reals, 1, 1, opt)
array = SinGAN_anchor_generate(Gs,
Zs,
reals,
NoiseAmp,
opt,
gen_start_scale=opt.gen_start_scale,
anchor_image=anchor_image,
direction=direction,
transfer=transfer,
noise_solutions=noise_solutions,
factor=factor,
base=base,
insert_limit=insert_limit)
return array
def invert_model(test_image,
model_name,
scales2invert=None,
penalty=1e-3,
show=True):
'''test_image is an array, model_name is a name'''
Noise_Solutions = []
parser = get_arguments()
parser.add_argument('--input_dir',
help='input image dir',
default='Input/Images')
parser.add_argument('--mode', default='RandomSamples')
opt = parser.parse_args("")
opt.input_name = model_name
opt.reg = penalty
if model_name == 'islands2_basis_2.jpg': #HARDCODED
opt.scale_factor = 0.6
opt = functions.post_config(opt)
### Loading in Generators
Gs, Zs, reals, NoiseAmp = functions.load_trained_pyramid(opt)
for G in Gs:
G = functions.reset_grads(G, False)
G.eval()
### Loading in Ground Truth Test Images
reals = [] #deleting old real images
real = functions.np2torch(test_image, opt)
functions.adjust_scales2image(real, opt)
real_ = functions.np2torch(test_image, opt)
real = imresize(real_, opt.scale1, opt)
reals = functions.creat_reals_pyramid(real, reals, opt)
### General Padding
pad_noise = int(((opt.ker_size - 1) * opt.num_layer) / 2)
m_noise = nn.ZeroPad2d(int(pad_noise))
pad_image = int(((opt.ker_size - 1) * opt.num_layer) / 2)
m_image = nn.ZeroPad2d(int(pad_image))
I_prev = None
REC_ERROR = 0
if scales2invert is None:
scales2invert = opt.stop_scale + 1
for scale in range(scales2invert):
#for scale in range(3):
#Get X, G
X = reals[scale]
G = Gs[scale]
noise_amp = NoiseAmp[scale]
#Defining Dimensions
opt.nc_z = X.shape[1]
opt.nzx = X.shape[2]
opt.nzy = X.shape[3]
#getting parameters for prior distribution penalty
pdf = torch.distributions.Normal(0, 1)
alpha = opt.reg
#alpha = 1e-2
#Defining Z
if scale == 0:
z_init = functions.generate_noise(
[1, opt.nzx, opt.nzy], device=opt.device) #only 1D noise
else:
z_init = functions.generate_noise(
[3, opt.nzx, opt.nzy],
device=opt.device) #otherwise move up to 3d noise
z_init = Variable(z_init.cuda(),
requires_grad=True) #variable to optimize
#Building I_prev
if I_prev == None: #first scale scenario
in_s = torch.full(reals[0].shape, 0, device=opt.device) #all zeros
I_prev = in_s
I_prev = m_image(I_prev) #padding
else: #otherwise take the output from the previous scale and upsample
I_prev = imresize(I_prev, 1 / opt.scale_factor, opt) #upsamples
I_prev = m_image(I_prev)
I_prev = I_prev[:, :, 0:X.shape[2] + 10, 0:X.shape[
3] + 10] #making sure that precision errors don't mess anything up
I_prev = functions.upsampling(I_prev, X.shape[2] + 10, X.shape[3] +
10) #seems to be redundant
LR = [2e-3, 2e-2, 2e-1, 2e-1, 2e-1, 2e-1, 2e-1, 2e-1, 2e-1, 2e-1, 2e-1]
Zoptimizer = torch.optim.RMSprop([z_init],
lr=LR[scale]) #Defining Optimizer
x_loss = [] #for plotting
epochs = [] #for plotting
niter = [
200, 400, 400, 400, 400, 400, 400, 400, 400, 400, 400, 400, 400,
400, 400
]
for epoch in range(niter[scale]): #Gradient Descent on Z
if scale == 0:
noise_input = m_noise(z_init.expand(1, 3, opt.nzx,
opt.nzy)) #expand and padd
else:
noise_input = m_noise(z_init) #padding
z_in = noise_amp * noise_input + I_prev
G_z = G(z_in, I_prev)
x_recLoss = F.mse_loss(G_z, X) #MSE loss
logProb = pdf.log_prob(z_init).mean() #Gaussian loss
loss = x_recLoss - (alpha * logProb.mean())
Zoptimizer.zero_grad()
loss.backward()
Zoptimizer.step()
#losses['rec'].append(x_recLoss.data[0])
#print('Image loss: [%d] loss: %0.5f' % (epoch, x_recLoss.item()))
#print('Noise loss: [%d] loss: %0.5f' % (epoch, z_recLoss.item()))
x_loss.append(loss.item())
epochs.append(epoch)
REC_ERROR = x_recLoss
if show:
plt.plot(epochs, x_loss, label='x_loss')
plt.legend()
plt.show()
I_prev = G_z.detach(
) #take final output, maybe need to edit this line something's very very fishy
_ = show_image(X, show, 'target')
reconstructed_image = show_image(I_prev, show, 'output')
_ = show_image(noise_input.detach().cpu(), show, 'noise')
Noise_Solutions.append(noise_input.detach())
return Noise_Solutions, reconstructed_image, REC_ERROR
