Commits

Mikhail Korobov committed b378b79 Merge

Merge pull request #14 from kmike/cmyk

CMYK support (thanks @megabuz for original patch)

  • Participants
  • Parent commits 6bdff46, 828b0f4

Comments (0)

Files changed (2)

File src/psd_tools/constants.py

     LAB = 9
 
 class ChannelID(Enum):
-    RED = 0
-    GREEN = 1
-    BLUE = 2
     TRANSPARENCY_MASK = -1
     USER_LAYER_MASK = -2
     REAL_USER_LAYER_MASK = -3

File src/psd_tools/user_api/pil_support.py

     layers = decoded_data.layer_and_mask_data.layers
     layer = layers.layer_records[layer_index]
 
-    channels_data = layers.channel_image_data[layer_index]
-    size = layer.width(), layer.height()
-    channel_types = [info.id for info in layer.channels]
+    return _channel_data_to_PIL(
+        channel_data = layers.channel_image_data[layer_index],
+        channel_ids = _get_layer_channel_ids(layer),
+        color_mode = decoded_data.header.color_mode,  # XXX?
+        size = (layer.width(), layer.height()),
+        depth = decoded_data.header.depth,
+        icc_profile = get_icc_profile(decoded_data)
+    )
 
-    return _channels_data_to_PIL(
-        channels_data, channel_types, size,
-        decoded_data.header.depth, get_icc_profile(decoded_data))
 
 def extract_composite_image(decoded_data):
     """
     """
     header = decoded_data.header
     size = header.width, header.height
-    if header.color_mode == ColorMode.RGB:
-
-        if header.number_of_channels == 3:
-            channel_types = [0, 1, 2]
-        elif header.number_of_channels == 4:
-            channel_types = [0, 1, 2, -1]
-        else:
-            warnings.warn("This number of channels (%d) is unsupported for this color mode (%s)" % (
-                         header.number_of_channels, header.color_mode))
-            return
-
-    else:
-        warnings.warn("Unsupported color mode (%s)" % header.color_mode)
+    if size == (0, 0):
         return
 
