def main():
#true = cv2.imread(os.path.join("..", "DSIDS", "HR", "niklas_animal_0010.jpg"))
true = cv2.cvtColor(
cv2.imread(os.path.join("..", "DSIDS", "test", "20190702_210258.jpg")),
cv2.COLOR_BGR2RGB)
lr = cv2.resize(true, (0, 0),
fx=1 / 4,
fy=1 / 4,
interpolation=cv2.INTER_CUBIC)
pred_nearest = cv2.resize(lr, (0, 0),
fx=4,
fy=4,
interpolation=cv2.INTER_CUBIC)
model = load_model(os.path.join("..", "models",
"initialized_gen_model20.h5"),
custom_objects={"tf": tf})
_in = Input(shape=lr.shape)
_out = model(_in)
_model = Model(_in, _out)
sr = Utils.denormalize(
np.squeeze(_model.predict(
np.expand_dims(rescale_imgs_to_neg1_1(lr), axis=0)),
axis=0))
print("Nearest Neighbor:")
print("MSE: ", MSE(true, pred_nearest))
print("PSNR: ", PSNR(true, pred_nearest))
print("SSIM: ", SSIM(true, pred_nearest))
print("MSSIM: ", MSSIM(true, pred_nearest))
print("Init SRGAN:")
print("MSE: ", MSE(true, sr))
print("PSNR: ", PSNR(true, sr))
print("SSIM: ", SSIM(true, sr))
print("MSSIM: ", MSSIM(true, sr))
print("Black:")
print("MSE: ", MSE(true, np.zeros(true.shape)))
print("PSNR: ", PSNR(true, np.zeros(true.shape)))
print("SSIM: ", SSIM(true, np.zeros(true.shape)))
print("MSSIM: ", MSSIM(true, np.zeros(true.shape)))
print("HR, HR:")
print("MSE: ", MSE(true, true))
print("PSNR: ", PSNR(true, true))
print("SSIM: ", SSIM(true, true))
print("MSSIM: ", MSSIM(true, true))
cv2.imwrite("nearest.jpg", cv2.cvtColor(pred_nearest, cv2.COLOR_RGB2BGR))
cv2.imwrite("sr.jpg", cv2.cvtColor(sr, cv2.COLOR_RGB2BGR))
def main():
# paths to the models
model_paths = [
os.path.join("..", "models", "SRDense-Type-3_ep80.h5"),
os.path.join("..", "models", "srdense-norm.h5"),
os.path.join("..", "models", "srresnet85.h5"),
os.path.join("..", "models", "gen_model90.h5"),
os.path.join("..", "models", "srgan60.h5"),
os.path.join("..", "models", "srgan-mse-20.h5"), "Nearest"
]
# corresponding names of the models
model_names = [
"SRDense", "SRDense-norm", "SRResNet", "SRGAN-from-scratch",
"SRGAN-percept.-loss", "SRGAN-mse", "NearestNeighbor"
]
# corresponding tile shapes
tile_shapes = [((168, 168), (42, 42)), ((168, 168), (42, 42)),
((504, 504), (126, 126)), ((336, 336), (84, 84)),
((504, 504), (126, 126)), ((504, 504), (126, 126)),
((336, 336), (84, 84))]
# used to load the models with custom loss functions
loss = VGG_LOSS((504, 504, 3))
custom_objects = [{}, {
"tf": tf
}, {
"tf": tf
}, {
"tf": tf
}, {
"tf": tf
}, {
"tf": tf
}, {}]
# creating a list of test images
# [(lr, hr)]
DOWN_SCALING_FACTOR = 4
INTERPOLATION = cv2.INTER_CUBIC
test_images = []
root = os.path.join("..", "DSIDS", "test")
# iterating over all files in the test folder
for img in os.listdir(root):
# chekcing if the file is an image
if not ".jpg" in img:
continue
hr = Utils.crop_into_lr_shape(cv2.cvtColor(
cv2.imread(os.path.join(root, img), cv2.IMREAD_COLOR),
cv2.COLOR_BGR2RGB),
shape=(3024, 4032))
lr = cv2.resize(hr, (0, 0),
fx=1 / DOWN_SCALING_FACTOR,
fy=1 / DOWN_SCALING_FACTOR,
interpolation=INTERPOLATION)
test_images.append((lr, hr))
if TILES:
'''
First calculating performance metrics on single image tiles
'''
tile_performance = {}
for i, mp in tqdm(enumerate(model_paths)):
keras.backend.clear_session()
# first step: load the model
if i < 6:
model = load_model(mp, custom_objects=custom_objects[i])
mse = []
psnr = []
ssim = []
mssim = []
# second step: iterate over the test images
for test_pair in tqdm(test_images):
# third step: tile the test image
lr_tiles = Utils.tile_image(test_pair[0],
shape=tile_shapes[i][1])
hr_tiles = Utils.tile_image(test_pair[1],
shape=tile_shapes[i][0])
m = []
p = []
s = []
ms = []
# fourth step: iterate over the tiles
for lr, hr in zip(lr_tiles, hr_tiles):
# fifth step: calculate the sr tile
if i < 2:
if i == 1:
lr = lr.astype(np.float64)
lr = lr / 255
tmp = np.squeeze(
model.predict(np.expand_dims(lr, axis=0)))
if i == 1:
tmp = tmp * 255
tmp[tmp < 0] = 0
tmp[tmp > 255] = 255
sr = tmp.astype(np.uint8)
elif i < 6:
sr = Utils.denormalize(
np.squeeze(model.predict(
np.expand_dims(rescale_imgs_to_neg1_1(lr),
axis=0)),
axis=0))
else:
sr = cv2.resize(lr, (0, 0),
fx=4,
fy=4,
interpolation=cv2.INTER_NEAREST)
# sixth step: append the calculated metric
m.append(metrics.MSE(hr, sr))
p.append(metrics.PSNR(hr, sr))
s.append(metrics.SSIM(hr, sr))
ms.append(metrics.MSSIM(hr, sr))
# seventh step: append the mean metric for this image
mse.append(np.mean(m))
psnr.append(np.mean(p))
ssim.append(np.mean(s))
mssim.append(np.mean(ms))
# eight step: append the mean metric for this model
tile_performance[model_names[i]] = (np.mean(mse), np.mean(psnr),
np.mean(ssim), np.mean(mssim))
# final output
print("Performance on single tiles:")
f = open("tile_performance.txt", "w")
for key in tile_performance:
print(
key + ": MSE = " + str(tile_performance[key][0]) +
", PSNR = " + str(tile_performance[key][1]) + ", SSIM = " +
str(tile_performance[key][2]),
", MSSIM = " + str(tile_performance[key][3]))
f.write(key + " " + str(tile_performance[key][0]) + " " +
str(tile_performance[key][1]) + " " +
str(tile_performance[key][2]) + " " +
str(tile_performance[key][3]) + "\n")
f.close()
if WHOLE_LR:
'''
Second calculating performance metrics on a single upscaled image
'''
img_performance = {}
for i, mp in tqdm(enumerate(model_paths)):
keras.backend.clear_session()
# first step: load the model
if i < 6:
model = load_model(mp, custom_objects=custom_objects[i])
# second step: changing the input layer
_in = Input(shape=test_images[0][0].shape)
_out = model(_in)
_model = Model(_in, _out)
mse = []
psnr = []
ssim = []
mssim = []
# third step: iterate over the test images
for test_pair in tqdm(test_images):
# fourth step: calculate the sr image
try:
if i < 2:
if i == 1:
lr = test_pair[0].astype(np.float64)
lr = lr / 255
else:
lr = test_pair[0]
tmp = np.squeeze(
_model.predict(np.expand_dims(lr, axis=0)))
if i == 1:
tmp = tmp * 255
tmp[tmp < 0] = 0
tmp[tmp > 255] = 255
sr = tmp.astype(np.uint8)
elif i < 6:
sr = Utils.denormalize(
np.squeeze(_model.predict(
np.expand_dims(rescale_imgs_to_neg1_1(
test_pair[0]),
axis=0)),
axis=0))
else:
sr = cv2.resize(test_pair[0], (0, 0),
fx=4,
fy=4,
interpolation=cv2.INTER_NEAREST)
# fifth step: append the metric for this image
mse.append(metrics.MSE(test_pair[1], sr))
psnr.append(metrics.PSNR(test_pair[1], sr))
ssim.append(metrics.SSIM(test_pair[1], sr))
mssim.append(metrics.MSSIM(test_pair[1], sr))
except:
mse.append("err")
psnr.append("err")
ssim.append("err")
mssim.append("err")
# sixth step: append the mean metric for this model
try:
img_performance[model_names[i]] = (np.mean(mse), np.mean(psnr),
np.mean(ssim),
np.mean(mssim))
except:
img_performance[model_names[i]] = ("err", "err", "err", "err")
# final output
print("Performance on whole lr:")
f = open("whole_lr_performance.txt", "w")
for key in img_performance:
print(
key + ": MSE = " + str(img_performance[key][0]) +
", PSNR = " + str(img_performance[key][1]) + ", SSIM = " +
str(img_performance[key][2]),
", MSSIM = " + str(img_performance[key][3]))
f.write(key + " " + str(img_performance[key][0]) + " " +
str(img_performance[key][1]) + " " +
str(img_performance[key][2]) + " " +
str(img_performance[key][3]) + "\n")
f.close()
if STITCHED:
'''
Second calculating performance metrics on a stitched image
'''
stitch_performance = {}
for i, mp in tqdm(enumerate(model_paths)):
keras.backend.clear_session()
# first step: load the model
if i < 6:
model = load_model(mp, custom_objects=custom_objects[i])
mse = []
psnr = []
ssim = []
mssim = []
o_mse = []
o_psnr = []
o_ssim = []
o_mssim = []
# second step: iterate over the test images
for test_pair in tqdm(test_images):
# third step: tile the test image
lr_tiles = Utils.tile_image(test_pair[0],
shape=tile_shapes[i][1])
lr_tiles_overlap = Utils.tile_image(test_pair[0],
shape=tile_shapes[i][1],
overlap=True)
sr_tiles = []
sr_tiles_overlap = []
# fourth step: iterate over the tiles
for lr in lr_tiles:
# fifth step: calculate the sr tiles
if i < 2:
if i == 1:
lr = lr.astype(np.float64)
lr = lr / 255
tmp = np.squeeze(
model.predict(np.expand_dims(lr, axis=0)))
if i == 1:
tmp = tmp * 255
tmp[tmp < 0] = 0
tmp[tmp > 255] = 255
sr = tmp.astype(np.uint8)
sr_tiles.append(sr)
elif i < 6:
sr_tiles.append(
Utils.denormalize(
np.squeeze(model.predict(
np.expand_dims(rescale_imgs_to_neg1_1(lr),
axis=0)),
axis=0)))
else:
sr_tiles.append(
cv2.resize(lr, (0, 0),
fx=4,
fy=4,
interpolation=cv2.INTER_NEAREST))
for lr in lr_tiles_overlap:
# fifth step: calculate the sr tiles
if i < 2:
if i == 1:
lr = lr.astype(np.float64)
lr = lr / 255
tmp = np.squeeze(
model.predict(np.expand_dims(lr, axis=0)))
if i == 1:
tmp = tmp * 255
tmp[tmp < 0] = 0
tmp[tmp > 255] = 255
sr = tmp.astype(np.uint8)
sr_tiles_overlap.append(sr)
elif i < 6:
sr_tiles_overlap.append(
Utils.denormalize(
np.squeeze(model.predict(
np.expand_dims(rescale_imgs_to_neg1_1(lr),
axis=0)),
axis=0)))
else:
sr_tiles_overlap.append(
cv2.resize(lr, (0, 0),
fx=4,
fy=4,
interpolation=cv2.INTER_NEAREST))
# sixth step: stitch the image
sr_simple = ImageStitching.stitch_images(
sr_tiles, test_pair[1].shape[1], test_pair[1].shape[0],
sr_tiles[0].shape[1], sr_tiles[0].shape[0],
test_pair[1].shape[1] // sr_tiles[0].shape[1],
test_pair[1].shape[0] // sr_tiles[0].shape[0])
sr_advanced = ImageStitching.stitching(
sr_tiles_overlap,
LR=None,
image_size=(test_pair[1].shape[0], test_pair[1].shape[1]),
adjustRGB=False,
overlap=True)
# seventh step: append the mean metric for this image
mse.append(metrics.MSE(test_pair[1], sr_simple))
psnr.append(metrics.PSNR(test_pair[1], sr_simple))
ssim.append(metrics.SSIM(test_pair[1], sr_simple))
mssim.append(metrics.MSSIM(test_pair[1], sr_simple))
o_mse.append(metrics.MSE(test_pair[1], sr_advanced))
o_psnr.append(metrics.PSNR(test_pair[1], sr_advanced))
o_ssim.append(metrics.SSIM(test_pair[1], sr_advanced))
o_mssim.append(metrics.MSSIM(test_pair[1], sr_advanced))
# ninth step: append the mean metric for this model
stitch_performance[model_names[i]] = [
(np.mean(mse), np.mean(psnr), np.mean(ssim), np.mean(mssim)),
(np.mean(o_mse), np.mean(o_psnr), np.mean(o_ssim),
np.mean(o_mssim))
]
# final output
print("Performance on stitched images:")
f = open("stitch_performance.txt", "w")
for key in stitch_performance:
print(
"simple stitch: " + key + ": MSE = " +
str(stitch_performance[key][0][0]) + ", PSNR = " +
str(stitch_performance[key][0][1]) + ", SSIM = " +
str(stitch_performance[key][0][2]),
", MSSIM = " + str(stitch_performance[key][0][3]))
print(
"advanced stitch: " + key + ": MSE = " +
str(stitch_performance[key][1][0]) + ", PSNR = " +
str(stitch_performance[key][1][1]) + ", SSIM = " +
str(stitch_performance[key][1][2]),
", MSSIM = " + str(stitch_performance[key][1][3]))
f.write(key + " " + str(stitch_performance[key][0][0]) + " " +
str(stitch_performance[key][0][1]) + " " +
str(stitch_performance[key][0][2]) + " " +
str(stitch_performance[key][0][3]) + "\n")
f.write(key + " " + str(stitch_performance[key][1][0]) + " " +
str(stitch_performance[key][1][1]) + " " +
str(stitch_performance[key][1][2]) + " " +
str(stitch_performance[key][1][3]) + "\n")
f.close()
import Utils_images
import matplotlib.pyplot as plt
image_shape = (336, 336, 3)
#input_dirs = [os.path.join('..', '..', 'DSIDS', 'HR', 'tiles_'+str(image_shape[0])),
# os.path.join('..', '..', 'DSIDS', 'LR', 'tiles_'+str(image_shape[0]) ,'4x_cubic')]
#test_images = []
#for img in sorted(os.listdir(os.path.join(input_dirs[1], 'ignore'))):
# if 'niklas_city_0009' in img:
# test_images.append(rescale_imgs_to_neg1_1(cv2.imread(os.path.join(input_dirs[1], 'ignore', img))))
test = rescale_imgs_to_neg1_1(
cv2.cvtColor(
cv2.imread(
os.path.join('..', '..', 'DSIDS', 'LR', '4x_cubic',
'niklas_city_0009.jpg'), cv2.IMREAD_COLOR),
cv2.COLOR_BGR2RGB))
#test = rescale_imgs_to_neg1_1(cv2.imread("1008.jpg",cv2.IMREAD_COLOR))
tile_list = Utils_images.crop_lr_image(test, hr_shape=(336, 336), overlap=True)
model = load_model(os.path.join("model", "gen_model90.h5"),
custom_objects={"tf": tf})
model.summary()
model.layers.pop(0)
_in = Input(shape=(756, 1008, 3))
#_in = Input(shape=(189, 252, 3))
_out = model(_in)
def main():
# paths to the models
model_paths = [os.path.join("..", "models", "SRDense-Type-3_ep80.h5"),
os.path.join("..", "models", "srdense-norm.h5"),
os.path.join("..", "models", "srresnet85.h5"),
os.path.join("..", "models", "gen_model90.h5"),
os.path.join("..", "models", "srgan60.h5"),
os.path.join("..", "models", "srgan-mse-20.h5"),
"Nearest"]
# corresponding names of the models
model_names = ["SRDense",
"SRDense-norm",
"SRResNet",
"SRGAN-from-scratch",
"SRGAN-percept-loss",
"SRGAN-mse",
"NearestNeighbor"]
custom_objects = [{},
{"tf": tf},
{"tf": tf},
{"tf": tf},
{"tf": tf},
{"tf": tf},
{}]
if not os.path.isdir(SAVE_PATH):
os.makedirs(SAVE_PATH)
test_image = cv2.cvtColor(cv2.imread(os.path.join(os.getcwd(), "test_image.jpg")), cv2.COLOR_BGR2RGB)
lr_size = ggt(test_image.shape[0], test_image.shape[1])
if lr_size == test_image.shape[0]:
lr_size = int(lr_size/10)
hr_size = lr_size*4
lr_tiles = Utils.tile_image(test_image, shape=(lr_size, lr_size))
lr_tiles_overlap = Utils.tile_image(test_image, shape=(lr_size, lr_size), overlap=True)
if WHOLE_LR:
'''
First upscaling whole lr image
'''
for i,mp in tqdm(enumerate(model_paths)):
keras.backend.clear_session()
# first step: load the model
if i < 6:
model = load_model(mp, custom_objects=custom_objects[i])
# second step: changing the input layer
_in = Input(shape=test_image.shape)
_out = model(_in)
_model = Model(_in, _out)
# third step propagating the image
if i < 2:
if i == 1:
lr = test_image.astype(np.float64)
lr = lr/255
else:
lr = test_image
tmp = np.squeeze(_model.predict(np.expand_dims(lr, axis=0)))
if i == 1:
tmp = tmp*255
tmp[tmp < 0] = 0
tmp[tmp > 255] = 255
sr = tmp.astype(np.uint8)
elif i < 6:
sr = Utils.denormalize(np.squeeze(_model.predict(np.expand_dims(rescale_imgs_to_neg1_1(test_image), axis=0)), axis=0))
else:
sr = cv2.resize(test_image, (0, 0),
fx=4,
fy=4,
interpolation=cv2.INTER_NEAREST)
# fourth step saving the image
cv2.imwrite(os.path.join(SAVE_PATH, model_names[i]+"_whole-lr.jpg"), cv2.cvtColor(sr, cv2.COLOR_RGB2BGR))
if STITCHED:
'''
second upscaling tiles and stitching them together
'''
for i,mp in tqdm(enumerate(model_paths)):
keras.backend.clear_session()
# first step: load the model
if i < 6:
model = load_model(mp, custom_objects=custom_objects[i])
# second step: changing the input layer
_in = Input(shape=(lr_size, lr_size, 3))
_out = model(_in)
_model = Model(_in, _out)
sr_tiles = []
# third step propagating the tiles
for tile in lr_tiles:
if i < 2:
if i == 1:
lr = tile.astype(np.float64)
lr = lr/255
else:
lr = tile
tmp = np.squeeze(_model.predict(np.expand_dims(lr, axis=0)))
if i == 1:
tmp = tmp*255
tmp[tmp < 0] = 0
tmp[tmp > 255] = 255
sr_tiles.append( tmp.astype(np.uint8) )
elif i < 6:
sr_tiles.append(Utils.denormalize(np.squeeze(_model.predict(np.expand_dims(rescale_imgs_to_neg1_1(tile), axis=0)), axis=0)))
else:
sr_tiles.append(cv2.resize(tile, (0, 0),
fx=4,
fy=4,
interpolation=cv2.INTER_NEAREST))
# fourth step stitch the tiles
sr_simple = ImageStitching.stitch_images(sr_tiles, hr_size*(test_image.shape[1]//lr_size), hr_size*(test_image.shape[0]//lr_size),
sr_tiles[0].shape[1], sr_tiles[0].shape[0], test_image.shape[1]//sr_tiles[0].shape[1],
test_image.shape[0]//sr_tiles[0].shape[0] )
# fourth step saving the image
cv2.imwrite(os.path.join(SAVE_PATH, model_names[i]+"_simple-stitched.jpg"), cv2.cvtColor(sr_simple, cv2.COLOR_RGB2BGR))
# the same again for overlapping stitching
sr_tiles_overlap = []
# propagating the tiles
for tile in lr_tiles_overlap:
if i < 2:
if i == 1:
lr = tile.astype(np.float64)
lr = lr/255
else:
lr = tile
tmp = np.squeeze(_model.predict(np.expand_dims(lr, axis=0)))
if i == 1:
tmp = tmp*255
tmp[tmp < 0] = 0
tmp[tmp > 255] = 255
sr_tiles_overlap.append( tmp.astype(np.uint8) )
elif i < 6:
sr_tiles_overlap.append(Utils.denormalize(np.squeeze(_model.predict(np.expand_dims(rescale_imgs_to_neg1_1(tile), axis=0)), axis=0)))
else:
sr_tiles_overlap.append(cv2.resize(tile, (0, 0),
fx=4,
fy=4,
interpolation=cv2.INTER_NEAREST))
# stitch the tiles
sr_overlap = ImageStitching.stitching(sr_tiles_overlap, LR=None,
image_size=(hr_size*(test_image.shape[0]//lr_size), hr_size*(test_image.shape[1]//lr_size)),
adjustRGB=False, overlap=True)
# save the image
cv2.imwrite(os.path.join(SAVE_PATH, model_names[i]+"_overlap-stitched.jpg"), cv2.cvtColor(sr_overlap, cv2.COLOR_RGB2BGR))
def train(img_shape,
epochs,
batch_size,
rescaling_factor,
input_dirs,
output_dir,
model_save_dir,
train_test_ratio,
gpu=1):
lr_shape = (img_shape[0] // rescaling_factor,
img_shape[1] // rescaling_factor, img_shape[2])
img_train_gen, img_test_gen = create_data_generator(
input_dirs[1],
input_dirs[0],
target_size_lr=(lr_shape[0], lr_shape[1]),
target_size_hr=(img_shape[0], img_shape[1]),
preproc_lr=rescale_imgs_to_neg1_1,
preproc_hr=rescale_imgs_to_neg1_1,
validation_split=train_test_ratio,
batch_size=batch_size)
batch_count = int(
(len(os.listdir(os.path.join(input_dirs[1], 'ignore'))) / batch_size) *
(1 - train_test_ratio))
test_image = []
for img in sorted(os.listdir(os.path.join(input_dirs[1], 'ignore'))):
if 'niklas_city_0009' in img:
test_image.append(
rescale_imgs_to_neg1_1(
cv2.imread(os.path.join(input_dirs[1], 'ignore', img))))
print("test length: ", len(test_image))
loss = VGG_LOSS(image_shape)
generator = Generator(lr_shape, rescaling_factor).generator()
print('memory usage generator: ',
get_model_memory_usage(batch_size, generator))
optimizer = Utils_model.get_optimizer()
if gpu > 1:
try:
print("multi_gpu_model generator")
par_generator = multi_gpu_model(generator, gpus=2)
except:
par_generator = generator
print("single_gpu_model generator")
else:
par_generator = generator
print("single_gpu_model generator")
par_generator.compile(loss=loss.loss, optimizer=optimizer)
par_generator.summary()
for e in range(1, epochs + 1):
print('-' * 15, 'Epoch %d' % e, '-' * 15)
for i in tqdm(range(batch_count)):
batch = next(img_train_gen)
image_batch_hr = batch[1]
image_batch_lr = batch[0]
if image_batch_hr.shape[0] == batch_size and image_batch_lr.shape[
0] == batch_size:
g_loss = par_generator.train_on_batch(image_batch_lr,
image_batch_hr)
else:
print("weird multi_gpu_model batch error dis: ")
print("hr batch shape: ", image_batch_hr.shape)
print("lr batch shape: ", image_batch_lr.shape)
#if e == 1 or e % 5 == 0:
#Utils.generate_test_image(output_dir, e, generator, test_image)
if e % 5 == 0:
generator.save(os.path.join(model_save_dir, 'srresnet%d.h5' % e))
generator.save(os.path.join(model_save_dir, 'srresnet.h5' % e))
def train(img_shape, epochs, batch_size, rescaling_factor, input_dirs,
output_dir, model_save_dir, train_test_ratio):
lr_shape = (img_shape[0] // rescaling_factor,
img_shape[1] // rescaling_factor, img_shape[2])
img_train_gen, img_test_gen = create_data_generator(
input_dirs[1],
input_dirs[0],
target_size_lr=(lr_shape[0], lr_shape[1]),
target_size_hr=(img_shape[0], img_shape[1]),
preproc_lr=rescale_imgs_to_neg1_1,
preproc_hr=rescale_imgs_to_neg1_1,
validation_split=train_test_ratio,
batch_size=batch_size)
loss = VGG_LOSS(image_shape)
batch_count = int(
(len(os.listdir(os.path.join(input_dirs[1], 'ignore'))) / batch_size) *
(1 - train_test_ratio))
test_image = []
for img in sorted(os.listdir(os.path.join(input_dirs[1], 'ignore'))):
if 'niklas_city_0009' in img:
test_image.append(
rescale_imgs_to_neg1_1(
cv2.imread(os.path.join(input_dirs[1], 'ignore', img))))
print("test length: ", len(test_image))
generator = Generator(lr_shape, rescaling_factor).generator()
discriminator = Discriminator(img_shape).discriminator()
print('memory usage generator: ',
get_model_memory_usage(batch_size, generator))
print('memory usage discriminator: ',
get_model_memory_usage(batch_size, discriminator))
optimizer = Utils_model.get_optimizer()
try:
print("multi_gpu_model generator")
par_generator = multi_gpu_model(generator, gpus=2)
except:
par_generator = generator
print("single_gpu_model generator")
try:
print("multi_gpu_model discriminator")
par_discriminator = multi_gpu_model(discriminator, gpus=2)
except:
par_discriminator = discriminator
print("single_gpu_model discriminator")
par_generator.compile(loss=loss.loss, optimizer=optimizer)
par_discriminator.compile(loss="binary_crossentropy", optimizer=optimizer)
gan, par_gan = get_gan_network(par_discriminator, lr_shape, par_generator,
optimizer, loss.loss, batch_size)
par_discriminator.summary()
par_generator.summary()
par_gan.summary()
loss_file = open(model_save_dir + 'losses.txt', 'w+')
loss_file.close()
for e in range(1, epochs + 1):
print('-' * 15, 'Epoch %d' % e, '-' * 15)
if e == 100:
optimizer.lr = 1e-5
for i in tqdm(range(batch_count)):
batch = next(img_train_gen)
image_batch_hr = batch[1]
image_batch_lr = batch[0]
generated_images_sr = generator.predict(image_batch_lr)
real_data_Y = np.ones(batch_size) - \
np.random.random_sample(batch_size)*0.2
fake_data_Y = np.random.random_sample(batch_size) * 0.2
par_discriminator.trainable = True
if image_batch_hr.shape[0] == batch_size and image_batch_lr.shape[
0] == batch_size:
d_loss_real = par_discriminator.train_on_batch(
image_batch_hr, real_data_Y)
d_loss_fake = par_discriminator.train_on_batch(
generated_images_sr, fake_data_Y)
discriminator_loss = 0.5 * np.add(d_loss_fake, d_loss_real)
else:
print("weird multi_gpu_model batch error dis: ")
print("hr batch shape: ", image_batch_hr.shape)
print("lr batch shape: ", image_batch_lr.shape)
print("gan y shape: ", gan_Y.shape)
batch = next(img_train_gen)
image_batch_hr = batch[1]
image_batch_lr = batch[0]
gan_Y = np.ones(batch_size) - \
np.random.random_sample(batch_size)*0.2
discriminator.trainable = False
if image_batch_hr.shape[0] == batch_size and image_batch_lr.shape[
0] == batch_size:
gan_loss = par_gan.train_on_batch(image_batch_lr,
[image_batch_hr, gan_Y])
else:
print("weird multi_gpu_model batch error gan: ")
print("hr batch shape: ", image_batch_hr.shape)
print("lr batch shape: ", image_batch_lr.shape)
print("gan y shape: ", gan_Y.shape)
print("discriminator_loss : %f" % discriminator_loss)
print("gan_loss :", gan_loss)
gan_loss = str(gan_loss)
loss_file = open(model_save_dir + '_losses.txt', 'a')
loss_file.write('epoch%d : gan_loss = %s ; discriminator_loss = %f\n' %
(e, gan_loss, discriminator_loss))
loss_file.close()
if e == 1 or e % 5 == 0:
Utils.generate_test_image(output_dir, e, generator, test_image)
if e % 5 == 0:
generator.save(os.path.join(model_save_dir, 'gen_model%d.h5' % e))
discriminator.save(
os.path.join(model_save_dir, 'dis_model%d.h5' % e))