def upload():
global zip_path, zip_file
target_img_path = ' '
if request.method == 'POST':
generator = Pix2Pix()
id_style = request.form["id_style"]
generator.set_weight(str(id_style))
files = request.files.getlist('file')
if len(files) > 1:
zip_path = './static/' + files[0].filename + 'style.zip'
zip_file = zipfile.ZipFile(zip_path, 'w')
for file in files:
if PRODUCTION:
input_img_path = os.path.join(app.config['UPLOAD_FOLDER'],
file.filename)
file.save(input_img_path)
target_img_path = generator.predict(input_img_path)
img = Image(path_input_img=input_img_path,
path_target_img=target_img_path,
style=id_style)
if len(files) > 1:
zip_file.write(target_img_path)
db.session.add(img)
db.session.commit()
if len(files) > 1:
target_img_path = zip_path
zip_file.close()
return jsonify(link=target_img_path[1:])
def test1_nobn_finetunep2p_bilin(mode):
assert mode in ["train", "interp","gen"]
from architectures import p2p, dcgan
# change the p2p discriminator
model = Pix2Pix(
gen_fn_dcgan=dcgan.default_generator,
disc_fn_dcgan=dcgan.default_discriminator,
gen_params_dcgan={'num_repeats':0, 'div':[2,2,4,4,8,8,8]},
disc_params_dcgan={'num_repeats':0, 'bn':False, 'nonlinearity':linear, 'div':[8,4,4,4,2,2,2]},
gen_fn_p2p=p2p.g_unet,
disc_fn_p2p=p2p.discriminator,
gen_params_p2p={'nf':64, 'act':tanh, 'num_repeats':0, 'bilinear_upsample':True},
disc_params_p2p={'nf':64, 'bn':False, 'num_repeats':0, 'act':linear, 'mul_factor':[1,2,4,8]},
in_shp=512,
latent_dim=1000,
is_a_grayscale=True,
is_b_grayscale=False,
lsgan=True,
opt=rmsprop,
opt_args={'learning_rate':theano.shared(floatX(1e-4))},
train_mode='p2p'
)
model.load_model("models/test1_repeatnod_fixp2p_nobn/1000.model.bak", mode='dcgan') # only load the dcgan
desert_h5 = "/data/lisa/data/cbeckham/textures_v2_brown500.h5"
bs = 4
it_train, it_val = get_iterators(desert_h5, bs, True, False, True)
name = "test1_repeatnod_fixp2p_nobn_finetunep2p_bilin"
# models/test1_repeatnod_fixp2p_nobn/600.model.bak --> good DCGAN?? (not good p2p!!)
# models/test1_repeatnod_fixp2p_nobn_finetunep2p_bilin/1000.model.bak --> good p2p (not good DCGAN??)
if mode == "train":
model.train(it_train, it_val, batch_size=bs, num_epochs=1000, out_dir="output/%s" % name, model_dir="models/%s" % name)
elif mode == "interp":
model.load_model("models/test1_repeatnod_fixp2p_nobn/1000.model.bak", mode='dcgan')
model.load_model("models/test1_repeatnod_fixp2p_nobn_finetunep2p_bilin/1000.model.bak", mode='p2p')
model.generate_interpolation_clip(100, 4, "output/%s/interp_clip_600_concat_bothdet/" % name, concat=True, deterministic=True)
def fit(train_ds, test_ds, steps, preview=0):
pix2pix = Pix2Pix()
pix2pix.compile()
example_input, example_target = next(iter(test_ds.take(1)))
start = time.time()
for step, batch in train_ds.repeat().take(steps).enumerate():
# Training step
losses = pix2pix.train_step(batch, verbose=True)
if (step + 1) % 10 == 0:
print(".", end="", flush=True)
if preview and ((step + 1) % preview == 0):
if step != 0:
print(
f"Time taken for {preview} steps: {time.time()-start:.2f} sec\n"
)
print("g_gan_loss: ", losses["g_gan_loss"])
print("g_l1_loss: ", losses["g_l1_loss"])
print("g_loss: ", losses["g_loss"])
print("d_loss: ", losses["d_loss"])
print("total_loss: ", losses["total_loss"])
val_losses = pix2pix.test_step(next(iter(test_ds)),
verbose=True)
print("val_g_gan_loss: ", val_losses["g_gan_loss"])
print("val_g_l1_loss: ", val_losses["g_l1_loss"])
print("val_g_loss: ", val_losses["g_loss"])
print("val_d_loss: ", val_losses["d_loss"])
print("val_total_loss: ", val_losses["total_loss"])
generate_images(pix2pix.generator, example_input,
example_target)
print(f"Step: {step + 1}")
start = time.time()
def test1_nobn_bilin_both(mode):
assert mode in ["train", "interp","gen"]
from architectures import p2p, dcgan
# change the p2p discriminator
model = Pix2Pix(
gen_fn_dcgan=dcgan.default_generator,
disc_fn_dcgan=dcgan.default_discriminator,
gen_params_dcgan={'num_repeats':0, 'div':[2,2,4,4,8,8,8]},
disc_params_dcgan={'num_repeats':0, 'bn':False, 'nonlinearity':linear, 'div':[8,4,4,4,2,2,2]},
gen_fn_p2p=p2p.g_unet,
disc_fn_p2p=p2p.discriminator,
gen_params_p2p={'nf':64, 'act':tanh, 'num_repeats':0, 'bilinear_upsample':True},
disc_params_p2p={'nf':64, 'bn':False, 'num_repeats':0, 'act':linear, 'mul_factor':[1,2,4,8]},
in_shp=512,
latent_dim=1000,
is_a_grayscale=True,
is_b_grayscale=False,
lsgan=True,
opt=rmsprop,
opt_args={'learning_rate':theano.shared(floatX(1e-4))},
train_mode='both'
)
desert_h5 = "/data/lisa/data/cbeckham/textures_v2_brown500.h5"
bs = 4
it_train, it_val = get_iterators(desert_h5, bs, True, False, True)
name = "test1_nobn_bilin_both_deleteme"
if mode == "train":
model.train(it_train, it_val, batch_size=bs, num_epochs=1000, out_dir="output/%s" % name, model_dir="models/%s" % name)
def test1_nobn(mode):
assert mode in ["train", "interp", "gen"]
from architectures import p2p, dcgan
model = Pix2Pix(
gen_fn_dcgan=dcgan.default_generator,
disc_fn_dcgan=dcgan.default_discriminator,
gen_params_dcgan={'num_repeats':0, 'div':[2,2,4,4,8,8,8]},
disc_params_dcgan={'num_repeats':0, 'bn':False, 'nonlinearity':linear, 'div':[8,4,4,4,2,2,2]},
gen_fn_p2p=p2p.g_unet,
disc_fn_p2p=p2p.discriminator,
gen_params_p2p={'nf':64, 'act':tanh, 'num_repeats':0},
disc_params_p2p={'nf':64, 'bn':False, 'num_repeats':0, 'act':linear, 'mul_factor':[1,2,4,8]},
in_shp=512,
latent_dim=1000,
is_a_grayscale=True,
is_b_grayscale=False,
lsgan=True,
opt=rmsprop,
opt_args={'learning_rate':theano.shared(floatX(1e-4))}
)
desert_h5 = "/data/lisa/data/cbeckham/textures_v2_brown500.h5"
bs = 4
it_train, it_val = get_iterators(desert_h5, bs, True, False, True)
name = "test1_repeatnod_fixp2p_nobn"
if mode == "train":
model.train(it_train, it_val, batch_size=bs, num_epochs=1000, out_dir="output/%s" % name, model_dir="models/%s" % name)
elif mode == "interp":
model.load_model("models/%s/600.model.bak" % name)
zs = model.sampler(2, model.latent_dim)
z1, z2 = floatX(zs[0]), floatX(zs[1])
model.generate_interpolation(z1, z2, "/tmp/test.png", mode='matrix')
elif mode == "gen":
model.load_model("models/%s/600.model.bak" % name)
model.generate_gz(100, 10, "deleteme")
def main():
parser = argparse.ArgumentParser()
parser.add_argument('x_dir', type=str)
parser.add_argument('y_dir', type=str)
parser.add_argument('--batch_size', '-bs', type=int, default=64)
parser.add_argument('--nb_epoch', '-e', type=int, default=1000)
parser.add_argument('--height', '-ht', type=int, default=256)
parser.add_argument('--width', '-wd', type=int, default=256)
parser.add_argument('--save_steps', '-ss', type=int, default=10)
parser.add_argument('--visualize_steps', '-vs', type=int, default=10)
parser.add_argument('--gp_weight', '-gp', type=float, default=10.)
parser.add_argument('--l1_weight', '-l1', type=float, default=1.)
parser.add_argument('--initial_steps', '-is', type=int, default=20)
parser.add_argument('--initial_critics', '-ic', type=int, default=20)
parser.add_argument('--normal_critics', '-nc', type=int, default=5)
parser.add_argument('--model_dir', '-md', type=str, default="./params")
parser.add_argument('--result_dir', '-rd', type=str, default="./result")
parser.add_argument('--noise_mode', '-nm', type=str, default="uniform")
parser.add_argument('--upsampling', '-up', type=str, default="deconv")
parser.add_argument('--dis_norm', '-dn', type=str, default=None)
args = parser.parse_args()
os.makedirs(args.result_dir, exist_ok=True)
os.makedirs(args.model_dir, exist_ok=True)
image_sampler = ImageSampler(target_size=(args.width, args.height),
color_mode_x='rgb',
color_mode_y='rgb',
normalization_x='tanh',
normalization_y='tanh',
is_flip=False)
generator = UNet((args.height, args.width, 3),
color_mode='rgb',
upsampling=args.upsampling,
is_training=True)
discriminator = ResidualDiscriminator((args.height, args.width, 6),
normalization=args.dis_norm,
is_training=True)
pix2pix = Pix2Pix(generator,
discriminator,
l1_weight=args.l1_weight,
gradient_penalty_weight=args.gp_weight,
is_training=True)
pix2pix.fit(image_sampler.flow_from_directory(args.x_dir,
args.y_dir,
batch_size=args.batch_size),
result_dir=args.result_dir,
model_dir=args.model_dir,
save_steps=args.save_steps,
visualize_steps=args.visualize_steps,
nb_epoch=args.nb_epoch,
initial_steps=args.initial_steps,
initial_critics=args.initial_critics,
normal_critics=args.normal_critics)
def __init__(self, parent, **kwargs):
super().__init__(parent, **kwargs)
#button.pack()
self._pix = Pix2Pix(restore=True)
self._generator = self._pix._generator
self._data_loader = ImageDataLoader()
self._data_loader.load_datasets(
"output_edges/train/edges.tfrecords",
"output_edges/test/test_edges.tfrecords")
self.bind("<Button-1>", self.save_posn)
self.bind("<B1-Motion>", self.add_line)
def __init__(self, flags):
run_config = tf.ConfigProto()
run_config.gpu_options.allow_growth = True
self.sess = tf.Session(config=run_config)
self.flags = flags
self.image_size = (256, 256, 3)
self.dataset = dataset(self.flags.dataset, image_size=self.image_size)
self.model = Pix2Pix(self.sess, self.flags, image_size=self.image_size)
self._make_folders()
self.saver = tf.train.Saver()
self.sess.run(tf.global_variables_initializer())
def gan_repository(sess, flags, dataset):
if flags.gan_model == 'vanilla_gan':
print('Initializing Vanilla GAN...')
return GAN(sess, flags, dataset.image_size)
elif flags.gan_model == 'dcgan':
print('Initializing DCGAN...')
return DCGAN(sess, flags, dataset.image_size)
elif flags.gan_model == 'pix2pix':
print('Initializing pix2pix...')
return Pix2Pix(sess, flags, dataset.image_size)
elif flags.gan_model == 'pix2pix-patch':
print('Initializing pix2pix-patch...')
return Pix2PixPatch(sess, flags, dataset.image_size)
elif flags.gan_model == 'wgan':
print('Initializing WGAN...')
return WGAN(sess, flags, dataset)
elif flags.gan_model == 'cyclegan':
print('Initializing cyclegan...')
return CycleGAN(sess, flags, dataset.image_size, dataset())
elif flags.gan_model == 'mrigan':
print('Initializing mrigan...')
return MRIGAN(sess, flags, dataset.image_size, dataset())
elif flags.gan_model == 'mrigan02':
print('Initializing mrigan02...')
return MRIGAN02(sess, flags, dataset.image_size, dataset())
elif flags.gan_model == 'mrigan03':
print('Initializing mrigan03...')
return MRIGAN03(sess, flags, dataset.image_size, dataset())
elif flags.gan_model == 'mrigan01_lsgan':
print('Initializing mrigan01_lsgan...')
return MRIGAN01_LSGAN(sess, flags, dataset.image_size, dataset())
elif flags.gan_model == 'mrigan02_lsgan':
print('Initializing mrigan02_lsgan...')
return MRIGAN02_LSGAN(sess, flags, dataset.image_size, dataset())
elif flags.gan_model == 'mrigan03_lsgan':
print('Initializing mrigan03_lsgan...')
return MRIGAN03_LSGAN(sess, flags, dataset.image_size, dataset())
elif flags.gan_model == 'mrigan_01':
print('Initializing mrigan_01...')
return MRIGAN_01(sess, flags, dataset.image_size, dataset())
elif flags.gan_model == 'mrigan_02':
print('Initializing mrigan_02...')
return MRIGAN_02(sess, flags, dataset.image_size, dataset())
else:
raise NotImplementedError
def __init__(self, flags):
run_config = tf.ConfigProto()
run_config.gpu_options.allow_growth = True
self.sess = tf.Session(config=run_config)
self.flags = flags
self.iter_time = 0
self._make_folders()
self._init_logger()
self.dataset = Dataset(self.flags.dataset, self.flags, log_path=self.log_out_dir)
self.model = Pix2Pix(self.sess, self.flags, self.dataset.image_size, self.dataset(self.flags.is_train),
log_path=self.log_out_dir)
self.saver = tf.train.Saver()
self.sess.run(tf.global_variables_initializer())
tf_utils.show_all_variables()
def build_model(self) -> tf.keras.models.Model:
model = Pix2Pix()
# Wrap the model
model = self.context.wrap_model(model)
# Create and wrap the optimizers
g_optimizer = self.context.wrap_optimizer(
make_generator_optimizer(
lr=self.context.get_hparam("generator_lr"),
beta_1=self.context.get_hparam("generator_beta_1"),
))
d_optimizer = self.context.wrap_optimizer(
make_discriminator_optimizer(
lr=self.context.get_hparam("discriminator_lr"),
beta_1=self.context.get_hparam("discriminator_beta_1"),
))
model.compile(
discriminator_optimizer=d_optimizer,
generator_optimizer=g_optimizer,
)
return model
WORKING_DIR = os.path.dirname(__file__)
DEFAULT_MODEL_PATH = os.path.join(WORKING_DIR, './weights')
DEFAULT_MODEL_PATH_OUT = os.path.join(WORKING_DIR, './out')
def arguments():
parser = argparse.ArgumentParser()
parser.add_argument('-v', dest='validation')
parser.add_argument('-t', dest='train')
parser.add_argument('-history', dest='history', action='count', default=0)
parser.add_argument('-m',
dest='model',
default=DEFAULT_MODEL_PATH,
help="Folder to the weight files")
parser.add_argument('-out',
dest='out_model',
default=DEFAULT_MODEL_PATH_OUT,
help="Folder to store the output weight files")
return parser.parse_args()
if __name__ == '__main__':
args = arguments()
gan = Pix2Pix(args.train, args.validation, args.model, args.out_model)
if (args.history > 0):
gan.load_history()
gan.plot_history()
gan.train(epochs=50, batch_size=2, sample_interval=25)
def train(opt, model_name):
if model_name == 'pix2pix': #This is for Arpit's version of Pix2Pix
model = Pix2Pix(opt)
# best_model_wts = copy.deepcopy(model.state_dict())
# best_acc = 0.0
for epoch in range(opt.epochs):
since = time.time()
print('Epoch ' + str(epoch) + ' running')
for phase in range(2):
val_dice = 0
count = 0
for i, Data in enumerate(dataloader[phase]):
inputs, masks = Data
inputs, masks = inputs.to(device), masks.to(device)
inputs = normalization(inputs) ##Changes made here
masks = normalization(masks)
inputs, masks = Variable(inputs), Variable(
masks) ##Ye Variable kyu likha hai ek baar batana
Data = inputs, masks ## --> kyuki isse computation fast ho jaata
## memoization ki vajah se
with torch.set_grad_enabled(phase == 0):
model.get_input(Data)
if phase == 0:
model.optimize()
else:
pred_mask = model.forward(inputs)
for j in range(pred_mask.size()[0]):
cv2.imwrite(
os.path.join(
'../results/pred_masks',
'mask_{}_{}_{}.png'.format(
i, j, epoch)),
np.array(
denormalize(pred_mask[j]).cpu().detach(
)).reshape(256, 256, 3))
cv2.imwrite(
os.path.join(
'../results/inputs',
'input_{}_{}_{}.png'.format(
i, j, epoch)),
np.array(
denormalize(
inputs[j]).cpu().detach()).reshape(
256, 256, 3))
val_dice += dice_coeff(
denormalize(pred_mask, flag=1),
denormalize(masks, flag=1))
count += 1
print("Validation Dice Coefficient is " + str(val_dice / count))
time_elapsed = time.time() - since
print('Epoch completed in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
elif model_name == 'CycleGAN':
model = cycleGan(cg_opt)
print_freq = 10
train_iter = iter(cg_train_loader)
val_iter = iter(cg_val_loader)
fixed_X, fixed_Y = val_iter.next()
fixed_X = normalization(fixed_X).to(device)
fixed_Y = normalization(fixed_Y).to(device)
loss_Gl = []
loss_DXl = []
loss_DYl = []
num_batches = len(train_iter)
for epoch in range(200):
if epoch == 35:
model.change_lr(model.opt.lr / 2)
if epoch == 80:
model.change_lr(model.opt.lr / 2)
if epoch == 130:
model.change_lr(model.opt.lr / 2)
since = time.time()
print("Epoch ", epoch, " entering ")
train_iter = iter(cg_train_loader)
for batch in range(num_batches):
print("Epoch ", epoch, "Batch ", batch,
" running with learning rate ", model.opt.lr)
inputX, inputY = train_iter.next()
inputX = normalization(inputX).to(device)
inputY = normalization(inputY).to(device)
model.get_input(inputX, inputY)
model.optimize()
# print("Dx Loss : {:.6f} Dy Loss: {:.6f} Generator Loss: {:.6f} ".format(model.dx_loss, model.dy_loss, model.gen_loss))
print("Model dx loss ", float(model.loss_D_X), "Model dy loss",
float(model.loss_D_Y), "model_gen_loss",
float(model.loss_G))
if (epoch + 1) % 10 == 0:
# torch.set_grad_enabled(False)
depth_map = model.G_XtoY.forward(fixed_X)
for j in range(depth_map.size()[0]):
if cg_opt.n_blocks == 6:
cv2.imwrite(
os.path.join(
'../cgresults/pred_masks',
'mask_{}_{}_{}.png'.format(batch, j, epoch)),
np.array(denormalize(
depth_map[j]).cpu().detach()).reshape(
256, 256, 3))
if epoch == 9:
cv2.imwrite(
os.path.join(
'../cgresults/inputs',
'input_{}_{}_{}.png'.format(
batch, j, epoch)),
np.array(
denormalize(
fixed_X[j]).cpu().detach()).reshape(
256, 256, 3))
else:
cv2.imwrite(
os.path.join(
'../cgresults/r-9-pred_masks',
'mask_{}_{}_{}.png'.format(batch, j, epoch)),
np.array(denormalize(
depth_map[j]).cpu().detach()).reshape(
256, 256, 3))
if epoch == 9:
cv2.imwrite(
os.path.join(
'../cgresults/r-9-inputs',
'input_{}_{}_{}.png'.format(
batch, j, epoch)),
np.array(
denormalize(
fixed_X[j]).cpu().detach()).reshape(
256, 256, 3))
# torch.set_grad_enabled(True)
print("Time to finish epoch ", time.time() - since)
torch.save(model, '../CGmodel/best_model5.pt')
loss_Gl.append(float(model.loss_G))
loss_DXl.append(float(model.loss_D_X))
loss_DYl.append(float(model.loss_D_Y))
with open('../CGloss/lossG5.pk', 'wb') as f:
pickle.dump(loss_Gl, f)
with open('../CGloss/lossD_X5.pk', 'wb') as f:
pickle.dump(loss_DXl, f)
with open('../CGloss/lossd_Y5.pk', 'wb') as f:
pickle.dump(loss_DYl, f)
elif model_name == 'P2P': # This is for Khem's version of Pix2Pix
model = Pix2Pix(p2p_opt)
print_freq = 10
train_iter = iter(p2p_train_loader)
val_iter = iter(p2p_val_loader)
fixed_X, fixed_Y = val_iter.next()
fixed_X = normalization(fixed_X).to(device)
fixed_Y = normalization(fixed_Y).to(device)
loss_Gl = []
loss_Dl = []
num_batches = len(train_iter)
for epoch in range(3000):
if epoch == 299:
model.change_lr(model.opt.lr / 2)
if epoch == 499:
model.change_lr(model.opt.lr / 2)
since = time.time()
print("Epoch ", epoch, " entering ")
train_iter = iter(p2p_train_loader)
for batch in range(num_batches):
print("Epoch ", epoch, "Batch ", batch,
" running with learning rate ", model.opt.lr)
inputX, inputY = train_iter.next()
inputX = normalization(inputX).to(device)
inputY = normalization(inputY).to(device)
model.get_input(inputX, inputY)
model.optimize()
# print("Dx Loss : {:.6f} Dy Loss: {:.6f} Generator Loss: {:.6f} ".format(model.dx_loss, model.dy_loss, model.gen_loss))
print("Model D Loss ", float(model.loss_D), "Model G loss",
float(model.loss_G))
if (epoch + 1) % 10 == 0:
# torch.set_grad_enabled(False)
depth_map = model.G.forward(fixed_X)
for j in range(depth_map.size()[0]):
cv2.imwrite(
os.path.join(
'../p2presults/pred_masks',
'mask_{}_{}_{}.png'.format(batch, j, epoch)),
np.array(denormalize(
depth_map[j]).cpu().detach()).reshape(256, 256, 3))
if epoch == 9:
cv2.imwrite(
os.path.join(
'../p2presults/inputs',
'input_{}_{}_{}.png'.format(batch, j, epoch)),
np.array(denormalize(
fixed_X[j]).cpu().detach()).reshape(
256, 256, 3))
cv2.imwrite(
os.path.join(
'../p2presults/inputs',
'ground_depth_{}_{}_{}.png'.format(
batch, j, epoch)),
np.array(denormalize(
fixed_Y[j]).cpu().detach()).reshape(
256, 256, 3))
# torch.set_grad_enabled(True)
print("Time to finish epoch ", time.time() - since)
torch.save(model, '../P2Pmodel/best_model8.pt')
loss_Gl.append(float(model.loss_G))
loss_Dl.append(float(model.loss_D))
with open('../P2Ploss/lossG8.pk', 'wb') as f:
pickle.dump(loss_Gl, f)
with open('../P2Ploss/lossD8.pk', 'wb') as f:
pickle.dump(loss_Dl, f)
import os
from pix2pix import Pix2Pix
if __name__ == "__main__":
DATA_PATH = './data/npy'
NUM_EPOCH = 1000
pix2pix = Pix2Pix(img_dim=(128, 256, 3), batch_size=16, patch_size=64)
print('Training.....')
pix2pix.train(num_epoch=NUM_EPOCH, data_path=DATA_PATH)
overlapping = cv2.addWeighted(bottom, alpha, img2, beta, gamma)
cv2.imwrite(save_pth, overlapping)
def predict_single_image(pix2pix,wdsr, image_path, save_path):
pix2pix.generator.load_weights('./weights/generator_weights.h5')
wdsr.load_weights('./weights/wdsr-b-32-x4.h5')
image_B = imageio.imread(image_path, pilmode='RGB').astype(np.float)
image_B = skimage.transform.resize(image_B, (pix2pix.nW, pix2pix.nH))
images_B = []
images_B.append(image_B)
images_B = np.array(images_B)/127.5 - 1.
generates_A = pix2pix.generator.predict(images_B)
generate_A = generates_A[0]
generate_A = np.uint8((np.array(generate_A) * 0.5 + 0.5) * 255)
generate_A = Image.fromarray(generate_A)
generated_image = Image.new('RGB', (pix2pix.nW, pix2pix.nH))
generated_image.paste(generate_A, (0, 0, pix2pix.nW, pix2pix.nH))
lr = np.asarray(generated_image)
x = np.array([lr])
y = wdsr.predict(x)
y = np.clip(y, 0, 255)
y = y.astype('uint8')
sr = Image.fromarray(y[0])
sr.save(save_path)
combine(image_path,save_path,0.5,0.5,0,save_path)
pass
gan = Pix2Pix()
wdsr = wdsr_b(scale=4, num_res_blocks=32)
predict_single_image(gan,wdsr, '1.jpg', 'test_1.jpg')
parser.add_argument('--test_folder', nargs='?', default='val', type=str)
parser.add_argument('--in_channels', nargs='?', default=3, type=int)
parser.add_argument('--batch_size', nargs='?', default=16, type=int)
parser.add_argument('--gen_filters', nargs='?', default=1024, type=int)
parser.add_argument('--disc_filters', nargs='?', default=512, type=int)
parser.add_argument('--img_output_size', nargs='?', default=256, type=int)
parser.add_argument('--gen_layers', nargs='?', default=6, type=int)
parser.add_argument('--disc_layers', nargs='?', default=4, type=int)
parser.add_argument('--test_perc', nargs='?', default=1, type=float)
parser.add_argument('--lr_disc', nargs='?', default=1e-4, type=float)
parser.add_argument('--lr_gen', nargs='?', default=1e-4, type=float)
parser.add_argument('--lr_cycle_mult', nargs='?', default=2, type=int)
parser.add_argument('--beta1', nargs='?', default=.5, type=float)
parser.add_argument('--beta2', nargs='?', default=.999, type=float)
parser.add_argument('--alpha', nargs='?', default=10, type=int)
parser.add_argument('--train_epoch', nargs='?', default=6, type=int)
parser.add_argument('--ids', type=int, nargs='+', default=[1, 53, 62])
parser.add_argument('--save_root', nargs='?', default='shoes', type=str)
parser.add_argument('--load_state', nargs='?', type=str)
params = vars(parser.parse_args())
# if load_state arg is not used, then train model from scratch
if __name__ == '__main__':
p2p = Pix2Pix(params)
if params['load_state']:
p2p.load_state(params['load_state'])
else:
print('Starting From Scratch')
p2p.train()
from pix2pix import Pix2Pix
from os import path
import sys
if __name__ == "__main__":
print("Initializing and loading model...")
gan = Pix2Pix("img_data")
print("Starting training ...")
# e is number of epochs. It should be increased to ~15 to 20 at least.
e = int(sys.argv[1]) if len(sys.argv) > 1 else 2
gan.train(epochs=e, batch_size=1, include_val=False, step_print=100)
folder = "saved_models"
print("Training complete. Saving ...")
gan.generator.save(path.join(folder, "gen_model_" + str(e) + ".h5"))
gan.discriminator.save(path.join(folder, "disc_model_" + str(e) + ".h5"))
gan.combined.save(path.join(folder, "gan_model_" + str(e) + ".h5"))
print("Models saved.")
Author: skconan
Date created: 2019/04/13
Python Version: 3.6
"""
import cv2 as cv
import numpy as np
import tensorflow as tf
# from bg_subtraction import bg_subtraction
from pix2pix import Pix2Pix
from utilities import *
from operator import itemgetter
import os
model_file = "./pix2pix_checkpoints_addbg_ver2/200-ckpt-20"
pix2pix = Pix2Pix(postfix='', train=False)
pix2pix.checkpoint.restore(model_file)
def segmentation_pix2pix(img):
with tf.device('/device:GPU:0'):
# with tf.device('/device:cpu:0'):
img = tf.convert_to_tensor(np.float32(img))
img = resize(img, 256, 256)
# img = normalize(img)
img, result = pix2pix.predict_images(img)
# hsv = cv.cvtColor(result, cv.COLOR_BGR2HSV)
# cv.imshow("result_tf_color",result.copy())
# lower = np.array([22, 200, 200], np.uint8)
# upper = np.array([38, 255, 255], np.uint8)
# mask = cv.inRange(hsv, lower, upper)
def __init__(self):
self.mode = 0 # 0 = paint, 1 = erase
self.x = 0
self.y = 0
self.root = tk.Tk()
# Declare Variables Used in GUI
self.size = tk.IntVar() # Paint Brush Size
self.size.set(5)
self.selectedModel = tk.StringVar() # What to generate in our GAN
self.selectedModel.set('trees')
self.selectedGAN = tk.StringVar() # What GAN to use for real time display
self.selectedGAN.set('cyclegan')
self.canvasTitle = tk.StringVar() # Title for drawing area
self.canvasTitle.set('Drawing (Trees)')
self.root.title("Pix2Pix GAN GUI")
self.root.resizable(width=False, height=False)
self.root.geometry("1075x645+436+218") # 1050x590 window centered on 1920x1080 displays
# Next 3 lines to bring window to front on launch, but allow other apps in front afterwards
self.root.lift()
self.root.attributes('-topmost', True)
self.root.after_idle(self.root.attributes, '-topmost', False)
# Create all needed buttons
self.drawButton = ttk.Button(self.root, text="Draw", command=self.setDraw).grid(row=0, column=0)
self.eraseButton = ttk.Button(self.root, text="Erase", command=self.setErase).grid(row=0, column=1)
self.clearAllButton = ttk.Button(self.root, text="Clear All", command=self.clearAll).grid(row=1, column=0)
self.saveButton = ttk.Button(self.root, text="Save", command=self.save).grid(row=1, column=1)
# Create Model Selector above image viewer
self.modelSelectorLabel = ttk.Label(self.root, text='Select What to Generate:').grid(row=0, column=3)
self.treeButton = ttk.Radiobutton(self.root, text="Trees", variable=self.selectedModel, value='trees', command=self.selectModel)
self.treeButton.grid(row=0, column=4)
self.quickDraw = ttk.Radiobutton(self.root, text="QuickDraw", variable=self.selectedModel, value='quickdraw_trees', command=self.selectModel)
self.quickDraw.grid(row=0, column=5)
self.pizzaButton = ttk.Radiobutton(self.root, text="Pizza", variable=self.selectedModel, value='pizza', command=self.selectModel)
self.pizzaButton.grid(row=0, column=6)
self.appleButton = ttk.Radiobutton(self.root, text="Apples", variable=self.selectedModel, value='apples', command=self.selectModel)
self.appleButton.grid(row=0, column=7)
# Create Paint Size slider
self.scaleLabel = tk.Label(self.root, text="Size").grid(row=1, column=2)
self.scale = tk.Scale(self.root, variable=self.size, from_=1, to=30, cursor='target', orient=tk.HORIZONTAL).grid(row=0, column=2)
# Create canvas to draw on, but don't pack it until adding event bindings
self.canvasLabel = tk.Label(self.root, textvariable=self.canvasTitle, font=("Arial", 16)).grid(row=2, column=0, columnspan=3, pady=(25, 0))
self.canvas = tk.Canvas(self.root, cursor='dot', height=512, width=512, bg='white', highlightthickness=3, highlightbackground='grey')
self.canvas.bind('<ButtonPress-1>', self.setStartPoint) # Gets initial left-click
self.canvas.bind('<B1-Motion>', self.draw) # Gets motion from mouse while left-click is held
self.canvas.bind('<ButtonRelease-1>', self.updateImageViewer) # Update image in real time on left-click release
self.canvas.grid(row=3, column=0, columnspan=3, padx=5, pady=3)
# Create Canvas to view images produced by GAN
self.imageViewerLabel = tk.Label(self.root, text="Generated Image (Pix2Pix)", font=("Arial", 16)).grid(row=2, column=3, columnspan=2, pady=(25, 0))
self.imageViewer = tk.Canvas(self.root, cursor='arrow', height=256, width=256, bg='white', highlightthickness=3, highlightbackground='grey')
self.imageViewer.grid(row=3, column=3, columnspan=2, sticky=tk.NW, pady=3, padx=5)
self.imageViewerLabel2 = tk.Label(self.root, text="Generated Image (CycleGAN)", font=("Arial", 16)).grid(row=2, column=5, columnspan=3, pady=(25, 0))
self.imageViewer2 = tk.Canvas(self.root, cursor='arrow', height=256, width=256, bg='white', highlightthickness=3, highlightbackground='grey')
self.imageViewer2.grid(row=3, column=5, columnspan=3, sticky=tk.NW, pady=3)
# Initialize models
self.cycle_gan_models = {}
for name in ['trees', 'pizza', 'apples', 'quickdraw_trees']:
self.cycle_gan_models[name] = CycleGan('cyclegan_%s'%name, epoch='latest')
self.pix2pix_models = {}
for name, epoch in zip(['trees', 'pizza', 'apples'], ['latest', 'latest', 'latest']):
self.pix2pix_models[name] = Pix2Pix('pix2pix_%s'%name,epoch=epoch)
self.pix2pix_models['quickdraw_trees'] = self.pix2pix_models['trees']
self.root.mainloop()
self.results_dir = "/data/kdabi/CS698O/Autopainter/CS698-cartoon-painter/saved_models"
opt = Options()
# opt = TrainOptions().parse()
data_loader = CustomDatasetDataLoader()
data_loader.initialize(opt)
dataset = data_loader.load_data()
dataset_size = len(data_loader)
# print('#training images = %d' % dataset_size)
model = Pix2Pix(opt)
visualizer = Visualizer(opt)
total_steps = 0
for epoch in range(opt.epoch_count, opt.niter + opt.niter_decay + 1):
epoch_start_time = time.time()
epoch_iter = 0
for i, data in enumerate(dataset):
iter_start_time = time.time()
total_steps += opt.batchSize
epoch_iter += opt.batchSize
model.set_input(data)
model.optimize_parameters()
if total_steps % opt.display_freq == 0: