def chenschmidt(alg, content, style, outfile, size, stylescale, algparams):
"""Runs Chen and Schmidt fast style-transfer algorithm
References:
* https://arxiv.org/pdf/1612.04337.pdf
* https://github.com/rtqichen/style-swap
"""
if alg not in ["chen-schmidt", "chen-schmidt-inverse"]:
raise ValueError(
"Unnaceptable subalgorithm %s for Chen-Schmidt family")
# Rescale style as requested
instyle = NamedTemporaryFile()
copyfile(style, instyle.name)
resize(instyle.name, int(stylescale * shape(style)[0]))
# Run algorithm
outdir = TemporaryDirectory()
runalgorithm(alg, [
"--save", outdir.name, "--content", content, "--style", instyle.name,
"--maxContentSize", size if size is not None else shape(content)[0],
"--maxStyleSize", size if size is not None else shape(content)[0],
*algparams
])
# Gather output results
output = outdir.name + "/" + filename(content) + "_stylized" + fileext(
content)
convert(output, outfile)
instyle.close()
def gatys(content, style, outfile, size, weight, stylescale, algparams):
"""Runs Gatys et al style-transfer algorithm
References:
* https://arxiv.org/abs/1508.06576
* https://github.com/jcjohnson/neural-style
"""
# Gatys can only process one combination of content, style, weight and scale at a time, so we need to iterate
tmpout = NamedTemporaryFile(suffix=".png")
runalgorithm(
"gatys",
[
"-content_image",
content,
"-style_image",
style,
"-style_weight",
weight * 100, # Because content weight is 100
"-style_scale",
stylescale,
"-output_image",
tmpout.name,
"-image_size",
size if size is not None else shape(content)[0],
*algparams
])
# Transform to original file format
convert(tmpout.name, outfile)
tmpout.close()
def test_convert_nolayers():
"""Convert a single image with no layers works as expected"""
for content in [CONTENTS + f for f in ["docker.png", "goldengate.jpg"]]:
for ext in [".png", ".jpg", ".psd", ".tga"]:
tmpdir = TemporaryDirectory()
outname = tmpdir.name + "/" + "output" + ext
convert(content, outname)
assert len(glob(tmpdir.name + "/" + filename(outname) + ext)) == 1
assert shape(outname) == shape(content)
def styletransfer_single(content,
style,
outfile,
size=None,
alg="gatys",
weight=5.0,
stylescale=1.0,
algparams=None):
"""General style transfer routine over a single set of options"""
workdir = TemporaryDirectory()
# Cut out alpha channel from content
rgbfile = workdir.name + "/" + "rgb.png"
alphafile = workdir.name + "/" + "alpha.png"
extractalpha(content, rgbfile, alphafile)
# Transform style to png, as some algorithms don't understand other formats
stylepng = workdir.name + "/" + "style.png"
convert(style, stylepng)
# Call style transfer algorithm
algfile = workdir.name + "/" + "algoutput.png"
if alg == "gatys":
gatys(rgbfile, stylepng, algfile, size, weight, stylescale, algparams)
elif alg == "gatys-multiresolution":
gatys_multiresolution(rgbfile, stylepng, algfile, size, weight,
stylescale, algparams)
elif alg in ["chen-schmidt", "chen-schmidt-inverse"]:
chenschmidt(alg, rgbfile, stylepng, algfile, size, stylescale,
algparams)
# Enforce correct size
correctshape(algfile, content, size)
# Recover alpha channel
correctshape(alphafile, content, size)
mergealpha(algfile, alphafile, outfile)
def gatys_multiresolution(content,
style,
outfile,
size,
weight,
stylescale,
algparams,
startres=256):
"""Runs a multiresolution version of Gatys et al method
The multiresolution strategy starts by generating a small image, then using that image as initializer
for higher resolution images. This procedure is repeated up to the tilesize.
Once the maximum tile size attainable by L-BFGS is reached, more iterations are run by using Adam. This allows
to produce larger images using this method than the basic Gatys.
References:
* Gatys et al - Controlling Perceptual Factors in Neural Style Transfer (https://arxiv.org/abs/1611.07865)
* https://gist.github.com/jcjohnson/ca1f29057a187bc7721a3a8c418cc7db
"""
# Multiresolution strategy: list of rounds, each round composed of a optimization method and a number of
# upresolution steps.
# Using "adam" as optimizer means that Adam will be used when necessary to attain higher resolutions
strategy = [["lbfgs", 7], ["lbfgs", 7], ["lbfgs", 7], ["lbfgs", 7],
["lbfgs", 7]]
LOGGER.info("Starting gatys-multiresolution with strategy " +
str(strategy))
# Initialization
workdir = TemporaryDirectory()
maxres = targetshape(content, size)[0]
if maxres < startres:
LOGGER.warning(
"Target resolution (%d) might too small for the multiresolution method to work well"
% maxres)
startres = maxres / 2.0
seed = None
tmpout = workdir.name + "/tmpout.png"
# Iterate over rounds
for roundnumber, (optimizer, steps) in enumerate(strategy):
LOGGER.info(
"gatys-multiresolution round %d with %s optimizer and %d steps" %
(roundnumber, optimizer, steps))
roundmax = min(maxtile("gatys"),
maxres) if optimizer == "lbfgs" else maxres
resolutions = np.linspace(startres, roundmax, steps, dtype=int)
iters = 1000
for stepnumber, res in enumerate(resolutions):
stepopt = "adam" if res > maxtile("gatys") else "lbfgs"
LOGGER.info("Step %d, resolution %d, optimizer %s" %
(stepnumber, res, stepopt))
passparams = algparams[:]
passparams.extend([
"-num_iterations", iters, "-tv_weight", "0", "-print_iter",
"0", "-optimizer", stepopt
])
if seed is not None:
passparams.extend(["-init", "image", "-init_image", seed])
gatys(content, style, tmpout, res, weight, stylescale, passparams)
seed = workdir.name + "/seed.png"
copyfile(tmpout, seed)
iters = max(iters / 2.0, 100)
convert(tmpout, outfile)
def neuraltile(content,
style,
outfile,
size=None,
overlap=100,
alg="gatys",
weight=5.0,
stylescale=1.0,
algparams=None):
"""Strategy to generate a high resolution image by running style transfer on overlapping image tiles"""
LOGGER.info("Starting tiling strategy")
if algparams is None:
algparams = []
workdir = TemporaryDirectory()
# Gather size info from original image
fullshape = targetshape(content, size)
# Compute number of tiles required to map all the image
xtiles, ytiles = tilegeometry(fullshape, alg, overlap)
# First scale image to target resolution
firstpass = workdir.name + "/" + "lowres.png"
convert(content, firstpass)
resize(firstpass, fullshape)
# Chop the styled image into tiles with the specified overlap value.
lowrestiles = choptiles(firstpass,
xtiles=xtiles,
ytiles=ytiles,
overlap=overlap,
outname=workdir.name + "/" + "lowres_tiles")
# High resolution pass over each tile
highrestiles = []
for i, tile in enumerate(lowrestiles):
name = workdir.name + "/" + "highres_tiles_" + str(i) + ".png"
styletransfer_single(tile,
style,
name,
size=None,
alg=alg,
weight=weight,
stylescale=stylescale,
algparams=algparams)
highrestiles.append(name)
# Feather tiles
featheredtiles = []
for i, tile in enumerate(highrestiles):
name = workdir.name + "/" + "feathered_tiles_" + str(i) + ".png"
feather(tile, name)
featheredtiles.append(name)
# Smush the feathered tiles together
smushedfeathered = workdir.name + "/" + "feathered_smushed.png"
smush(featheredtiles, xtiles, ytiles, overlap, overlap, smushedfeathered)
# Smush also the non-feathered tiles
smushedhighres = workdir.name + "/" + "highres_smushed.png"
smush(highrestiles, xtiles, ytiles, overlap, overlap, smushedhighres)
# Combine feathered and un-feathered output images to disguise feathering
composite([smushedfeathered, smushedhighres], outfile)
# Adjust back to desired size
assertshape(outfile, fullshape)