Merge pull request #15 from albarji/feature/moreimageformats

tga and psd image processing

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.

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(" ")]
 
 

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

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()