Commits

Mikhail Korobov  committed 6521f15

Layer.as_pymaging() method

  • Participants
  • Parent commits 9d6daa0

Comments (0)

Files changed (4)

File src/psd_tools/user_api/psd_image.py

         """ Returns a PIL image for this layer. """
         return self._psd._layer_as_PIL(self._index)
 
+    def as_pymaging(self):
+        """ Returns a pymaging.Image for this PSD file. """
+        return self._psd._layer_as_pymaging(self._index)
+
     @property
     def bbox(self):
         """ BBox(x1, y1, x2, y2) namedtuple with layer bounding box. """
 
     def as_pymaging(self):
         """
-        Retirns a pymaging.Image for this PSD file.
+        Returns a pymaging.Image for this PSD file.
         """
         return pymaging_support.extract_composite_image(self.decoded_data)
 
     def _layer_as_PIL(self, index):
         return layer_to_PIL(self.decoded_data, index)
 
+    def _layer_as_pymaging(self, index):
+        return pymaging_support.extract_layer_image(self.decoded_data, index)
+
 
 def _combined_bbox(layers):
     """

File src/psd_tools/user_api/pymaging_support.py

     Image = None
     packbits = None
 
-from psd_tools.constants import ColorMode, Compression
+from psd_tools.constants import ColorMode, Compression, ChannelID
+
 
 def extract_composite_image(decoded_data):
     """
     Converts a composite (merged) image from the ``decoded_data``
     to a pymaging.Image.
     """
-    if Image is None or packbits is None:
-        raise Exception("This module requires `pymaging` and `packbits` packages.")
-
     header = decoded_data.header
     size = header.width, header.height
+    depth, mode = _validate_header(header)
 
-    if header.color_mode != ColorMode.RGB:
-        raise NotImplementedError(
-            "This color mode (%s) is not supported yet" % ColorMode.name_of(header.color_mode)
-        )
+    return _channels_data_to_image(decoded_data.image_data, mode, size, depth)
 
-    if header.depth != 8:
-        raise NotImplementedError("Only 8bit images are currently supported with pymaging.")
+def extract_layer_image(decoded_data, layer_index):
+    """
+    Converts a layer from the ``decoded_data`` to a ``pymaging.Image``.
+    """
+    layers = decoded_data.layer_and_mask_data.layers
+    layer = layers.layer_records[layer_index]
 
-    if header.number_of_channels == 3:
-        mode = RGB
-    elif header.number_of_channels == 4:
-        mode = RGBA
-    else:
-        raise NotImplementedError("This number of channels (%d) is unsupported for this color mode (%s)" % (
-                         header.number_of_channels, header.color_mode))
+    channels_data = layers.channel_image_data[layer_index]
+    channel_types = [info.id for info in layer.channels]
+    size = layer.width(), layer.height()
 
-    return _channels_data_to_image(decoded_data.image_data,
-        mode,
-        size,
-        header.depth,
-    )
+    depth, _ = _validate_header(decoded_data.header)
+
+    # FIXME: support for layers with mask (there would be 5 channels in this case)
+    if channel_types[0] == ChannelID.TRANSPARENCY_MASK:
+        # move alpha channel to the end
+        channels_data = [channels_data[i] for i in [1, 2, 3, 0]]
+
+    print(layer.channels)
+    mode = _get_mode(len(channels_data))
+
+    return _channels_data_to_image(channels_data, mode, size, depth)
 
 
 def _channels_data_to_image(channels_data, mode, size, depth):
 
     return Image(pixels, mode)
 
+
+def _get_mode(number_of_channels):
+    mode = None
+    if number_of_channels == 3:
+        mode = RGB
+    elif number_of_channels == 4:
+        mode = RGBA
+    return mode
+
+
+def _validate_header(header):
+    """
+    Validates header and returns (depth, mode) tuple.
+    """
+    if Image is None or packbits is None:
+        raise Exception("This module requires `pymaging` and `packbits` packages.")
+
+    if header.color_mode != ColorMode.RGB:
+        raise NotImplementedError(
+            "This color mode (%s) is not supported yet" % ColorMode.name_of(header.color_mode)
+        )
+
+    mode = _get_mode(header.number_of_channels)
+    if mode is None:
+        raise NotImplementedError("This number of channels (%d) is unsupported for this color mode (%s)" % (
+                         header.number_of_channels, header.color_mode))
+
+    if header.depth != 8:
+        raise NotImplementedError("Only 8bit images are currently supported with pymaging.")
+
+
+
+    return 8, mode
+

File tests/test_pixels.py

 
 from psd_tools import PSDImage, Layer, Group
 
-from .utils import decode_psd
+from .utils import full_name
 
 PIXEL_COLORS = (
     # filename                  probe point    pixel value
 )
 
 LAYER_COLORS = (
+    ('1layer.psd',  0,  (5, 5),       (0x27, 0xBA, 0x0F)),
+    ('2layers.psd', 1,  (5, 5),       (0x27, 0xBA, 0x0F)),
+    ('2layers.psd', 1,  (70, 30),     (0x27, 0xBA, 0x0F)),
+    ('2layers.psd', 0,  (0, 0),       (0, 0, 0, 0)),
+    ('2layers.psd', 0,  (62, 26),     (0xF2, 0xF4, 0xC2, 0xFE)),
+)
+
+LAYER_COLORS_MULTIBYTE = (
     ('16bit5x5.psd', 1, (0, 0), (236, 242, 251, 255)),
     ('16bit5x5.psd', 1, (1, 3), (46, 196, 104, 255)),
     ('32bit5x5.psd', 1, (0, 0), (235, 241, 250, 255)), # why not equal to 16bit5x5.psd?
     ('32bit5x5.psd', 1, (1, 3), (46, 196, 104, 255)),
 )
 
-def _assert_PIL_image_pixel(filename, point, color):
-    psd = PSDImage(decode_psd(filename))
-    image = psd.as_PIL()
-    assert image.getpixel(point) == color
+def color_PIL(psd, point):
+    im = psd.as_PIL()
+    return im.getpixel(point)
 
-def _assert_pymaging_image_pixel(filename, point, color):
-    psd = PSDImage(decode_psd(filename))
-    image = psd.as_pymaging()
-    assert list(image.get_pixel(*point)) == list(color)
+def color_pymaging(psd, point):
+    im = psd.as_pymaging()
+    return tuple(im.get_pixel(*point))
 
-BACKENDS = [[_assert_PIL_image_pixel], [_assert_pymaging_image_pixel]]
+BACKENDS = [[color_PIL], [color_pymaging]]
 
 
-@pytest.mark.parametrize(["assert_function"], BACKENDS)
+@pytest.mark.parametrize(["get_color"], BACKENDS)
 @pytest.mark.parametrize(["filename", "point", "color"], PIXEL_COLORS)
-def test_composite(filename, point, color, assert_function):
-    assert_function(filename, point, color)
+def test_composite(filename, point, color, get_color):
+    psd = PSDImage.load(full_name(filename))
+    assert color == get_color(psd, point)
 
 @pytest.mark.parametrize(["filename", "point", "color"], PIXEL_COLORS_32BIT)
 def test_composite_32bit(filename, point, color):
-    _assert_PIL_image_pixel(filename, point, color)
+    psd = PSDImage.load(full_name(filename))
+    assert color == color_PIL(psd, point)
 
 @pytest.mark.parametrize(["filename", "point", "color"], PIXEL_COLORS_16BIT)
 def test_composite_16bit(filename, point, color):
-    _assert_PIL_image_pixel(filename, point, color)
+    psd = PSDImage.load(full_name(filename))
+    assert color == color_PIL(psd, point)
 
+@pytest.mark.parametrize(["filename", "layer_num", "point", "color"], LAYER_COLORS_MULTIBYTE)
+def test_layer_colors_multibyte(filename, layer_num, point, color):
+    psd = PSDImage.load(full_name(filename))
+    layer = psd.layers[layer_num]
+    assert color == color_PIL(layer, point)
+
+
+@pytest.mark.parametrize(["get_color"], BACKENDS)
 @pytest.mark.parametrize(["filename", "layer_num", "point", "color"], LAYER_COLORS)
-def test_layer_colors(filename, layer_num, point, color):
-    psd = PSDImage(decode_psd(filename))
+def test_layer_colors(filename, layer_num, point, color, get_color):
+    psd = PSDImage.load(full_name(filename))
     layer = psd.layers[layer_num]
-    image = layer.as_PIL()
-    assert image.getpixel(point) == color
+    assert color == get_color(layer, point)

File tests/utils.py

 
 DATA_PATH = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'psd_files')
 
+def full_name(filename):
+    return os.path.join(DATA_PATH, filename)
+
 def load_psd(filename):
-    full_filename = os.path.join(DATA_PATH, filename)
-    with open(full_filename, 'rb') as f:
+    with open(full_name(filename), 'rb') as f:
         return psd_tools.reader.parse(f)
 
 def decode_psd(filename):