-    return _channels_data_to_PIL(
-        decoded_data.image_data,
-        channel_types,
-        size,
-        header.depth,
-        get_icc_profile(decoded_data),
+    channel_ids = _get_header_channel_ids(header)
+    if channel_ids is None:
+        warnings.warn("This number of channels (%d) is unsupported for this color mode (%s)" % (
+                     header.number_of_channels, header.color_mode))
+        return
+
+    return _channel_data_to_PIL(
+        channel_data=decoded_data.image_data,
+        channel_ids=channel_ids,
+        color_mode=header.color_mode,
+        size=size,
+        depth=header.depth,
+        icc_profile=get_icc_profile(decoded_data)
     )
 
 
-def _channels_data_to_PIL(channels_data, channel_types, size, depth, icc_profile):
+def get_icc_profile(decoded_data):
+    """
+    Return ICC image profile if it exists and was correctly decoded
+    """
+    # fixme: move this function somewhere?
+    icc_profiles = [res.data for res in decoded_data.image_resource_blocks
+                   if res.resource_id == ImageResourceID.ICC_PROFILE]
+
+    if not icc_profiles:
+        return None
+
+    icc_profile = icc_profiles[0]
+
+    if isinstance(icc_profile, bytes): # profile was not decoded
+        return None
+
+    return icc_profile
+
+
+def apply_opacity(im, opacity):
+    """ Apply opacity to an image. """
+    if im.mode == 'RGB':
+        im.putalpha(opacity)
+        return im
+    elif im.mode == 'RGBA':
+        r, g, b, a = im.split()
+        opacity_scale = opacity / 255
+        a = a.point(lambda i: i*opacity_scale)
+        return Image.merge('RGBA', [r, g, b, a])
+    else:
+        raise NotImplementedError()
+
+
+def _channel_data_to_PIL(channel_data, channel_ids, color_mode, size, depth, icc_profile):
+    bands = _get_band_images(
+        channel_data=channel_data,
+        channel_ids=channel_ids,
+        color_mode=color_mode,
+        size=size,
+        depth=depth
+    )
+    return _merge_bands(bands, color_mode, size, icc_profile)
+
+
+def _merge_bands(bands, color_mode, size, icc_profile):
     if Image is None:
         raise Exception("This module requires PIL (or Pillow) installed.")
 
-    if size == (0, 0):
-        return
+    if color_mode == ColorMode.RGB:
+        merged_image = Image.merge('RGB', [bands[key] for key in 'RGB'])
+    elif color_mode == ColorMode.CMYK:
+        merged_image = Image.merge('CMYK', [bands[key] for key in 'CMYK'])
+        # colors are inverted in Photoshop CMYK images; invert them back
+        merged_image = frombytes('CMYK', size, merged_image.tobytes(), 'raw', 'CMYK;I')
+    else:
+        raise NotImplementedError()
 
+    if icc_profile is not None:
+        try:
+            sRGB_profile = ImageCms.createProfile('sRGB')
+            merged_image = ImageCms.profileToProfile(merged_image, icc_profile, sRGB_profile, outputMode='RGB')
+        except ImageCms.PyCMSError as e:
+            # PIL/Pillow/(old littlecms?) can't convert some ICC profiles
+            warnings.warn(repr(e))
+
+    if color_mode == ColorMode.CMYK:
+        merged_image = merged_image.convert('RGB')
+
+    alpha = bands.get('A')
+    if alpha:
+        merged_image.putalpha(alpha)
+
+    return merged_image
+
+
+def _get_band_images(channel_data, channel_ids, color_mode, size, depth):
     bands = {}
+    for channel, channel_id in zip(channel_data, channel_ids):
 
-    for channel, channel_type in zip(channels_data, channel_types):
-
-        pil_band = channel_id_to_PIL(channel_type)
+        pil_band = _channel_id_to_PIL(channel_id, color_mode)
         if pil_band is None:
-            warnings.warn("Unsupported channel type (%d)" % channel_type)
+            warnings.warn("Unsupported channel type (%d)" % channel_id)
             continue
 
         if channel.compression in [Compression.RAW, Compression.ZIP, Compression.ZIP_WITH_PREDICTION]:
             continue
 
         bands[pil_band] = im.convert('L')
+    return bands
 
-    mode = _get_mode(bands.keys())
-    merged_image = Image.merge(mode, [bands[band] for band in mode])
-
-    if icc_profile is not None:
-        display_profile = ImageCms.createProfile('sRGB') # XXX: ImageCms.get_display_profile()?
-        ImageCms.profileToProfile(merged_image, icc_profile, display_profile, inPlace=True)
-
-    return merged_image
-
-
-def channel_id_to_PIL(channel_id):
-    BANDS_MAP = {
-        ChannelID.RED: 'R',
-        ChannelID.GREEN: 'G',
-        ChannelID.BLUE: 'B',
-        ChannelID.TRANSPARENCY_MASK: 'A'
-    }
-    return BANDS_MAP.get(channel_id, None)
-
-
-def _get_mode(band_keys):
-    for mode in ['RGBA', 'RGB']:
-        if set(band_keys) == set(list(mode)):
-            return mode
 
 def _from_8bit_raw(data, size):
     return frombytes('L', size, data, "raw", 'L')
 
+
 def _from_16bit_raw(data, size):
     im = frombytes('I', size, data, "raw", 'I;16B')
     return im.point(lambda i: i * (1/(256.0)))
 
+
 def _from_32bit_raw(data, size):
     pixels = be_array_from_bytes("f", data)
     im = Image.new("F", size)
     im.putdata(pixels, 255, 0)
     return im
 
-def get_icc_profile(decoded_data):
-    """
-    Returns ICC image profile (if it exists and was correctly decoded)
-    """
-    # fixme: move this function somewhere?
-    icc_profiles = [res.data for res in decoded_data.image_resource_blocks
-                   if res.resource_id == ImageResourceID.ICC_PROFILE]
 
-    if not icc_profiles:
+def _channel_id_to_PIL(channel_id, color_mode):
+    if ChannelID.is_known(channel_id):
+        if channel_id == ChannelID.TRANSPARENCY_MASK:
+            return 'A'
+        warnings.warn("Channel %s (%s) is not handled" % (channel_id, ChannelID.name_of(channel_id)))
         return None
 
-    icc_profile = icc_profiles[0]
+    try:
+        assert channel_id >= 0
+        if color_mode == ColorMode.RGB:
+            return 'RGB'[channel_id]
+        elif color_mode == ColorMode.CMYK:
+            return 'CMYK'[channel_id]
 
-    if isinstance(icc_profile, bytes): # profile was not decoded
+    except IndexError:
+        # spot channel
+        warnings.warn("Spot channel %s is not handled" % channel_id)
         return None
 
-    return icc_profile
 
-def apply_opacity(im, opacity):
-    """
-    Applies opacity to an image.
-    """
-    if im.mode == 'RGB':
-        im.putalpha(opacity)
-        return im
-    elif im.mode == 'RGBA':
-        r, g, b, a = im.split()
-        opacity_scale = opacity / 255
-        a = a.point(lambda i: i*opacity_scale)
-        return Image.merge('RGBA', [r, g, b, a])
+def _get_header_channel_ids(header):
+    channel_ids = None
+    if header.color_mode == ColorMode.RGB:
+        if header.number_of_channels == 3:
+            channel_ids = [0, 1, 2]
+        elif header.number_of_channels == 4:
+            channel_ids = [0, 1, 2, -1]
+
+    elif header.color_mode == ColorMode.CMYK:
+        if header.number_of_channels == 4:
+            channel_ids = [0, 1, 2, 3]
+        elif header.number_of_channels == 5:
+            # XXX: how to distinguish
+            # "4 CMYK + 1 alpha" and "4 CMYK + 1 spot"?
+            channel_ids = [0, 1, 2, 3, -1]
+
     else:
-        raise NotImplementedError()
+        warnings.warn("Unsupported color mode (%s)" % header.color_mode)
+    return channel_ids
+
+
+def _get_layer_channel_ids(layer):
+    return [info.id for info in layer.channels]