[ { "path": ".github/FUNDING.yml", "content": "# These are supported funding model platforms\n\ngithub: [lengstrom] # Replace with up to 4 GitHub Sponsors-enabled usernames e.g., [user1, user2]\n" }, { "path": ".gitignore", "content": "t Byte-compiled / optimized / DLL files\ndeps.txt\narchive\nsaver\n*~\nstyles\npngs\npreds\n\n*.sw*\ndata\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nenv/\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\n*.egg-info/\n.installed.cfg\n*.egg\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*,cover\n.hypothesis/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\ntarget/\n\n# IPython Notebook\n.ipynb_checkpoints\n\n# pyenv\n.python-version\n\n# celery beat schedule file\ncelerybeat-schedule\n\n# dotenv\n.env\n\n# virtualenv\nvenv/\nENV/\n\n# Spyder project settings\n.spyderproject\n\n# Rope project settings\n.ropeproject\n\n# PyCharm\n.idea\n\n# checkpoint\ncheckpoint\n" }, { "path": "CITATION.cff", "content": "# YAML 1.2\n---\nauthors: \n -\n family-names: Engstrom\n given-names: Logan\ncff-version: \"1.1.0\"\ndate-released: 2016-10-31\nmessage: \"If you use this software, please cite it using these metadata.\"\nrepository-code: \"https://github.com/lengstrom/fast-style-transfer\"\ntitle: \"Fast Style Transfer\"\nversion: \"1.0\"\n...\n" }, { "path": "README.md", "content": "## Fast Style Transfer in [TensorFlow](https://github.com/tensorflow/tensorflow)\n\nAdd styles from famous paintings to any photo in a fraction of a second! [You can even style videos!](#video-stylization)\n\n

\n\n\n\n

\n

\nIt takes 100ms on a 2015 Titan X to style the MIT Stata Center (1024×680) like Udnie, by Francis Picabia.\n

\n\nOur implementation is based off of a combination of Gatys' [A Neural Algorithm of Artistic Style](https://arxiv.org/abs/1508.06576), Johnson's [Perceptual Losses for Real-Time Style Transfer and Super-Resolution](http://cs.stanford.edu/people/jcjohns/eccv16/), and Ulyanov's [Instance Normalization](https://arxiv.org/abs/1607.08022). \n\n### Sponsorship\nPlease consider sponsoring my work on this project!\n\n### License\nCopyright (c) 2016 Logan Engstrom. Contact me for commercial use (or rather any use that is not academic research) (email: engstrom at my university's domain dot edu). Free for research use, as long as proper attribution is given and this copyright notice is retained.\n\n## Video Stylization \nHere we transformed every frame in a video, then combined the results. [Click to go to the full demo on YouTube!](https://www.youtube.com/watch?v=xVJwwWQlQ1o) The style here is Udnie, as above.\n
\n \n Stylized fox video. Click to go to YouTube!\n \n
\n\nSee how to generate these videos [here](#stylizing-video)!\n\n## Image Stylization\nWe added styles from various paintings to a photo of Chicago. Click on thumbnails to see full applied style images.\n
\n\n
\n \n
\n\n\n\n\n
\n\n\n\n\n\n
\n\n\n\n\n
\n\n## Implementation Details\nOur implementation uses TensorFlow to train a fast style transfer network. We use roughly the same transformation network as described in Johnson, except that batch normalization is replaced with Ulyanov's instance normalization, and the scaling/offset of the output `tanh` layer is slightly different. We use a loss function close to the one described in Gatys, using VGG19 instead of VGG16 and typically using \"shallower\" layers than in Johnson's implementation (e.g. we use `relu1_1` rather than `relu1_2`). Empirically, this results in larger scale style features in transformations.\n## Virtual Environment Setup (Anaconda) - Windows/Linux\nTested on\n| Spec | |\n|-----------------------------|-------------------------------------------------------------|\n| Operating System | Windows 10 Home |\n| GPU | Nvidia GTX 2080 TI |\n| CUDA Version | 11.0 |\n| Driver Version | 445.75 |\n### Step 1：Install Anaconda\nhttps://docs.anaconda.com/anaconda/install/\n### Step 2：Build a virtual environment\nRun the following commands in sequence in Anaconda Prompt:\n```\nconda create -n tf-gpu tensorflow-gpu=2.1.0\nconda activate tf-gpu\nconda install jupyterlab\njupyter lab\n```\nRun the following command in the notebook or just conda install the package:\n```\n!pip install moviepy==1.0.2\n```\nFollow the commands below to use fast-style-transfer\n## Documentation\n### Training Style Transfer Networks\nUse `style.py` to train a new style transfer network. Run `python style.py` to view all the possible parameters. Training takes 4-6 hours on a Maxwell Titan X. [More detailed documentation here](docs.md#stylepy). **Before you run this, you should run `setup.sh`**. Example usage:\n\n python style.py --style path/to/style/img.jpg \\\n --checkpoint-dir checkpoint/path \\\n --test path/to/test/img.jpg \\\n --test-dir path/to/test/dir \\\n --content-weight 1.5e1 \\\n --checkpoint-iterations 1000 \\\n --batch-size 20\n\n### Evaluating Style Transfer Networks\nUse `evaluate.py` to evaluate a style transfer network. Run `python evaluate.py` to view all the possible parameters. Evaluation takes 100 ms per frame (when batch size is 1) on a Maxwell Titan X. [More detailed documentation here](docs.md#evaluatepy). Takes several seconds per frame on a CPU. **Models for evaluation are [located here](https://drive.google.com/drive/folders/0B9jhaT37ydSyRk9UX0wwX3BpMzQ?resourcekey=0-Z9LcNHC-BTB4feKwm4loXw&usp=sharing)**. Example usage:\n\n python evaluate.py --checkpoint path/to/style/model.ckpt \\\n --in-path dir/of/test/imgs/ \\\n --out-path dir/for/results/\n\n### Stylizing Video\nUse `transform_video.py` to transfer style into a video. Run `python transform_video.py` to view all the possible parameters. Requires `ffmpeg`. [More detailed documentation here](docs.md#transform_videopy). Example usage:\n\n python transform_video.py --in-path path/to/input/vid.mp4 \\\n --checkpoint path/to/style/model.ckpt \\\n --out-path out/video.mp4 \\\n --device /gpu:0 \\\n --batch-size 4\n\n### Requirements\nYou will need the following to run the above:\n- TensorFlow 0.11.0\n- Python 2.7.9, Pillow 3.4.2, scipy 0.18.1, numpy 1.11.2\n- If you want to train (and don't want to wait for 4 months):\n - A decent GPU\n - All the required NVIDIA software to run TF on a GPU (cuda, etc)\n- ffmpeg 3.1.3 if you want to stylize video\n\n### Citation\n```\n @misc{engstrom2016faststyletransfer,\n author = {Logan Engstrom},\n title = {Fast Style Transfer},\n year = {2016},\n howpublished = {\\url{https://github.com/lengstrom/fast-style-transfer/}},\n note = {commit xxxxxxx}\n }\n```\n\n### Attributions/Thanks\n- This project could not have happened without the advice (and GPU access) given by [Anish Athalye](http://www.anishathalye.com/). \n - The project also borrowed some code from Anish's [Neural Style](https://github.com/anishathalye/neural-style/)\n- Some readme/docs formatting was borrowed from Justin Johnson's [Fast Neural Style](https://github.com/jcjohnson/fast-neural-style)\n- The image of the Stata Center at the very beginning of the README was taken by [Juan Paulo](https://juanpaulo.me/)\n\n### Related Work\n- Michael Ramos ported this network [to use CoreML on iOS](https://medium.com/@rambossa/diy-prisma-fast-style-transfer-app-with-coreml-and-tensorflow-817c3b90dacd)\n" }, { "path": "docs.md", "content": "## style.py \n\n`style.py` trains networks that can transfer styles from artwork into images.\n\n**Flags**\n- `--checkpoint-dir`: Directory to save checkpoint in. Required.\n- `--style`: Path to style image. Required.\n- `--train-path`: Path to training images folder. Default: `data/train2014`.\n- `--test`: Path to content image to test network on at at every checkpoint iteration. Default: no image.\n- `--test-dir`: Path to directory to save test images in. Required if `--test` is passed a value.\n- `--epochs`: Epochs to train for. Default: `2`.\n- `--batch-size`: Batch size for training. Default: `4`.\n- `--checkpoint-iterations`: Number of iterations to go for between checkpoints. Default: `2000`.\n- `--vgg-path`: Path to VGG19 network (default). Can pass VGG16 if you want to try out other loss functions. Default: `data/imagenet-vgg-verydeep-19.mat`.\n- `--content-weight`: Weight of content in loss function. Default: `7.5e0`.\n- `--style-weight`: Weight of style in loss function. Default: `1e2`.\n- `--tv-weight`: Weight of total variation term in loss function. Default: `2e2`.\n- `--learning-rate`: Learning rate for optimizer. Default: `1e-3`.\n- `--slow`: For debugging loss function. Direct optimization on pixels using Gatys' approach. Uses `test` image as content value, `test_dir` for saving fully optimized images.\n\n\n## evaluate.py\n`evaluate.py` evaluates trained networks given a checkpoint directory. If evaluating images from a directory, every image in the directory must have the same dimensions.\n\n**Flags**\n- `--checkpoint`: Directory or `ckpt` file to load checkpoint from. Required.\n- `--in-path`: Path of image or directory of images to transform. Required.\n- `--out-path`: Out path of transformed image or out directory to put transformed images from in directory (if `in_path` is a directory). Required.\n- `--device`: Device used to transform image. Default: `/cpu:0`.\n- `--batch-size`: Batch size used to evaluate images. In particular meant for directory transformations. Default: `4`.\n- `--allow-different-dimensions`: Allow different image dimensions. Default: not enabled\n\n## transform_video.py\n`transform_video.py` transforms videos into stylized videos given a style transfer net.\n\n**Flags**\n- `--checkpoint-dir`: Directory or `ckpt` file to load checkpoint from. Required.\n- `--in-path`: Path to video to transfer style to. Required.\n- `--out-path`: Path to out video. Required.\n- `--tmp-dir`: Directory to put temporary processing files in. Will generate a dir if you do not pass it a path. Will delete tmpdir afterwards. Default: randomly generates invisible dir, then deletes it after execution completion.\n- `--device`: Device to evaluate frames with. Default: `/gpu:0`.\n- `--batch-size`: Batch size for evaluating images. Default: `4`.\n" }, { "path": "evaluate.py", "content": "from __future__ import print_function\nimport sys\nsys.path.insert(0, 'src')\nimport transform, numpy as np, vgg, pdb, os\nimport scipy.misc\nimport tensorflow as tf\nfrom utils import save_img, get_img, exists, list_files\nfrom argparse import ArgumentParser\nfrom collections import defaultdict\nimport time\nimport json\nimport subprocess\nimport numpy\nfrom moviepy.video.io.VideoFileClip import VideoFileClip\nimport moviepy.video.io.ffmpeg_writer as ffmpeg_writer\n\nBATCH_SIZE = 4\nDEVICE = '/gpu:0'\n\n\ndef ffwd_video(path_in, path_out, checkpoint_dir, device_t='/gpu:0', batch_size=4):\n video_clip = VideoFileClip(path_in, audio=False)\n video_writer = ffmpeg_writer.FFMPEG_VideoWriter(path_out, video_clip.size, video_clip.fps, codec=\"libx264\",\n preset=\"medium\", bitrate=\"2000k\",\n audiofile=path_in, threads=None,\n ffmpeg_params=None)\n\n g = tf.Graph()\n soft_config = tf.compat.v1.ConfigProto(allow_soft_placement=True)\n soft_config.gpu_options.allow_growth = True\n with g.as_default(), g.device(device_t), \\\n tf.compat.v1.Session(config=soft_config) as sess:\n batch_shape = (batch_size, video_clip.size[1], video_clip.size[0], 3)\n img_placeholder = tf.compat.v1.placeholder(tf.float32, shape=batch_shape,\n name='img_placeholder')\n\n preds = transform.net(img_placeholder)\n saver = tf.compat.v1.train.Saver()\n if os.path.isdir(checkpoint_dir):\n ckpt = tf.train.get_checkpoint_state(checkpoint_dir)\n if ckpt and ckpt.model_checkpoint_path:\n saver.restore(sess, ckpt.model_checkpoint_path)\n else:\n raise Exception(\"No checkpoint found...\")\n else:\n saver.restore(sess, checkpoint_dir)\n\n X = np.zeros(batch_shape, dtype=np.float32)\n\n def style_and_write(count):\n for i in range(count, batch_size):\n X[i] = X[count - 1] # Use last frame to fill X\n _preds = sess.run(preds, feed_dict={img_placeholder: X})\n for i in range(0, count):\n video_writer.write_frame(np.clip(_preds[i], 0, 255).astype(np.uint8))\n\n frame_count = 0 # The frame count that written to X\n for frame in video_clip.iter_frames():\n X[frame_count] = frame\n frame_count += 1\n if frame_count == batch_size:\n style_and_write(frame_count)\n frame_count = 0\n\n if frame_count != 0:\n style_and_write(frame_count)\n\n video_writer.close()\n\n\n# get img_shape\ndef ffwd(data_in, paths_out, checkpoint_dir, device_t='/gpu:0', batch_size=4):\n assert len(paths_out) > 0\n is_paths = type(data_in[0]) == str\n if is_paths:\n assert len(data_in) == len(paths_out)\n img_shape = get_img(data_in[0]).shape\n else:\n assert data_in.size[0] == len(paths_out)\n img_shape = X[0].shape\n\n g = tf.Graph()\n batch_size = min(len(paths_out), batch_size)\n curr_num = 0\n soft_config = tf.compat.v1.ConfigProto(allow_soft_placement=True)\n soft_config.gpu_options.allow_growth = True\n with g.as_default(), g.device(device_t), \\\n tf.compat.v1.Session(config=soft_config) as sess:\n batch_shape = (batch_size,) + img_shape\n img_placeholder = tf.compat.v1.placeholder(tf.float32, shape=batch_shape,\n name='img_placeholder')\n\n preds = transform.net(img_placeholder)\n saver = tf.compat.v1.train.Saver()\n if os.path.isdir(checkpoint_dir):\n ckpt = tf.train.get_checkpoint_state(checkpoint_dir)\n if ckpt and ckpt.model_checkpoint_path:\n saver.restore(sess, ckpt.model_checkpoint_path)\n else:\n raise Exception(\"No checkpoint found...\")\n else:\n saver.restore(sess, checkpoint_dir)\n\n num_iters = int(len(paths_out)/batch_size)\n for i in range(num_iters):\n pos = i * batch_size\n curr_batch_out = paths_out[pos:pos+batch_size]\n if is_paths:\n curr_batch_in = data_in[pos:pos+batch_size]\n X = np.zeros(batch_shape, dtype=np.float32)\n for j, path_in in enumerate(curr_batch_in):\n img = get_img(path_in)\n assert img.shape == img_shape, \\\n 'Images have different dimensions. ' + \\\n 'Resize images or use --allow-different-dimensions.'\n X[j] = img\n else:\n X = data_in[pos:pos+batch_size]\n\n _preds = sess.run(preds, feed_dict={img_placeholder:X})\n for j, path_out in enumerate(curr_batch_out):\n save_img(path_out, _preds[j])\n \n remaining_in = data_in[num_iters*batch_size:]\n remaining_out = paths_out[num_iters*batch_size:]\n if len(remaining_in) > 0:\n ffwd(remaining_in, remaining_out, checkpoint_dir, \n device_t=device_t, batch_size=1)\n\ndef ffwd_to_img(in_path, out_path, checkpoint_dir, device='/cpu:0'):\n paths_in, paths_out = [in_path], [out_path]\n ffwd(paths_in, paths_out, checkpoint_dir, batch_size=1, device_t=device)\n\ndef ffwd_different_dimensions(in_path, out_path, checkpoint_dir, \n device_t=DEVICE, batch_size=4):\n in_path_of_shape = defaultdict(list)\n out_path_of_shape = defaultdict(list)\n for i in range(len(in_path)):\n in_image = in_path[i]\n out_image = out_path[i]\n shape = \"%dx%dx%d\" % get_img(in_image).shape\n in_path_of_shape[shape].append(in_image)\n out_path_of_shape[shape].append(out_image)\n for shape in in_path_of_shape:\n print('Processing images of shape %s' % shape)\n ffwd(in_path_of_shape[shape], out_path_of_shape[shape], \n checkpoint_dir, device_t, batch_size)\n\ndef build_parser():\n parser = ArgumentParser()\n parser.add_argument('--checkpoint', type=str,\n dest='checkpoint_dir',\n help='dir or .ckpt file to load checkpoint from',\n metavar='CHECKPOINT', required=True)\n\n parser.add_argument('--in-path', type=str,\n dest='in_path',help='dir or file to transform',\n metavar='IN_PATH', required=True)\n\n help_out = 'destination (dir or file) of transformed file or files'\n parser.add_argument('--out-path', type=str,\n dest='out_path', help=help_out, metavar='OUT_PATH',\n required=True)\n\n parser.add_argument('--device', type=str,\n dest='device',help='device to perform compute on',\n metavar='DEVICE', default=DEVICE)\n\n parser.add_argument('--batch-size', type=int,\n dest='batch_size',help='batch size for feedforwarding',\n metavar='BATCH_SIZE', default=BATCH_SIZE)\n\n parser.add_argument('--allow-different-dimensions', action='store_true',\n dest='allow_different_dimensions', \n help='allow different image dimensions')\n\n return parser\n\ndef check_opts(opts):\n exists(opts.checkpoint_dir, 'Checkpoint not found!')\n exists(opts.in_path, 'In path not found!')\n if os.path.isdir(opts.out_path):\n exists(opts.out_path, 'out dir not found!')\n assert opts.batch_size > 0\n\ndef main():\n parser = build_parser()\n opts = parser.parse_args()\n check_opts(opts)\n\n if not os.path.isdir(opts.in_path):\n if os.path.exists(opts.out_path) and os.path.isdir(opts.out_path):\n out_path = \\\n os.path.join(opts.out_path,os.path.basename(opts.in_path))\n else:\n out_path = opts.out_path\n\n ffwd_to_img(opts.in_path, out_path, opts.checkpoint_dir,\n device=opts.device)\n else:\n files = list_files(opts.in_path)\n full_in = [os.path.join(opts.in_path,x) for x in files]\n full_out = [os.path.join(opts.out_path,x) for x in files]\n if opts.allow_different_dimensions:\n ffwd_different_dimensions(full_in, full_out, opts.checkpoint_dir, \n device_t=opts.device, batch_size=opts.batch_size)\n else :\n ffwd(full_in, full_out, opts.checkpoint_dir, device_t=opts.device,\n batch_size=opts.batch_size)\n\nif __name__ == '__main__':\n main()\n" }, { "path": "setup.sh", "content": "#! /bin/bash\n\nmkdir data\ncd data\nwget http://www.vlfeat.org/matconvnet/models/beta16/imagenet-vgg-verydeep-19.mat\nmkdir bin\nwget http://msvocds.blob.core.windows.net/coco2014/train2014.zip\nunzip -q train2014.zip\n" }, { "path": "src/optimize.py", "content": "from __future__ import print_function\nimport functools\nimport vgg, pdb, time\nimport tensorflow as tf, numpy as np, os\nimport transform\nfrom utils import get_img\n\nSTYLE_LAYERS = ('relu1_1', 'relu2_1', 'relu3_1', 'relu4_1', 'relu5_1')\nCONTENT_LAYER = 'relu4_2'\nDEVICES = 'CUDA_VISIBLE_DEVICES'\n\n# np arr, np arr\ndef optimize(content_targets, style_target, content_weight, style_weight,\n tv_weight, vgg_path, epochs=2, print_iterations=1000,\n batch_size=4, save_path='saver/fns.ckpt', slow=False,\n learning_rate=1e-3, debug=False):\n if slow:\n batch_size = 1\n mod = len(content_targets) % batch_size\n if mod > 0:\n print(\"Train set has been trimmed slightly..\")\n content_targets = content_targets[:-mod] \n\n style_features = {}\n\n batch_shape = (batch_size,256,256,3)\n style_shape = (1,) + style_target.shape\n print(style_shape)\n\n # precompute style features\n with tf.Graph().as_default(), tf.device('/cpu:0'), tf.compat.v1.Session() as sess:\n style_image = tf.compat.v1.placeholder(tf.float32, shape=style_shape, name='style_image')\n style_image_pre = vgg.preprocess(style_image)\n net = vgg.net(vgg_path, style_image_pre)\n style_pre = np.array([style_target])\n for layer in STYLE_LAYERS:\n features = net[layer].eval(feed_dict={style_image:style_pre})\n features = np.reshape(features, (-1, features.shape[3]))\n gram = np.matmul(features.T, features) / features.size\n style_features[layer] = gram\n\n with tf.Graph().as_default(), tf.compat.v1.Session() as sess:\n X_content = tf.compat.v1.placeholder(tf.float32, shape=batch_shape, name=\"X_content\")\n X_pre = vgg.preprocess(X_content)\n\n # precompute content features\n content_features = {}\n content_net = vgg.net(vgg_path, X_pre)\n content_features[CONTENT_LAYER] = content_net[CONTENT_LAYER]\n\n if slow:\n preds = tf.Variable(\n tf.random.normal(X_content.get_shape()) * 0.256\n )\n preds_pre = preds\n else:\n preds = transform.net(X_content/255.0)\n preds_pre = vgg.preprocess(preds)\n\n net = vgg.net(vgg_path, preds_pre)\n\n content_size = _tensor_size(content_features[CONTENT_LAYER])*batch_size\n assert _tensor_size(content_features[CONTENT_LAYER]) == _tensor_size(net[CONTENT_LAYER])\n content_loss = content_weight * (2 * tf.nn.l2_loss(\n net[CONTENT_LAYER] - content_features[CONTENT_LAYER]) / content_size\n )\n\n style_losses = []\n for style_layer in STYLE_LAYERS:\n layer = net[style_layer]\n bs, height, width, filters = map(lambda i:i,layer.get_shape())\n size = height * width * filters\n feats = tf.reshape(layer, (bs, height * width, filters))\n feats_T = tf.transpose(a=feats, perm=[0,2,1])\n grams = tf.matmul(feats_T, feats) / size\n style_gram = style_features[style_layer]\n style_losses.append(2 * tf.nn.l2_loss(grams - style_gram)/style_gram.size)\n\n style_loss = style_weight * functools.reduce(tf.add, style_losses) / batch_size\n\n # total variation denoising\n tv_y_size = _tensor_size(preds[:,1:,:,:])\n tv_x_size = _tensor_size(preds[:,:,1:,:])\n y_tv = tf.nn.l2_loss(preds[:,1:,:,:] - preds[:,:batch_shape[1]-1,:,:])\n x_tv = tf.nn.l2_loss(preds[:,:,1:,:] - preds[:,:,:batch_shape[2]-1,:])\n tv_loss = tv_weight*2*(x_tv/tv_x_size + y_tv/tv_y_size)/batch_size\n\n loss = content_loss + style_loss + tv_loss\n\n # overall loss\n train_step = tf.compat.v1.train.AdamOptimizer(learning_rate).minimize(loss)\n sess.run(tf.compat.v1.global_variables_initializer())\n import random\n uid = random.randint(1, 100)\n print(\"UID: %s\" % uid)\n for epoch in range(epochs):\n num_examples = len(content_targets)\n iterations = 0\n while iterations * batch_size < num_examples:\n start_time = time.time()\n curr = iterations * batch_size\n step = curr + batch_size\n X_batch = np.zeros(batch_shape, dtype=np.float32)\n for j, img_p in enumerate(content_targets[curr:step]):\n X_batch[j] = get_img(img_p, (256,256,3)).astype(np.float32)\n\n iterations += 1\n assert X_batch.shape[0] == batch_size\n\n feed_dict = {\n X_content:X_batch\n }\n\n train_step.run(feed_dict=feed_dict)\n end_time = time.time()\n delta_time = end_time - start_time\n if debug:\n print(\"UID: %s, batch time: %s\" % (uid, delta_time))\n is_print_iter = int(iterations) % print_iterations == 0\n if slow:\n is_print_iter = epoch % print_iterations == 0\n is_last = epoch == epochs - 1 and iterations * batch_size >= num_examples\n should_print = is_print_iter or is_last\n if should_print:\n to_get = [style_loss, content_loss, tv_loss, loss, preds]\n test_feed_dict = {\n X_content:X_batch\n }\n\n tup = sess.run(to_get, feed_dict = test_feed_dict)\n _style_loss,_content_loss,_tv_loss,_loss,_preds = tup\n losses = (_style_loss, _content_loss, _tv_loss, _loss)\n if slow:\n _preds = vgg.unprocess(_preds)\n else:\n saver = tf.compat.v1.train.Saver()\n res = saver.save(sess, save_path)\n yield(_preds, losses, iterations, epoch)\n\ndef _tensor_size(tensor):\n from operator import mul\n return functools.reduce(mul, (d for d in tensor.get_shape()[1:]), 1)\n" }, { "path": "src/transform.py", "content": "import tensorflow as tf, pdb\n\nWEIGHTS_INIT_STDEV = .1\n\ndef net(image):\n conv1 = _conv_layer(image, 32, 9, 1)\n conv2 = _conv_layer(conv1, 64, 3, 2)\n conv3 = _conv_layer(conv2, 128, 3, 2)\n resid1 = _residual_block(conv3, 3)\n resid2 = _residual_block(resid1, 3)\n resid3 = _residual_block(resid2, 3)\n resid4 = _residual_block(resid3, 3)\n resid5 = _residual_block(resid4, 3)\n conv_t1 = _conv_tranpose_layer(resid5, 64, 3, 2)\n conv_t2 = _conv_tranpose_layer(conv_t1, 32, 3, 2)\n conv_t3 = _conv_layer(conv_t2, 3, 9, 1, relu=False)\n preds = tf.nn.tanh(conv_t3) * 150 + 255./2\n return preds\n\ndef _conv_layer(net, num_filters, filter_size, strides, relu=True):\n weights_init = _conv_init_vars(net, num_filters, filter_size)\n strides_shape = [1, strides, strides, 1]\n net = tf.nn.conv2d(input=net, filters=weights_init, strides=strides_shape, padding='SAME')\n net = _instance_norm(net)\n if relu:\n net = tf.nn.relu(net)\n\n return net\n\ndef _conv_tranpose_layer(net, num_filters, filter_size, strides):\n weights_init = _conv_init_vars(net, num_filters, filter_size, transpose=True)\n\n batch_size, rows, cols, in_channels = [i for i in net.get_shape()]\n new_rows, new_cols = int(rows * strides), int(cols * strides)\n # new_shape = #tf.pack([tf.shape(net)[0], new_rows, new_cols, num_filters])\n\n new_shape = [batch_size, new_rows, new_cols, num_filters]\n tf_shape = tf.stack(new_shape)\n strides_shape = [1,strides,strides,1]\n\n net = tf.nn.conv2d_transpose(net, weights_init, tf_shape, strides_shape, padding='SAME')\n net = _instance_norm(net)\n return tf.nn.relu(net)\n\ndef _residual_block(net, filter_size=3):\n tmp = _conv_layer(net, 128, filter_size, 1)\n return net + _conv_layer(tmp, 128, filter_size, 1, relu=False)\n\ndef _instance_norm(net, train=True):\n batch, rows, cols, channels = [i for i in net.get_shape()]\n var_shape = [channels]\n mu, sigma_sq = tf.nn.moments(x=net, axes=[1,2], keepdims=True)\n shift = tf.Variable(tf.zeros(var_shape))\n scale = tf.Variable(tf.ones(var_shape))\n epsilon = 1e-3\n normalized = (net-mu)/(sigma_sq + epsilon)**(.5)\n return scale * normalized + shift\n\ndef _conv_init_vars(net, out_channels, filter_size, transpose=False):\n _, rows, cols, in_channels = [i for i in net.get_shape()]\n if not transpose:\n weights_shape = [filter_size, filter_size, in_channels, out_channels]\n else:\n weights_shape = [filter_size, filter_size, out_channels, in_channels]\n\n weights_init = tf.Variable(tf.random.truncated_normal(weights_shape, stddev=WEIGHTS_INIT_STDEV, seed=1), dtype=tf.float32)\n return weights_init\n" }, { "path": "src/utils.py", "content": "import scipy.misc, numpy as np, os, sys\nimport imageio\nfrom PIL import Image\n\ndef save_img(out_path, img):\n img = np.clip(img, 0, 255).astype(np.uint8)\n imageio.imwrite(out_path, img)\n\ndef scale_img(style_path, style_scale):\n scale = float(style_scale)\n o0, o1, o2 = imageio.imread(style_path, pilmode='RGB').shape\n scale = float(style_scale)\n new_shape = (int(o0 * scale), int(o1 * scale), o2)\n style_target = _get_img(style_path, img_size=new_shape)\n return style_target\n\ndef get_img(src, img_size=False):\n img = imageio.imread(src, pilmode='RGB') # misc.imresize(, (256, 256, 3))\n if not (len(img.shape) == 3 and img.shape[2] == 3):\n img = np.dstack((img,img,img))\n if img_size != False:\n img = np.array(Image.fromarray(img).resize(img_size[:2]))\n return img\n\ndef exists(p, msg):\n assert os.path.exists(p), msg\n\ndef list_files(in_path):\n files = []\n for (dirpath, dirnames, filenames) in os.walk(in_path):\n files.extend(filenames)\n break\n\n return files\n" }, { "path": "src/vgg.py", "content": "# Copyright (c) 2015-2016 Anish Athalye. Released under GPLv3.\n\nimport tensorflow as tf\nimport numpy as np\nimport scipy.io\nimport pdb\n\nMEAN_PIXEL = np.array([ 123.68 , 116.779, 103.939])\n\ndef net(data_path, input_image):\n layers = (\n 'conv1_1', 'relu1_1', 'conv1_2', 'relu1_2', 'pool1',\n\n 'conv2_1', 'relu2_1', 'conv2_2', 'relu2_2', 'pool2',\n\n 'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2', 'conv3_3',\n 'relu3_3', 'conv3_4', 'relu3_4', 'pool3',\n\n 'conv4_1', 'relu4_1', 'conv4_2', 'relu4_2', 'conv4_3',\n 'relu4_3', 'conv4_4', 'relu4_4', 'pool4',\n\n 'conv5_1', 'relu5_1', 'conv5_2', 'relu5_2', 'conv5_3',\n 'relu5_3', 'conv5_4', 'relu5_4'\n )\n\n data = scipy.io.loadmat(data_path)\n mean = data['normalization'][0][0][0]\n mean_pixel = np.mean(mean, axis=(0, 1))\n weights = data['layers'][0]\n\n net = {}\n current = input_image\n for i, name in enumerate(layers):\n kind = name[:4]\n if kind == 'conv':\n kernels, bias = weights[i][0][0][0][0]\n # matconvnet: weights are [width, height, in_channels, out_channels]\n # tensorflow: weights are [height, width, in_channels, out_channels]\n kernels = np.transpose(kernels, (1, 0, 2, 3))\n bias = bias.reshape(-1)\n current = _conv_layer(current, kernels, bias)\n elif kind == 'relu':\n current = tf.nn.relu(current)\n elif kind == 'pool':\n current = _pool_layer(current)\n net[name] = current\n\n assert len(net) == len(layers)\n return net\n\n\ndef _conv_layer(input, weights, bias):\n conv = tf.nn.conv2d(input=input, filters=tf.constant(weights), strides=(1, 1, 1, 1),\n padding='SAME')\n return tf.nn.bias_add(conv, bias)\n\n\ndef _pool_layer(input):\n return tf.nn.max_pool2d(input=input, ksize=(1, 2, 2, 1), strides=(1, 2, 2, 1),\n padding='SAME')\n\n\ndef preprocess(image):\n return image - MEAN_PIXEL\n\n\ndef unprocess(image):\n return image + MEAN_PIXEL\n" }, { "path": "style.py", "content": "from __future__ import print_function\nimport sys, os, pdb\nsys.path.insert(0, 'src')\nimport numpy as np, scipy.misc \nfrom optimize import optimize\nfrom argparse import ArgumentParser\nfrom utils import save_img, get_img, exists, list_files\nimport evaluate\n\nCONTENT_WEIGHT = 7.5e0\nSTYLE_WEIGHT = 1e2\nTV_WEIGHT = 2e2\n\nLEARNING_RATE = 1e-3\nNUM_EPOCHS = 2\nCHECKPOINT_DIR = 'checkpoints'\nCHECKPOINT_ITERATIONS = 2000\nVGG_PATH = 'data/imagenet-vgg-verydeep-19.mat'\nTRAIN_PATH = 'data/train2014'\nBATCH_SIZE = 4\nDEVICE = '/gpu:0'\nFRAC_GPU = 1\n\ndef build_parser():\n parser = ArgumentParser()\n parser.add_argument('--checkpoint-dir', type=str,\n dest='checkpoint_dir', help='dir to save checkpoint in',\n metavar='CHECKPOINT_DIR', required=True)\n\n parser.add_argument('--style', type=str,\n dest='style', help='style image path',\n metavar='STYLE', required=True)\n\n parser.add_argument('--train-path', type=str,\n dest='train_path', help='path to training images folder',\n metavar='TRAIN_PATH', default=TRAIN_PATH)\n\n parser.add_argument('--test', type=str,\n dest='test', help='test image path',\n metavar='TEST', default=False)\n\n parser.add_argument('--test-dir', type=str,\n dest='test_dir', help='test image save dir',\n metavar='TEST_DIR', default=False)\n\n parser.add_argument('--slow', dest='slow', action='store_true',\n help='gatys\\' approach (for debugging, not supported)',\n default=False)\n\n parser.add_argument('--epochs', type=int,\n dest='epochs', help='num epochs',\n metavar='EPOCHS', default=NUM_EPOCHS)\n\n parser.add_argument('--batch-size', type=int,\n dest='batch_size', help='batch size',\n metavar='BATCH_SIZE', default=BATCH_SIZE)\n\n parser.add_argument('--checkpoint-iterations', type=int,\n dest='checkpoint_iterations', help='checkpoint frequency',\n metavar='CHECKPOINT_ITERATIONS',\n default=CHECKPOINT_ITERATIONS)\n\n parser.add_argument('--vgg-path', type=str,\n dest='vgg_path',\n help='path to VGG19 network (default %(default)s)',\n metavar='VGG_PATH', default=VGG_PATH)\n\n parser.add_argument('--content-weight', type=float,\n dest='content_weight',\n help='content weight (default %(default)s)',\n metavar='CONTENT_WEIGHT', default=CONTENT_WEIGHT)\n \n parser.add_argument('--style-weight', type=float,\n dest='style_weight',\n help='style weight (default %(default)s)',\n metavar='STYLE_WEIGHT', default=STYLE_WEIGHT)\n\n parser.add_argument('--tv-weight', type=float,\n dest='tv_weight',\n help='total variation regularization weight (default %(default)s)',\n metavar='TV_WEIGHT', default=TV_WEIGHT)\n \n parser.add_argument('--learning-rate', type=float,\n dest='learning_rate',\n help='learning rate (default %(default)s)',\n metavar='LEARNING_RATE', default=LEARNING_RATE)\n\n return parser\n\ndef check_opts(opts):\n exists(opts.checkpoint_dir, \"checkpoint dir not found!\")\n exists(opts.style, \"style path not found!\")\n exists(opts.train_path, \"train path not found!\")\n if opts.test or opts.test_dir:\n exists(opts.test, \"test img not found!\")\n exists(opts.test_dir, \"test directory not found!\")\n exists(opts.vgg_path, \"vgg network data not found!\")\n assert opts.epochs > 0\n assert opts.batch_size > 0\n assert opts.checkpoint_iterations > 0\n assert os.path.exists(opts.vgg_path)\n assert opts.content_weight >= 0\n assert opts.style_weight >= 0\n assert opts.tv_weight >= 0\n assert opts.learning_rate >= 0\n\ndef _get_files(img_dir):\n files = list_files(img_dir)\n return [os.path.join(img_dir,x) for x in files]\n\n \ndef main():\n parser = build_parser()\n options = parser.parse_args()\n check_opts(options)\n\n style_target = get_img(options.style)\n if not options.slow:\n content_targets = _get_files(options.train_path)\n elif options.test:\n content_targets = [options.test]\n\n kwargs = {\n \"slow\":options.slow,\n \"epochs\":options.epochs,\n \"print_iterations\":options.checkpoint_iterations,\n \"batch_size\":options.batch_size,\n \"save_path\":os.path.join(options.checkpoint_dir,'fns.ckpt'),\n \"learning_rate\":options.learning_rate\n }\n\n if options.slow:\n if options.epochs < 10:\n kwargs['epochs'] = 1000\n if options.learning_rate < 1:\n kwargs['learning_rate'] = 1e1\n\n args = [\n content_targets,\n style_target,\n options.content_weight,\n options.style_weight,\n options.tv_weight,\n options.vgg_path\n ]\n\n for preds, losses, i, epoch in optimize(*args, **kwargs):\n style_loss, content_loss, tv_loss, loss = losses\n\n print('Epoch %d, Iteration: %d, Loss: %s' % (epoch, i, loss))\n to_print = (style_loss, content_loss, tv_loss)\n print('style: %s, content:%s, tv: %s' % to_print)\n if options.test:\n assert options.test_dir != False\n preds_path = '%s/%s_%s.png' % (options.test_dir,epoch,i)\n if not options.slow:\n ckpt_dir = os.path.dirname(options.checkpoint_dir)\n evaluate.ffwd_to_img(options.test,preds_path,\n options.checkpoint_dir)\n else:\n save_img(preds_path, img)\n ckpt_dir = options.checkpoint_dir\n cmd_text = 'python evaluate.py --checkpoint %s ...' % ckpt_dir\n print(\"Training complete. For evaluation:\\n `%s`\" % cmd_text)\n\nif __name__ == '__main__':\n main()\n" }, { "path": "transform_video.py", "content": "from __future__ import print_function\nfrom argparse import ArgumentParser\nimport sys\nsys.path.insert(0, 'src')\nimport os, random, subprocess, evaluate, shutil\nfrom utils import exists, list_files\nimport pdb\n\nTMP_DIR = '.fns_frames_%s/' % random.randint(0,99999)\nDEVICE = '/gpu:0'\nBATCH_SIZE = 4\n\ndef build_parser():\n parser = ArgumentParser()\n parser.add_argument('--checkpoint', type=str,\n dest='checkpoint', help='checkpoint directory or .ckpt file',\n metavar='CHECKPOINT', required=True)\n\n parser.add_argument('--in-path', type=str,\n dest='in_path', help='in video path',\n metavar='IN_PATH', required=True)\n \n parser.add_argument('--out-path', type=str,\n dest='out', help='path to save processed video to',\n metavar='OUT', required=True)\n \n parser.add_argument('--tmp-dir', type=str, dest='tmp_dir',\n help='tmp dir for processing', metavar='TMP_DIR',\n default=TMP_DIR)\n\n parser.add_argument('--device', type=str, dest='device',\n help='device for eval. CPU discouraged. ex: \\'/gpu:0\\'',\n metavar='DEVICE', default=DEVICE)\n\n parser.add_argument('--batch-size', type=int,\n dest='batch_size',help='batch size for eval. default 4.',\n metavar='BATCH_SIZE', default=BATCH_SIZE)\n\n parser.add_argument('--no-disk', type=bool, dest='no_disk',\n help='Don\\'t save intermediate files to disk. Default False',\n metavar='NO_DISK', default=False)\n return parser\n\ndef check_opts(opts):\n exists(opts.checkpoint)\n exists(opts.out)\n\ndef main():\n parser = build_parser()\n opts = parser.parse_args()\n evaluate.ffwd_video(opts.in_path, opts.out, opts.checkpoint, opts.device, opts.batch_size)\n\n \nif __name__ == '__main__':\n main()\n\n\n" } ]