def __init__(self, input_name, load_existing_model):
self.input_name = input_name
self.load_existing_model = load_existing_model
self.opt = Config(input_name)
self.dir2save = functions.generate_dir2save(self.opt)
self.real = functions.read_image(self.opt)
functions.adjust_scales2image(self.real, self.opt)
dir_exists = os.path.exists(self.dir2save)
assert (not load_existing_model) or dir_exists, "cannot find trained model"
if load_existing_model:
print("Trained model has been loaded (not really)")
self.Gs = torch.load(f'{self.dir2save}/Gs.pth')
self.Zs = torch.load(f'{self.dir2save}/Zs.pth')
self.reals = torch.load(f'{self.dir2save}/reals.pth')
self.NoiseAmp = torch.load(f'{self.dir2save}/NoiseAmp.pth')
self.is_loaded = True
else:
if dir_exists:
user_input = input("Trained model has been found, type \"yes\" to overwrite: ")
assert user_input == 'yes', "train aborted"
rmtree(self.dir2save)
print("train directory has been deleted")
try:
os.makedirs(self.dir2save)
except OSError:
pass
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 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 main(opt):
Gs = []
Zs = []
reals = []
NoiseAmp = []
dir2save = functions.generate_dir2save(opt)
if os.path.exists(dir2save):
logger.info("Trained model directory already exists")
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 main(opt, generate=True):
Gs = []
Zs: List[Tuple] = []
reals1 = []
reals2 = []
NoiseAmp = []
dir2save = functions.generate_dir2save(opt)
if (os.path.exists(dir2save)):
print('trained model already exist')
else:
try:
os.makedirs(dir2save)
except OSError:
pass
_configure_logger(dir2save)
# dump configuration file to json
with open(os.path.join(f"{dir2save}", "config.json"), "w") as fp:
config_dict = {k: str(v) for k, v in opt.__dict__.items()}
json.dump(config_dict, fp)
try:
real1 = functions.read_image(opt, image_name=opt.input_name1)
real2 = functions.read_image(opt, image_name=opt.input_name2)
functions.adjust_scales2image(real1, opt)
functions.adjust_scales2image(real2, opt)
train(opt, Gs, Zs, reals1, reals2, NoiseAmp)
logger.info("Done training")
if generate:
logger.info("Generating random samples")
SinGAN_generate(Gs, Zs, reals1, reals2, NoiseAmp, opt)
except Exception as e:
logger.exception("Failed")
raise
finally:
logger.info("Cleaning logger")
_cleanup_logger()
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
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)
train(opt, Gs, Zs, reals, NoiseAmp)
SinGAN_generate(Gs, Zs, reals, NoiseAmp, opt)
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('--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) # adjust the pyramid parameters according to real image scale
train(opt, Gs, Zs, reals, NoiseAmp)
SinGAN_generate(Gs,Zs,reals,NoiseAmp,opt)
Zs = []
reals = []
NoiseAmp = []
dir2save = functions.generate_dir2save(opt)
if dir2save is None:
print('task does not exist')
else:
try:
os.makedirs(dir2save)
except OSError:
pass
#real = functions.read_image(opt)
#real = functions.adjust_scales2image(real, opt)
#Gs, Zs, reals, NoiseAmp = functions.load_trained_pyramid(opt)
real1, real2 = functions.read_image(opt)
real1 = functions.adjust_scales2image(real1, opt)
real2 = functions.adjust_scales2image(real2, opt)
Gs, Zs, reals1, reals2, NoiseAmp = functions.load_trained_pyramid(opt)
if (opt.paint_start_scale < 1) | (opt.paint_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)
if ref.shape[3] != real1.shape[3]:
ref = imresize_to_shape(ref, [real1.shape[2], real1.shape[3]], opt)
ref = ref[:, :, :real1.shape[2], :real1.shape[3]]
N = len(reals1) - 1
n = opt.paint_start_scale
in_s = imresize(ref, pow(opt.scale_factor, (N - n + 1)), opt)
in_s = in_s[:, :, :reals1[n - 1].shape[2], :reals1[n - 1].shape[3]]
in_s = imresize(in_s, 1 / opt.scale_factor, 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
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) # opt.input_dir 과 opt.input_name을 이용하여 이미지를 array형식으로 받아옴
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) # opt.input_dir 과 opt.input_name을 이용하여 이미지를 array형식으로 받아옴
functions.adjust_scales2image(real, opt) #opt를 설정해줌
Gs, Zs, reals, NoiseAmp = functions.load_trained_pyramid(
opt) # 총 4개의 모델을 불러옴
