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 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 test_styletransfer_keepsize():
"""Style transfer keeps the original image size if no size paramenter is given"""
for alg in ALGORITHMS.keys():
tmpdir = TemporaryDirectory()
img = CONTENTS + "dockersmall.png"
styletransfer([img], [STYLES + "cubism.jpg"], tmpdir.name, alg=alg)
files = glob(tmpdir.name + "/" + filename(img) + "*cubism*")
print("Expected size", shape(img))
print("Actual shape", shape(files[0]))
assert len(files) == 1
assert shape(files[0]) == shape(img)
def test_neuraltile():
"""The neural tiling procedure can be run without issues"""
tmpdir = TemporaryDirectory()
content = CONTENTS + "avila-walls.jpg"
outfile = tmpdir.name + "/tiled.png"
neuraltile(content,
STYLES + "cubism.jpg",
outfile,
alg="chen-schmidt-inverse",
overlap=100)
assert shape(outfile) == shape(content)
def test_extractalpha():
"""Extracting the alpha channel from an image works as expected"""
tmpdir = TemporaryDirectory()
img = CONTENTS + "/alphasample.png"
alphafile = tmpdir.name + "/alpha.png"
rgbfile = tmpdir.name + "/rgb.png"
extractalpha(img, rgbfile, alphafile)
assert shape(rgbfile) == shape(img)
assert shape(alphafile) == shape(img)
assert not equalimages(img, rgbfile)
assert not equalimages(img, alphafile)
assert not equalimages(rgbfile, alphafile)
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_resize_changeproportions():
"""Resizing an image changing proportions works correctly"""
tmpdir = TemporaryDirectory()
fname = tmpdir.name + "/docker.png"
copyfile(CONTENTS + "docker.png", fname)
resize(fname, [700, 300])
assert shape(fname) == [700, 300]
def test_resize_keepproportions():
"""Resizing an image without changing proportions works correctly"""
tmpdir = TemporaryDirectory()
fname = tmpdir.name + "/docker.png"
copyfile(CONTENTS + "docker.png", fname)
resize(fname, 1016)
assert shape(fname) == [1016, 886]
def targetshape(content, size=None):
"""Computes the shape the resultant image will have after a reshape of the size given
If size is None, return original shape.
"""
contentshape = shape(content)
if size is None:
return contentshape
else:
return [size, int(size * contentshape[1] / contentshape[0])]
def test_styletransfer_size():
"""Style transfer works for varying image sizes, producing correctly scaled images"""
for alg in ALGORITHMS.keys():
for size in [50, 100, 200]:
for img in ["docker.png", "obama.jpg"]:
originalshape = shape(CONTENTS + img)
tmpdir = TemporaryDirectory()
styletransfer([CONTENTS + img], [STYLES + "cubism.jpg"],
tmpdir.name,
alg=alg,
size=size)
files = glob(tmpdir.name + "/" + filename(img) + "*cubism*")
resultshape = shape(files[0])
rescalefactor = size / originalshape[0]
expectedshape = [size, int(rescalefactor * originalshape[1])]
print("Expected shape", expectedshape)
print("Actual shape", resultshape)
assert len(files) == 1
assert expectedshape == resultshape
def test_shape():
"""The shape of an image can be correctly recovered"""
tests = [ # Inputs, expected outputs
(CONTENTS + "docker.png", [508, 443]),
(CONTENTS + "goldengate.jpg", [1920, 1080])
]
for imfile, expected in tests:
result = shape(imfile)
print("Input", imfile)
print("Expected", expected)
print("Output", result)
assert result == expected
