Skip to content

Commit

Permalink
More cleaning up and adding introduction text.
Browse files Browse the repository at this point in the history
  • Loading branch information
@alexjc
alexjc committed Mar 8, 2016
1 parent 6901716 commit 2cfa409
Showing 1 changed file with 54 additions and 28 deletions.
82 changes: 54 additions & 28 deletions doodle.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,9 @@
#
# Copyright (c) 2016, Alex J. Champandard.
#
# Research and development sponsored by the nucl.ai Conference!
# http://events.nucl.ai/
# July 18-20, 2016 in Vienna/Austria.
#

import os
Expand All @@ -8,10 +12,6 @@
import pickle
import argparse

import numpy as np
import scipy.optimize
import skimage.transform


# Configure all options first so we can custom load other libraries (Theano) based on device specified by user.
parser = argparse.ArgumentParser(description='Generate a new image by applying style onto a content image.',
Expand All @@ -37,30 +37,47 @@
args = parser.parse_args()


#----------------------------------------------------------------------------------------------------------------------

# Color coded output helps visualize the information a little better, plus looks cool!
class ansi:
BOLD = '\033[1;97m'
WHITE = '\033[0;97m'
YELLOW = '\033[0;33m'
RED = '\033[0;31m'
RED_B = '\033[1;31m'
BLUE = '\033[0;94m'
BLUE_B = '\033[1;94m'
CYAN = '\033[0;36m'
CYAN_B = '\033[1;36m'
ENDC = '\033[0m'

print('{}Neural Doodle for semantic style transfer.{}'.format(ansi.CYAN_B, ansi.ENDC))


# Load the underlying deep learning libraries based on the device specified. If you specify THEANO_FLAGS manually,
# the code assumes you know what you are doing and they are not overriden!
os.environ.setdefault('THEANO_FLAGS', 'device=%s,floatX=float32,print_active_device=False' % (args.device))

# Scientific Libraries
import numpy as np
import scipy.optimize
import skimage.transform

# Numeric Computing (GPU)
import theano
import theano.tensor as T
import theano.tensor.nnet.neighbours

# Deep Learning Framework
import lasagne
from lasagne.layers import Conv2DLayer as ConvLayer, Pool2DLayer as PoolLayer
from lasagne.layers import InputLayer, ConcatLayer


class ansi:
BOLD = '\033[1;97m'
WHITE = '\033[0;97m'
YELLOW = '\033[0;33m'
RED = '\033[0;31m'
RED_B = '\033[1;31m'
BLUE = '\033[0;94m'
BLUE_B = '\033[1;94m'
ENDC = '\033[0m'


#----------------------------------------------------------------------------------------------------------------------
# Convolutional Neural Network
#----------------------------------------------------------------------------------------------------------------------
class Model(object):
"""Store all the data related to the neural network (aka. "model"). This is currently based on VGG19.
"""
Expand Down Expand Up @@ -156,6 +173,9 @@ def finalize_image(self, image, resolution):
return np.clip(image, 0, 255).astype('uint8')


#----------------------------------------------------------------------------------------------------------------------
# Semantic Style Transfer
#----------------------------------------------------------------------------------------------------------------------
class NeuralGenerator(object):
"""This is the main part of the application that generates an image using optimization and LBFGS.
The images will be processed at increasing resolutions in the run() method.
Expand All @@ -170,37 +190,43 @@ def __init__(self):
if args.output is not None and os.path.isfile(args.output):
os.remove(args.output)

self.style_image_original, self.style_map_original = self.load_images(args.style)
self.content_image_original, self.content_map_original = self.load_images(args.content or args.output)
print(ansi.CYAN, end='')
target = args.content or args.output
self.content_img_original, self.content_map_original = self.load_images('content', target)
self.style_img_original, self.style_map_original = self.load_images('style', args.style)
print(ansi.ENDC, end='')

if self.content_map_original is None:
self.content_map_original = np.zeros(self.content_image_original.shape[:2]+(1,))
self.content_map_original = np.zeros(self.content_img_original.shape[:2]+(1,))
self.semantic_weight = 0.0

if self.content_image_original is None:
self.content_image_original = np.zeros(self.content_map_original.shape[:2]+(3,))
if self.content_img_original is None:
self.content_img_original = np.zeros(self.content_map_original.shape[:2]+(3,))
args.content_weight = 0.0

assert self.content_map_original.shape[2] == self.style_map_original.shape[2],\
"Mismatch in number of channels for style and content semantic map.s"

def load_images(self, filename):
"""If the image and mask files exist, load them. Otherwise they'll be set to default values later.
def load_images(self, name, filename):
"""If the image and map files exist, load them. Otherwise they'll be set to default values later.
"""
basename, _ = os.path.splitext(filename)
mapname = basename + args.semantic_ext
img = scipy.ndimage.imread(filename, mode='RGB') if os.path.exists(filename) else None
mask = scipy.ndimage.imread(mapname) if os.path.exists(mapname) else None
return img, mask
map = scipy.ndimage.imread(mapname) if os.path.exists(mapname) else None

if img is not None: print(' - Loading {} image data from {}.'.format(name, filename))
if map is not None: print(' - Loading {} semantic map from {}.'.format(name, mapname))
return img, map

#------------------------------------------------------------------------------------------------------------------
# Initialization & Setup
#------------------------------------------------------------------------------------------------------------------

def prepare_content(self, scale=1.0):
"""Called each phase of the optimization, rescale the original content image and its mask to use as inputs.
"""Called each phase of the optimization, rescale the original content image and its map to use as inputs.
"""
content_image = skimage.transform.rescale(self.content_image_original, scale) * 255.0
content_image = skimage.transform.rescale(self.content_img_original, scale) * 255.0
self.content_image = self.model.prepare_image(content_image)

content_map = skimage.transform.rescale(self.content_map_original * args.semantic_weight, scale) * 255.0
Expand All @@ -210,7 +236,7 @@ def prepare_style(self, scale=1.0):
"""Called each phase of the optimization, process the style image according to the scale, then run it
through the model to extract intermediate outputs (e.g. sem4_1) and turn them into patches.
"""
style_image = skimage.transform.rescale(self.style_image_original, scale) * 255.0
style_image = skimage.transform.rescale(self.style_img_original, scale) * 255.0
self.style_image = self.model.prepare_image(style_image)

style_map = skimage.transform.rescale(self.style_map_original * args.semantic_weight, scale) * 255.0
Expand Down Expand Up @@ -374,7 +400,7 @@ def run(self):
self.error = 255.0
scale = 1.0 / 2.0 ** (args.resolutions - 1 - i)

shape = self.content_image_original.shape
shape = self.content_img_original.shape
print('\n{}Phase #{}: resolution {}x{} scale {}{}'\
.format(ansi.BLUE_B, i, int(shape[1]*scale), int(shape[0]*scale), scale, ansi.BLUE))

Expand Down

0 comments on commit 2cfa409

Please sign in to comment.