def GeneratorErrorMap(img_dim_tuple,
batch_size=10,
modelEpoch=130,
threshold=0.1,
f=None,
indexes=1448):
if (f == None):
path = "../../data/nyu_depth_v2_labeled.mat"
f = h5py.File(path)
gen, length = fDataSet_generator(img_dim, batch_size, f=f)
model = generator_unet_upsampling(img_dim, bn_mode, batch_size)
model.load_weights('../../models/CNN/pix2depthgen_weights_epoch%d.h5' %
modelEpoch) #gen_weights_epoch45.h5
def GenFunction():
while True:
img, dep = gen.__next__()
dmap = model.predict(img)
# errorMap=(np.divide(np.fabs(dmap-dep), dep)<threshold)
# errorMap=(np.divide(np.fabs(dmap-dep), dep)<threshold).astype(int)
# errorMap=(div0(np.fabs(dmap-dep), dep)<threshold).astype(int)
errorMap = (np.fabs(dmap - dep) < threshold).astype(int)
yield (img, errorMap)
return (GenFunction(), 1448) #todo: accurate Value
def GeneratorErrorMap(img_dim_tuple,batch_size=10, modelEpoch=130, threshold=0.1,f=None,indexes=1448):
if(f==None):
path = "../../data/nyu_depth_v2_labeled.mat"
f = h5py.File(path)
gen , length = fDataSet_generator(img_dim, batch_size,f=f)
model = generator_unet_upsampling(img_dim, bn_mode, batch_size)
model.load_weights('../../models/CNN/pix2depthgen_weights_epoch%d.h5' % modelEpoch) #gen_weights_epoch45.h5
def GenFunction():
while True:
img, dep = gen.__next__()
dmap = model.predict(img)
# errorMap=(np.divide(np.fabs(dmap-dep), dep)<threshold)
# errorMap=(np.divide(np.fabs(dmap-dep), dep)<threshold).astype(int)
# errorMap=(div0(np.fabs(dmap-dep), dep)<threshold).astype(int)
errorMap=(np.fabs(dmap-dep)<threshold).astype(int)
yield (img,errorMap)
return(GenFunction(),1448) #todo: accurate Value
img = f['images'][i]
depth = f['depths'][i]
img_ = np.empty([img_dim, img_dim, 3])
img_[:, :, 0] = cv2.resize(img[0, :, :].T, (img_dim, img_dim))
img_[:, :, 1] = cv2.resize(img[1, :, :].T, (img_dim, img_dim))
img_[:, :, 2] = cv2.resize(img[2, :, :].T, (img_dim, img_dim))
depth_ = np.empty([img_dim, img_dim, 3])
depth_[:, :, 0] = cv2.resize(depth[:, :].T, (img_dim, img_dim))
depth_[:, :, 1] = cv2.resize(depth[:, :].T, (img_dim, img_dim))
depth_[:, :, 2] = cv2.resize(depth[:, :].T, (img_dim, img_dim))
img_ = img_ / 255.0
depth_ = depth_ / 10.0
img_batch[index] = img_
depth_batch[index] = depth_
yield img_batch, depth_batch
gen = facades_generator(img_dim, batch_size=1)
img, dep = next(gen)
cv2.imwrite("real.jpg", img[0] * 255)
for i in range(0, 200, 5):
model = generator_unet_upsampling(img_dim, bn_mode, batch_size)
model.load_weights('../../models/CNN/gen_weights_epoch%d.h5' %
i) #gen_weights_epoch45.h5
dmap = model.predict(dep)[0]
print((dep[0][0][0]))
print((dep.shape))
print((dmap[0][0][0]))
cv2.imwrite("test_%d.jpg" % i, np.hstack((dmap * 255, dep[0] * 255)))
def main(dataset, batch_size, patch_size, epochs, label_smoothing,
label_flipping):
print(project_dir)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True # dynamically grow the memory used on the GPU
sess = tf.Session(config=config)
K.tensorflow_backend.set_session(
sess) # set this TensorFlow session as the default session for Keras
image_data_format = "channels_first"
K.set_image_data_format(image_data_format)
save_images_every_n_batches = 30
save_model_every_n_epochs = 0
# configuration parameters
print("Config params:")
print(" dataset = {}".format(dataset))
print(" batch_size = {}".format(batch_size))
print(" patch_size = {}".format(patch_size))
print(" epochs = {}".format(epochs))
print(" label_smoothing = {}".format(label_smoothing))
print(" label_flipping = {}".format(label_flipping))
print(" save_images_every_n_batches = {}".format(
save_images_every_n_batches))
print(" save_model_every_n_epochs = {}".format(save_model_every_n_epochs))
model_name = datetime.strftime(datetime.now(), '%y%m%d-%H%M')
model_dir = os.path.join(project_dir, "models", model_name)
fig_dir = os.path.join(project_dir, "reports", "figures")
logs_dir = os.path.join(project_dir, "reports", "logs", model_name)
os.makedirs(model_dir)
# Load and rescale data
ds_train_gen = data_utils.DataGenerator(file_path=dataset,
dataset_type="train",
batch_size=batch_size)
ds_train_disc = data_utils.DataGenerator(file_path=dataset,
dataset_type="train",
batch_size=batch_size)
ds_val = data_utils.DataGenerator(file_path=dataset,
dataset_type="val",
batch_size=batch_size)
enq_train_gen = OrderedEnqueuer(ds_train_gen,
use_multiprocessing=True,
shuffle=True)
enq_train_disc = OrderedEnqueuer(ds_train_disc,
use_multiprocessing=True,
shuffle=True)
enq_val = OrderedEnqueuer(ds_val, use_multiprocessing=True, shuffle=False)
img_dim = ds_train_gen[0][0].shape[-3:]
n_batch_per_epoch = len(ds_train_gen)
epoch_size = n_batch_per_epoch * batch_size
print("Derived params:")
print(" n_batch_per_epoch = {}".format(n_batch_per_epoch))
print(" epoch_size = {}".format(epoch_size))
print(" n_batches_val = {}".format(len(ds_val)))
# Get the number of non overlapping patch and the size of input image to the discriminator
nb_patch, img_dim_disc = data_utils.get_nb_patch(img_dim, patch_size)
tensorboard = TensorBoard(log_dir=logs_dir,
histogram_freq=0,
batch_size=batch_size,
write_graph=True,
write_grads=True,
update_freq='batch')
try:
# Create optimizers
opt_dcgan = Adam(lr=1E-3, beta_1=0.9, beta_2=0.999, epsilon=1e-08)
# opt_discriminator = SGD(lr=1E-3, momentum=0.9, nesterov=True)
opt_discriminator = Adam(lr=1E-3,
beta_1=0.9,
beta_2=0.999,
epsilon=1e-08)
# Load generator model
generator_model = models.generator_unet_upsampling(img_dim)
generator_model.summary()
plot_model(generator_model,
to_file=os.path.join(fig_dir, "generator_model.png"),
show_shapes=True,
show_layer_names=True)
# Load discriminator model
# TODO: modify disc to accept real input as well
discriminator_model = models.DCGAN_discriminator(
img_dim_disc, nb_patch)
discriminator_model.summary()
plot_model(discriminator_model,
to_file=os.path.join(fig_dir, "discriminator_model.png"),
show_shapes=True,
show_layer_names=True)
# TODO: pretty sure this is unnecessary
generator_model.compile(loss='mae', optimizer=opt_discriminator)
discriminator_model.trainable = False
DCGAN_model = models.DCGAN(generator_model, discriminator_model,
img_dim, patch_size, image_data_format)
# L1 loss applies to generated image, cross entropy applies to predicted label
loss = [models.l1_loss, 'binary_crossentropy']
loss_weights = [1E1, 1]
DCGAN_model.compile(loss=loss,
loss_weights=loss_weights,
optimizer=opt_dcgan)
discriminator_model.trainable = True
discriminator_model.compile(loss='binary_crossentropy',
optimizer=opt_discriminator)
tensorboard.set_model(DCGAN_model)
# Start training
enq_train_gen.start(workers=1, max_queue_size=20)
enq_train_disc.start(workers=1, max_queue_size=20)
enq_val.start(workers=1, max_queue_size=20)
out_train_gen = enq_train_gen.get()
out_train_disc = enq_train_disc.get()
out_val = enq_val.get()
print("Start training")
for e in range(1, epochs + 1):
# Initialize progbar and batch counter
progbar = generic_utils.Progbar(epoch_size)
start = time.time()
for batch_counter in range(1, n_batch_per_epoch + 1):
X_transformed_batch, X_orig_batch = next(out_train_disc)
# Create a batch to feed the discriminator model
X_disc, y_disc = data_utils.get_disc_batch(
X_transformed_batch,
X_orig_batch,
generator_model,
batch_counter,
patch_size,
label_smoothing=label_smoothing,
label_flipping=label_flipping)
# Update the discriminator
disc_loss = discriminator_model.train_on_batch(X_disc, y_disc)
# Create a batch to feed the generator model
X_gen_target, X_gen = next(out_train_gen)
y_gen = np.zeros((X_gen.shape[0], 2), dtype=np.uint8)
# Set labels to 1 (real) to maximize the discriminator loss
y_gen[:, 1] = 1
# Freeze the discriminator
discriminator_model.trainable = False
gen_loss = DCGAN_model.train_on_batch(X_gen,
[X_gen_target, y_gen])
# Unfreeze the discriminator
discriminator_model.trainable = True
metrics = [("D logloss", disc_loss), ("G tot", gen_loss[0]),
("G L1", gen_loss[1]), ("G logloss", gen_loss[2])]
progbar.add(batch_size, values=metrics)
logs = {k: v for (k, v) in metrics}
logs["size"] = batch_size
tensorboard.on_batch_end(batch_counter, logs=logs)
# Save images for visualization
if batch_counter % save_images_every_n_batches == 0:
# Get new images from validation
data_utils.plot_generated_batch(
X_transformed_batch, X_orig_batch, generator_model,
os.path.join(logs_dir, "current_batch_training.png"))
X_transformed_batch, X_orig_batch = next(out_val)
data_utils.plot_generated_batch(
X_transformed_batch, X_orig_batch, generator_model,
os.path.join(logs_dir, "current_batch_validation.png"))
print("")
print('Epoch %s/%s, Time: %s' % (e, epochs, time.time() - start))
tensorboard.on_epoch_end(e, logs=logs)
if (save_model_every_n_epochs >= 1 and e % save_model_every_n_epochs == 0) or \
(e == epochs):
print("Saving model for epoch {}...".format(e), end="")
sys.stdout.flush()
gen_weights_path = os.path.join(
model_dir, 'gen_weights_epoch{:03d}.h5'.format(e))
generator_model.save_weights(gen_weights_path, overwrite=True)
disc_weights_path = os.path.join(
model_dir, 'disc_weights_epoch{:03d}.h5'.format(e))
discriminator_model.save_weights(disc_weights_path,
overwrite=True)
DCGAN_weights_path = os.path.join(
model_dir, 'DCGAN_weights_epoch{:03d}.h5'.format(e))
DCGAN_model.save_weights(DCGAN_weights_path, overwrite=True)
print("done")
except KeyboardInterrupt:
pass
enq_train_gen.stop()
enq_train_disc.stop()
enq_val.stop()
def train(_run, train_cnf, generator_cnf, discriminator_cnf):
generator = train_cnf["generator"]
discriminator = train_cnf["discriminator"]
batch_size = train_cnf["batch_size"]
nb_epoch = train_cnf["nb_epoch"]
learning_rate = train_cnf["learning_rate"]
beta_1 = train_cnf["beta_1"]
beta_2 = train_cnf["beta_2"]
l1_weight = train_cnf["l1_weight"]
dset = train_cnf["dset"]
base_dir = train_cnf["base_dir"]
model_save_epoch = train_cnf["model_save_epoch"]
fig_save_epoch = train_cnf["fig_save_epoch"]
reverse_dset = train_cnf["reverse_dset"]
# Setup environment (logging directory etc)
# Just creates the figures and models dirs
experiment_name = 'busgan/exp' + str(_run._id)
experiment_dir = os.path.join(base_dir, 'experiments', experiment_name)
general_utils.setup_logging(experiment_dir)
# Load and rescale data
X_full_train, X_sketch_train, X_full_val, X_sketch_val = data_utils.load_data(dset, base_dir, reverse_dset)
img_dim = X_full_train.shape[-3:]
if generator == 'generator_unet_upsampling':
generator_model = models.generator_unet_upsampling(img_dim=img_dim,
**generator_cnf)
elif generator == 'generator_upsampling':
generator_model = models.generator_upsampling(img_dim=img_dim,
**generator_cnf)
else:
raise ValueError('Generator name not recognised')
if discriminator == 'DCGAN_discriminator':
discriminator_model = models.DCGAN_discriminator(img_dim=img_dim,
**discriminator_cnf)
else:
raise ValueError('Discriminator name not recognised')
batches_per_e = X_full_train.shape[0] // batch_size
try:
# Create optimizers
opt_dcgan = Adam(lr=learning_rate,
beta_1=beta_1, beta_2=beta_2, epsilon=1e-08)
opt_discriminator = Adam(lr=learning_rate,
beta_1=beta_1, beta_2=beta_2,
epsilon=1e-08)
generator_model.compile(loss='mae', optimizer=opt_discriminator)
discriminator_model.trainable = False
GAN_model = models.GAN(generator_model, discriminator_model, img_dim)
loss = [l1_loss, 'binary_crossentropy']
loss_weights = [l1_weight, 1]
GAN_model.compile(loss=loss, loss_weights=loss_weights,
optimizer=opt_dcgan)
discriminator_model.trainable = True
discriminator_model.compile(loss='binary_crossentropy',
optimizer=opt_discriminator)
# Start training
for e in range(1, nb_epoch + 1):
print(f'Epoch {e}/{nb_epoch}')
batch_counter = 1
for X_full_batch, X_sketch_batch in data_utils.gen_batch(X_full_train, X_sketch_train, batch_size):
print(f'Batch {batch_counter}/{batches_per_e}')
# Create a batch to feed the discriminator model
X_disc, y_disc = data_utils.get_disc_batch(X_full_batch,
X_sketch_batch,
generator_model,
batch_counter)
# Update the discriminator
disc_loss = discriminator_model.train_on_batch(X_disc, y_disc)
# Create a batch to feed the generator model
X_gen_target, X_gen = X_full_batch, X_sketch_batch
y_gen = np.ones((X_gen.shape[0], 1), dtype=np.uint8)
# Freeze the discriminator
discriminator_model.trainable = False
gen_loss = GAN_model.train_on_batch(X_gen,
[X_gen_target, y_gen])
# Unfreeze the discriminator
discriminator_model.trainable = True
# In the last iteration and if epoch is multiple of image
# saving interval
if batch_counter == batches_per_e and e % fig_save_epoch == 0:
# Save figures
figures_dir = os.path.join(experiment_dir, 'figures')
# At the end of each epoch save images for visualization
# Get new images from validation
data_utils.plot_generated_batch(X_full_batch,
X_sketch_batch,
generator_model,
"training",
figures_dir, e)
idx = np.random.choice(X_full_val.shape[0], batch_size,
replace=False)
data_utils.plot_generated_batch(X_full_val[idx],
X_sketch_val[idx],
generator_model,
"validation",
figures_dir, e)
# In the last batch also log metrics
if batch_counter == batches_per_e:
_run.log_scalar('d.logloss', disc_loss, e)
_run.log_scalar('g.tot', gen_loss[0], e)
_run.log_scalar('g.l1', gen_loss[1], e)
_run.log_scalar('g.logloss', gen_loss[2], e)
batch_counter += 1
if e % model_save_epoch == 0:
models_dir = os.path.join(experiment_dir, 'models')
gen_weights_path = os.path.join(models_dir,
'gen_weights_epoch%s.h5' % e)
generator_model.save_weights(gen_weights_path, overwrite=True)
disc_weights_path = os.path.join(models_dir,
'disc_weights_epoch%s.h5' % e)
discriminator_model.save_weights(disc_weights_path,
overwrite=True)
GAN_weights_path = os.path.join(models_dir,
'GAN_weights_epoch%s.h5' % e)
GAN_model.save_weights(GAN_weights_path, overwrite=True)
# Clear GPU + RAM
K.clear_session()
del GAN_model
del generator_model
del discriminator_model
except KeyboardInterrupt:
pass
def main(model,
input,
truth,
output,
image_size,
concat,
batch_size,
dataset):
K.set_image_data_format("channels_first")
# Load the model
img_dim = (3,) + image_size
generator_model = models.generator_unet_upsampling(img_dim)
generator_model.load_weights(str(model))
# Setup the data generator
if h5py.is_hdf5(input):
generator = data_utils.DataGenerator(file_path=input,
dataset_type=dataset,
batch_size=batch_size)
batch_gen = get_batch_from_hdf5(generator)
count = len(generator) * batch_size
elif os.path.isdir(input):
# Directory of images
input_paths = [str(img) for img in sorted(Path(input).glob('**/*.jpg'))]
if truth is not None:
truth_paths = [str(img) for img in sorted(Path(truth).glob('**/*.jpg'))]
else:
truth_paths = []
batch_gen = get_batch_from_images(input_paths, truth_paths, batch_size, image_size)
count = len(input_paths)
else:
# Single image file
input_paths = [input]
if truth is not None:
truth_paths = [truth]
else:
truth_paths = []
batch_size = 1
batch_gen = get_batch_from_images(input_paths, truth_paths, batch_size, image_size)
count = 1
# Setup the output
if output is not None:
# If input is multiple images, create directory for output images
if count > 1:
if os.path.exists(output):
raise IOError('Output directory already exists')
os.makedirs(output)
# Inference
with tqdm(None, total=count) as progress_bar:
for batch_record in batch_gen:
input_batch = batch_record.input
truth_batch = batch_record.truth
image_names = batch_record.image_names
out_batch = generator_model.predict(input_batch)
progress_bar.update(input_batch.shape[0])
if output is not None:
# Individually save each image in the batch
for i in range(len(image_names)):
input_b = input_batch[i:i+1]
output_b = out_batch[i:i+1]
truth_b = truth_batch[i:i+1] if truth_batch is not None else None
image_name = image_names[i]
out_image = generate_output_image(input_b, output_b, truth_b, concat)
if count > 1:
out_path = os.path.join(str(output), image_name)
else:
out_path = output
plt.imsave(out_path, out_image)
else:
# Show the image
out_image = generate_output_image(input_batch, out_batch, truth_batch, concat)
plt.figure()
plt.imshow(out_image)
plt.show()
plt.clf()
plt.close()
img = f['images'][i]
depth = f['depths'][i]
img_=np.empty([img_dim,img_dim,3])
img_[:,:,0] = cv2.resize(img[0,:,:].T,(img_dim,img_dim))
img_[:,:,1] = cv2.resize(img[1,:,:].T,(img_dim,img_dim))
img_[:,:,2] = cv2.resize(img[2,:,:].T,(img_dim,img_dim))
depth_ = np.empty([img_dim, img_dim, 3])
depth_[:,:,0] = cv2.resize(depth[:,:].T,(img_dim,img_dim))
depth_[:,:,1] = cv2.resize(depth[:,:].T,(img_dim,img_dim))
depth_[:,:,2] = cv2.resize(depth[:,:].T,(img_dim,img_dim))
img_ = img_/255.0
depth_ = depth_/10.0
img_batch[index] = img_
depth_batch[index] = depth_
yield (img_batch,depth_batch)
return (innerGenerator(),1448) #TODO: the corect size
if __name__ == "__main__":
gen, length =fDataSet_generator(img_dim, batch_size=1)
img, dep = gen.__next__()
cv2.imwrite("real.jpg",img[0]*255)
i=130
model = generator_unet_upsampling(img_dim, bn_mode, batch_size)
model.load_weights('../../models/CNN/pix2depthgen_weights_epoch%d.h5' % i) #gen_weights_epoch45.h5
dmap = model.predict(img)[0]
print (dep[0][1][0])
print (dep.shape)
print (dmap[1][0])
cv2.imwrite("test_%d.jpg" % i,np.hstack((dmap*255,dep[0]*255)))
