def train_model(input_name):
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 = functions.post_config(opt)
Gs = []
Zs = []
reals = []
NoiseAmp = []
dir2save = functions.generate_dir2save(opt)
if (os.path.exists(dir2save)):
print('trained model already exist')
else:
try:
os.makedirs(dir2save)
except OSError:
pass
real = functions.read_image(opt)
functions.adjust_scales2image(real, opt)
train(opt, Gs, Zs, reals, NoiseAmp)
SinGAN_generate(Gs, Zs, reals, NoiseAmp, opt)
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 cli():
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 = functions.post_config(opt)
main(opt)
def parse_arguments(namespace=None):
parser = get_arguments()
parser.add_argument('--input_dir',
help='input image dir',
default='Input/Images')
parser.add_argument('--input_name1',
help='input image name 1',
required=namespace is None)
parser.add_argument('--input_name2',
help='input image name 2',
required=namespace is None)
parser.add_argument('--mode', help='task to be done', default='train')
opt = parser.parse_args(namespace=namespace)
opt = functions.post_config(opt)
return opt
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
from config import get_arguments
from SinGAN.manipulate import *
from SinGAN.training import *
import SinGAN.functions as functions
if __name__ == '__main__':
parser = get_arguments()
parser.add_argument('--input_dir',
help='input image dir',
default='Input/Images')
parser.add_argument('--model_name',
help='input image name -1',
required=True)
parser.add_argument('--mode', help='task to be done', default='train')
opt = parser.parse_args()
opt = functions.post_config(opt)
Gs = []
Zs = []
reals = []
NoiseAmp = []
dir2save = functions.generate_dir2save(opt)
if (os.path.exists(dir2save)):
print('trained model already exist')
else:
try:
os.makedirs(dir2save)
except OSError:
pass
real = functions.read_images(opt)
functions.adjust_scales2image(real, opt)
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
