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 choptiles(imfile, xtiles, ytiles, overlap, outname):
"""Chops an image file into a geometry of overlapping tiles. Returns ordered list of generated tiles image files"""
command = 'convert %s -crop %dx%d+%d+%d@ +repage +adjoin %s_%%d.png' % (
imfile, xtiles, ytiles, overlap, overlap, outname)
run(command, shell=True, check=True)
return sorted(glob(outname + "_*.png"),
key=lambda x: int(filename(x).split("_")[-1]))
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_choptiles():
"""Chopping an image into tiles works as expected"""
tmpdir = TemporaryDirectory()
content = CONTENTS + "/goldengate.jpg"
tiles = choptiles(content,
xtiles=2,
ytiles=3,
overlap=50,
outname=tmpdir.name + "/tiles")
print("Generated tiles", tiles)
assert len(tiles) == 6
for i, tile in enumerate(tiles):
assert int(filename(tile).split("_")[-1]) == i
def test_styletransfer_sw():
"""Style transfer works for varying style weights"""
styleweights = [1, 5, 10]
alg = "gatys"
img = "docker.png"
tmpdir = TemporaryDirectory()
styletransfer([CONTENTS + img], [STYLES + "cubism.jpg"],
tmpdir.name,
alg=alg,
size=100,
weights=styleweights)
assertalldifferent(tmpdir.name + "/" + filename(img) + "*cubism*",
len(styleweights))
def test_styletransfer_ss():
"""Style transfer works for varying style scales"""
stylescales = [0.75, 1, 1.25]
for alg in ALGORITHMS.keys():
img = "docker.png"
tmpdir = TemporaryDirectory()
styletransfer([CONTENTS + img], [STYLES + "cubism.jpg"],
tmpdir.name,
alg=alg,
size=100,
stylescales=stylescales)
assertalldifferent(tmpdir.name + "/" + filename(img) + "*cubism*",
len(stylescales))
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 outname(savefolder, content, style, alg, scale, weight=None, ext=None):
"""Creates an output filename that reflects the style transfer parameters"""
return (savefolder + "/" + filename(content) + "_" + filename(style) +
"_" + alg + "_ss" + str(scale) +
("_sw" + str(weight) if weight is not None else "") +
(ext if ext is not None else fileext(content)))