diff --git a/neuralstyle/algorithms.py b/neuralstyle/algorithms.py index 5412ce8..7ff810f 100644 --- a/neuralstyle/algorithms.py +++ b/neuralstyle/algorithms.py @@ -85,17 +85,21 @@ def styletransfer_single(content, style, outfile, size=None, alg="gatys", weight """General style transfer routine over a single set of options""" workdir = TemporaryDirectory() - # Cut out alpha channel - rgbfile = workdir.name + "rgb.png" - alphafile = workdir.name + "alpha.png" + # 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" + algfile = workdir.name + "/" + "algoutput.png" if alg == "gatys": - gatys(rgbfile, style, algfile, size, weight, stylescale, algparams) + gatys(rgbfile, stylepng, algfile, size, weight, stylescale, algparams) elif alg in ["chen-schmidt", "chen-schmidt-inverse"]: - chenschmidt(alg, rgbfile, style, algfile, size, stylescale, algparams) + chenschmidt(alg, rgbfile, stylepng, algfile, size, stylescale, algparams) # Enforce correct size correctshape(algfile, content, size) @@ -118,9 +122,9 @@ def neuraltile(content, style, outfile, size=None, maxtilesize=400, overlap=50, # Compute number of tiles required to map all the image xtiles, ytiles = tilegeometry(fullshape, maxtilesize, overlap) - # First scale image to target resolution. Also transform to PNG to avoid alpha channel operations + # First scale image to target resolution firstpass = workdir.name + "/" + "lowres.png" - copyfile(content, firstpass) + convert(content, firstpass) resize(firstpass, fullshape) # Chop the styled image into tiles with the specified overlap value. diff --git a/neuralstyle/imagemagick.py b/neuralstyle/imagemagick.py index e17d690..e250973 100644 --- a/neuralstyle/imagemagick.py +++ b/neuralstyle/imagemagick.py @@ -1,18 +1,28 @@ # Convenience functions to perform Image Magicks from subprocess import run, PIPE from glob import glob +from tempfile import TemporaryDirectory from neuralstyle.utils import filename def convert(origin, dest): - """Transforms the format of an image in a file, by creating a new file with the new format""" - command = "convert " + origin + " " + dest + """Transforms the format of an image in a file, by creating a new file with the new format + + If the input file has several layers, they are flattened. + """ + command = "convert " + origin + " -flatten " + dest run(command, shell=True) def shape(imfile): - """Returns the shape of an image file""" - result = run("convert " + imfile + ' -format "%w %h" info:', shell=True, check=True, stdout=PIPE) + """Returns the shape of an image file + + If the input file has several layers, it is flattened before computing the shape. + """ + tmpdir = TemporaryDirectory() + tmpname = tmpdir.name + "/" + "image.png" + convert(imfile, tmpname) + result = run("convert " + tmpname + ' -format "%w %h" info:', shell=True, check=True, stdout=PIPE) return [int(x) for x in result.stdout.decode("utf-8").split(" ")] diff --git a/tests/algorithms_tests.py b/tests/algorithms_tests.py index 285b6be..7c5da32 100644 --- a/tests/algorithms_tests.py +++ b/tests/algorithms_tests.py @@ -108,3 +108,19 @@ def test_neuraltile(): outfile = tmpdir.name + "/tiled.png" neuraltile(content, STYLES + "cubism.jpg", outfile, alg="chen-schmidt-inverse", maxtilesize=400) assert shape(outfile) == shape(content) + + +def test_formattga(): + """TGA format images can be processed correctly""" + contents = [CONTENTS + f for f in ["tgasample.tga", "marbles.tga"]] + tmpdir = TemporaryDirectory() + styletransfer(contents, [STYLES + "cubism.jpg"], tmpdir.name, alg="chen-schmidt-inverse") + assert len(glob(tmpdir.name + "/*cubism*")) == 2 + + +def test_formatpsd(): + """PSD format images can be processed correctly""" + contents = [CONTENTS + f for f in ["oldtelephone.psd"]] + tmpdir = TemporaryDirectory() + styletransfer(contents, [STYLES + "cubism.jpg"], tmpdir.name, alg="chen-schmidt-inverse") + assert len(glob(tmpdir.name + "/*cubism*")) == 1 diff --git a/tests/imagemagick_tests.py b/tests/imagemagick_tests.py index dc2d33b..4eb7076 100644 --- a/tests/imagemagick_tests.py +++ b/tests/imagemagick_tests.py @@ -3,7 +3,8 @@ # from tempfile import TemporaryDirectory from shutil import copyfile -from neuralstyle.imagemagick import shape, resize, choptiles, feather, extractalpha, mergealpha, equalimages +from glob import glob +from neuralstyle.imagemagick import shape, resize, choptiles, feather, extractalpha, mergealpha, equalimages, convert from neuralstyle.utils import filename CONTENTS = "/app/entrypoint/tests/contents/" @@ -35,6 +36,28 @@ def test_shape(): assert result == expected +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_convert_layers(): + """Convert a single image with layers works as expected""" + for content in [CONTENTS + f for f in ["oldtelephone.psd"]]: + 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_resize_keepproportions(): """Resizing an image without changing proportions works correctly""" tmpdir = TemporaryDirectory()