From e688c2847b7fd66fb6e66120af92daa101b90bba Mon Sep 17 00:00:00 2001 From: HuiLi Date: Fri, 6 Oct 2017 13:17:01 +0800 Subject: [PATCH 1/2] add the operation when the image shape is () --- train.py | 5 +++++ utils.py | 2 ++ 2 files changed, 7 insertions(+) diff --git a/train.py b/train.py index cc4ea81..616d84c 100644 --- a/train.py +++ b/train.py @@ -113,6 +113,11 @@ def train(style_weight, content_imgs_path, style_imgs_path, encoder_path, save_p content_batch = get_train_images(content_batch_path, crop_height=HEIGHT, crop_width=WIDTH) style_batch = get_train_images(style_batch_path, crop_height=HEIGHT, crop_width=WIDTH) + + if content_batch == () or style_batch == (): + continue + + # run the training step sess.run(train_op, feed_dict={content: content_batch, style: style_batch}) diff --git a/utils.py b/utils.py index ee0b527..06a68ae 100644 --- a/utils.py +++ b/utils.py @@ -24,6 +24,8 @@ def get_train_images(paths, resize_len=512, crop_height=256, crop_width=256): images = [] for path in paths: image = imread(path, mode='RGB') + if image.shape == (): + return image.shape height, width, _ = image.shape if height < width: From 4767e43bdb933140bd32d7abcd9e696698a86339 Mon Sep 17 00:00:00 2001 From: HuiLi Date: Fri, 6 Oct 2017 13:20:43 +0800 Subject: [PATCH 2/2] add training data path --- main.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/main.py b/main.py index ae52d53..737d68e 100644 --- a/main.py +++ b/main.py @@ -32,8 +32,14 @@ def main(): if IS_TRAINING: - content_imgs_path = list_images('../MS_COCO') # path to training content dataset - style_imgs_path = list_images('../WikiArt') # path to training style dataset + content_imgs_path = list_images('D:/ImageDatabase/Microsoft_COCO2014/train2014') # path to training content dataset + style_imgs_path = list_images('D:/ImageDatabase/WikiArt_database/all') # path to training style dataset + + # content_imgs_path = list_images('D:/ImageDatabase/Microsoft_COCO2014/train2014') # path to training content dataset + # style_imgs_path = list_images('D:/ImageDatabase/WikiArt_database/train_1') # path to training style dataset + + # content_imgs_path = list_images('D:/ImageDatabase/train_data_temp/MS_COCO_1000') # path to training content dataset + # style_imgs_path = list_images('D:/ImageDatabase/train_data_temp/WikiArt_1000') # path to training style dataset for style_weight, model_save_path in zip(STYLE_WEIGHTS, MODEL_SAVE_PATHS): print('\nBegin to train the network with the style weight: %.2f ...\n' % style_weight)