diff options
Diffstat (limited to 'main.py')
| -rw-r--r-- | main.py | 7 |
1 files changed, 5 insertions, 2 deletions
@@ -1,7 +1,9 @@ import os import sys import argparse + import tensorflow as tf +from skimage import transform import hyperparameters as hp from losses import YourModel @@ -37,7 +39,8 @@ def parse_args(): return parser.parse_args() def train(model): - for _ in range(hp.num_epochs): + for i in range(hp.num_epochs): + print('batch', i) model.train_step() def main(): @@ -51,7 +54,7 @@ def main(): content_image = imread(ARGS.content) style_image = imread(ARGS.style) - style_image = transform.resize(style_image, content_image.shape) + style_image = transform.resize(style_image, content_image.shape).astype('uint8') my_model = YourModel(content_image=content_image, style_image=style_image) my_model.vgg16.build([1, 255, 255, 3]) my_model.vgg16.load_weights('vgg16_imagenet.h5', by_name=True) |