Source code for psd_tools.api.psd_image

"""
PSD Image module.
"""

from __future__ import annotations

import logging
import os
from typing import Any, BinaryIO, Callable, Literal

try:
    from typing import Self
except ImportError:
    from typing_extensions import Self

import numpy as np
from PIL.Image import Image as PILImage

from psd_tools.api import adjustments
from psd_tools.api.layers import (
    Artboard,
    FillLayer,
    Group,
    GroupMixin,
    PixelLayer,
    ShapeLayer,
    SmartObjectLayer,
    TypeLayer,
)
from psd_tools.api.pil_io import get_pil_channels, get_pil_mode
from psd_tools.constants import (
    BlendMode,
    ChannelID,
    ColorMode,
    CompatibilityMode,
    Compression,
    Resource,
    SectionDivider,
    Tag,
)
from psd_tools.psd import (
    PSD,
    ChannelImageData,
    FileHeader,
    GlobalLayerMaskInfo,
    ImageData,
    ImageResources,
    LayerInfo,
    LayerRecords,
    TaggedBlocks,
)

logger = logging.getLogger(__name__)


[docs] class PSDImage(GroupMixin): """ Photoshop PSD/PSB file object. The low-level data structure is accessible at :py:attr:`PSDImage._record`. Example:: from psd_tools import PSDImage psd = PSDImage.open('example.psd') image = psd.compose() for layer in psd: layer_image = layer.compose() """ def __init__(self, data: PSD): from psd_tools.api.layers import Layer assert isinstance(data, PSD) self._record = data self._layers: list[Layer] = [] self._compatibility_mode = CompatibilityMode.DEFAULT self._updated_layers = False # Flag to check if the layer tree is edited. self._init()
[docs] @classmethod def new( cls, mode: str, size: tuple[int, int], color: int = 0, depth: Literal[8, 16, 32] = 8, **kwargs: Any, ): """ Create a new PSD document. :param mode: The color mode to use for the new image. :param size: A tuple containing (width, height) in pixels. :param color: What color to use for the image. Default is black. :return: A :py:class:`~psd_tools.api.psd_image.PSDImage` object. """ header = cls._make_header(mode, size, depth) image_data = ImageData.new(header, color=color, **kwargs) # TODO: Add default metadata. return cls( PSD( header=header, image_data=image_data, image_resources=ImageResources.new(), ) )
[docs] @classmethod def frompil(cls, image: PILImage, compression=Compression.RLE) -> Self: """ Create a new PSD document from PIL Image. :param image: PIL Image object. :param compression: ImageData compression option. See :py:class:`~psd_tools.constants.Compression`. :return: A :py:class:`~psd_tools.api.psd_image.PSDImage` object. """ header = cls._make_header(image.mode, image.size) # TODO: Add default metadata. # TODO: Perhaps make this smart object. image_data = ImageData(compression=compression) image_data.set_data([channel.tobytes() for channel in image.split()], header) return cls( PSD( header=header, image_data=image_data, image_resources=ImageResources.new(), ) )
[docs] @classmethod def open(cls, fp: BinaryIO | str | bytes | os.PathLike, **kwargs: Any) -> Self: """ Open a PSD document. :param fp: filename or file-like object. :param encoding: charset encoding of the pascal string within the file, default 'macroman'. Some psd files need explicit encoding option. :return: A :py:class:`~psd_tools.api.psd_image.PSDImage` object. """ if isinstance(fp, (str, bytes, os.PathLike)): with open(fp, "rb") as f: self = cls(PSD.read(f, **kwargs)) else: self = cls(PSD.read(fp, **kwargs)) return self
[docs] def save(self, fp: BinaryIO | str | bytes | os.PathLike, mode: str = "wb", **kwargs: Any) -> None: """ Save the PSD file. Updates the ImageData section if the layer structure has been updated. :param fp: filename or file-like object. :param encoding: charset encoding of the pascal string within the file, default 'macroman'. :param mode: file open mode, default 'wb'. """ self._update_record() if self._updated_layers: composited_psd = self.composite(force=True) self._record.image_data.set_data([channel.tobytes() for channel in composited_psd.split()], self._record.header) if isinstance(fp, (str, bytes, os.PathLike)): with open(fp, mode) as f: self._record.write(f, **kwargs) else: self._record.write(fp, **kwargs)
[docs] def topil( self, channel: int | ChannelID | None = None, apply_icc: bool = True ) -> PILImage | None: """ Get PIL Image. :param channel: Which channel to return; e.g., 0 for 'R' channel in RGB image. See :py:class:`~psd_tools.constants.ChannelID`. When `None`, the method returns all the channels supported by PIL modes. :param apply_icc: Whether to apply ICC profile conversion to sRGB. :return: :py:class:`PIL.Image`, or `None` if the composed image is not available. """ from .pil_io import convert_image_data_to_pil self._update_record() if self.has_preview(): return convert_image_data_to_pil(self, channel, apply_icc) return None
[docs] def numpy( self, channel: Literal["color", "shape", "alpha", "mask"] | None = None ) -> np.ndarray: """ Get NumPy array of the layer. :param channel: Which channel to return, can be 'color', 'shape', 'alpha', or 'mask'. Default is 'color+alpha'. :return: :py:class:`numpy.ndarray` """ from .numpy_io import get_array return get_array(self, channel)
[docs] def composite( self, viewport: tuple[int, int, int, int] | None = None, force: bool = False, color: float | tuple[float, ...] = 1.0, alpha: float = 0.0, layer_filter: Callable | None = None, ignore_preview: bool = False, apply_icc: bool = True, ): """ Composite the PSD image. :param viewport: Viewport bounding box specified by (x1, y1, x2, y2) tuple. Default is the viewbox of the PSD. :param ignore_preview: Boolean flag to whether skip compositing when a pre-composited preview is available. :param force: Boolean flag to force vector drawing. :param color: Backdrop color specified by scalar or tuple of scalar. The color value should be in [0.0, 1.0]. For example, (1., 0., 0.) specifies red in RGB color mode. :param alpha: Backdrop alpha in [0.0, 1.0]. :param layer_filter: Callable that takes a layer as argument and returns whether if the layer is composited. Default is :py:func:`~psd_tools.api.layers.PixelLayer.is_visible`. :return: :py:class:`PIL.Image`. """ from psd_tools.composite import composite_pil self._update_record() if ( not (ignore_preview or force or layer_filter) and self.has_preview() and not self._updated_layers ): return self.topil(apply_icc=apply_icc) return composite_pil( self, color, alpha, viewport, layer_filter, force, apply_icc=apply_icc )
[docs] def is_visible(self) -> bool: """ Returns visibility of the element. :return: `bool` """ return self.visible
@property def parent(self) -> None: """Parent of this layer.""" return None
[docs] def is_group(self) -> bool: """ Return True if the layer is a group. :return: `bool` """ return isinstance(self, GroupMixin)
[docs] def has_preview(self) -> bool: """ Returns if the document has real merged data. When True, `topil()` returns pre-composed data. """ version_info = self.image_resources.get_data(Resource.VERSION_INFO) if version_info: return version_info.has_composite return True # Assuming the image data is valid by default.
@property def name(self) -> str: """ Element name. :return: `'Root'` """ return "Root" @property def kind(self) -> str: """ Kind. :return: `'psdimage'` """ return self.__class__.__name__.lower() @property def visible(self) -> bool: """ Visibility. :return: `True` """ return True @property def left(self) -> int: """ Left coordinate. :return: `0` """ return 0 @property def top(self) -> int: """ Top coordinate. :return: `0` """ return 0 @property def right(self) -> int: """ Right coordinate. :return: `int` """ return self.width @property def bottom(self) -> int: """ Bottom coordinate. :return: `int` """ return self.height @property def width(self) -> int: """ Document width. :return: `int` """ return self._record.header.width @property def height(self) -> int: """ Document height. :return: `int` """ return self._record.header.height @property def size(self) -> tuple[int, int]: """ (width, height) tuple. :return: `tuple` """ return self.width, self.height @property def offset(self) -> tuple[int, int]: """ (left, top) tuple. :return: `tuple` """ return self.left, self.top @property def bbox(self) -> tuple[int, int, int, int]: """ Minimal bounding box that contains all the visible layers. Use :py:attr:`~psd_tools.api.psd_image.PSDImage.viewbox` to get viewport bounding box. When the psd is empty, bbox is equal to the canvas bounding box. :return: (left, top, right, bottom) `tuple`. """ bbox = super(PSDImage, self).bbox if bbox == (0, 0, 0, 0): bbox = self.viewbox return bbox @property def viewbox(self) -> tuple[int, int, int, int]: """ Return bounding box of the viewport. :return: (left, top, right, bottom) `tuple`. """ return self.left, self.top, self.right, self.bottom @property def color_mode(self) -> ColorMode: """ Document color mode, such as 'RGB' or 'GRAYSCALE'. See :py:class:`~psd_tools.constants.ColorMode`. :return: :py:class:`~psd_tools.constants.ColorMode` """ return self._record.header.color_mode @property def channels(self) -> int: """ Number of color channels. :return: `int` """ return self._record.header.channels @property def depth(self) -> int: """ Pixel depth bits. :return: `int` """ return self._record.header.depth @property def version(self) -> int: """ Document version. PSD file is 1, and PSB file is 2. :return: `int` """ return self._record.header.version @property def image_resources(self) -> ImageResources: """ Document image resources. :py:class:`~psd_tools.psd.image_resources.ImageResources` is a dict-like structure that keeps various document settings. See :py:class:`psd_tools.constants.Resource` for available keys. :return: :py:class:`~psd_tools.psd.image_resources.ImageResources` Example:: from psd_tools.constants import Resource version_info = psd.image_resources.get_data(Resource.VERSION_INFO) slices = psd.image_resources.get_data(Resource.SLICES) Image resources contain an ICC profile. The following shows how to export a PNG file with embedded ICC profile:: from psd_tools.constants import Resource icc_profile = psd.image_resources.get_data(Resource.ICC_PROFILE) image = psd.compose(apply_icc=False) image.save('output.png', icc_profile=icc_profile) """ return self._record.image_resources @property def tagged_blocks(self) -> TaggedBlocks | None: """ Document tagged blocks that is a dict-like container of settings. See :py:class:`psd_tools.constants.Tag` for available keys. :return: :py:class:`~psd_tools.psd.tagged_blocks.TaggedBlocks` or `None`. Example:: from psd_tools.constants import Tag patterns = psd.tagged_blocks.get_data(Tag.PATTERNS1) """ return self._record.layer_and_mask_information.tagged_blocks @property def compatibility_mode(self) -> CompatibilityMode: """ Set the compositing and layer organization compatibility mode. Writable. :return: :py:class:`~psd_tools.constants.CompatibilityMode` """ return self._compatibility_mode @compatibility_mode.setter def compatibility_mode(self, value: CompatibilityMode) -> None: self._compatibility_mode = value self._compute_clipping_layers() @property def pil_mode(self) -> str: alpha = self.channels - get_pil_channels(get_pil_mode(self.color_mode)) return get_pil_mode(self.color_mode, alpha > 0)
[docs] def has_thumbnail(self) -> bool: """True if the PSDImage has a thumbnail resource.""" return ( Resource.THUMBNAIL_RESOURCE in self.image_resources or Resource.THUMBNAIL_RESOURCE_PS4 in self.image_resources )
[docs] def thumbnail(self) -> PILImage | None: """ Returns a thumbnail image in PIL.Image. When the file does not contain an embedded thumbnail image, returns None. """ from .pil_io import convert_thumbnail_to_pil if Resource.THUMBNAIL_RESOURCE in self.image_resources: return convert_thumbnail_to_pil( self.image_resources.get_data(Resource.THUMBNAIL_RESOURCE) ) elif Resource.THUMBNAIL_RESOURCE_PS4 in self.image_resources: return convert_thumbnail_to_pil( self.image_resources.get_data(Resource.THUMBNAIL_RESOURCE_PS4) ) return None
def __repr__(self) -> str: return ("%s(mode=%s size=%dx%d depth=%d channels=%d)") % ( self.__class__.__name__, self.color_mode, self.width, self.height, self._record.header.depth, self._record.header.channels, ) def _repr_pretty_(self, p, cycle): if cycle: return self.__repr__() def _pretty(layer, p): p.text(layer.__repr__()) if hasattr(layer, "clip_layers"): for clip_layer in layer.clip_layers or []: p.break_() p.text(" + ") p.pretty(clip_layer) if hasattr(layer, "__iter__"): with p.indent(2): for idx, child in enumerate(layer): p.break_() p.text("[%d] " % idx) _pretty(child, p) _pretty(self, p) @classmethod def _make_header( cls, mode: str, size: tuple[int, int], depth: Literal[8, 16, 32] = 8 ) -> FileHeader: from .pil_io import get_color_mode assert depth in (8, 16, 32), "Invalid depth: %d" % (depth) assert size[0] <= 300000, "Width too large > 300,000" assert size[1] <= 300000, "Height too large > 300,000" version = 1 if size[0] > 30000 or size[1] > 30000: logger.debug("Width or height larger than 30,000 pixels") version = 2 color_mode = get_color_mode(mode) alpha = int(mode.upper().endswith("A")) channels = ColorMode.channels(color_mode, alpha) return FileHeader( version=version, width=size[0], height=size[1], depth=depth, channels=channels, color_mode=color_mode, ) def _get_pattern(self, pattern_id): """Get pattern item by id.""" for key in (Tag.PATTERNS1, Tag.PATTERNS2, Tag.PATTERNS3): if key in self.tagged_blocks: data = self.tagged_blocks.get_data(key) for pattern in data: if pattern.pattern_id == pattern_id: return pattern return None def _clear_clipping_layers(self) -> None: for layer in self.descendants(): layer._clip_layers = [] layer._has_clip_target = True def _compute_clipping_layers(self) -> None: self._update_record() self._clear_clipping_layers() def rec_helper(layer): if not layer.is_group(): return stack = [] for sublayer in reversed(layer._layers): if sublayer.clipping_layer: stack.append(sublayer) else: if sublayer.blend_mode == BlendMode.PASS_THROUGH and ( self.compatibility_mode == CompatibilityMode.PAINT_TOOL_SAI or self.compatibility_mode == CompatibilityMode.CLIP_STUDIO_PAINT ): for clip_layer in stack: clip_layer._has_clip_target = False else: stack.reverse() sublayer._clip_layers = stack stack = [] rec_helper(sublayer) for clip_layer in stack: clip_layer._has_clip_target = False rec_helper(self) def _init(self) -> None: """Initialize layer structure.""" from psd_tools.api.layers import Layer group_stack: list[Group | Artboard | PSDImage] = [self] for record, channels in self._record._iter_layers(): current_group = group_stack[-1] blocks = record.tagged_blocks end_of_group = False layer: Layer | Group | Artboard | PSDImage | None = None divider = blocks.get_data(Tag.SECTION_DIVIDER_SETTING, None) divider = blocks.get_data(Tag.NESTED_SECTION_DIVIDER_SETTING, divider) if ( divider is not None and # Some PSDs contain dividers with SectionDivider.OTHER. # Ignoring them, allows the correct categorization of the layer. # Issue : https://github.com/psd-tools/psd-tools/issues/338 divider.kind is not SectionDivider.OTHER ): if divider.kind == SectionDivider.BOUNDING_SECTION_DIVIDER: layer = Group( # type: ignore psd=self, record=None, # type: ignore channels=None, # type: ignore parent=current_group, ) layer._set_bounding_records(record, channels) group_stack.append(layer) elif divider.kind in ( SectionDivider.OPEN_FOLDER, SectionDivider.CLOSED_FOLDER, ): layer = group_stack.pop() assert not isinstance(layer, PSDImage) layer._record = record layer._channels = channels for key in ( Tag.ARTBOARD_DATA1, Tag.ARTBOARD_DATA2, Tag.ARTBOARD_DATA3, ): if key in blocks: layer = Artboard._move(layer) end_of_group = True else: logger.warning("Divider %s found." % divider.kind) elif Tag.TYPE_TOOL_OBJECT_SETTING in blocks or Tag.TYPE_TOOL_INFO in blocks: layer = TypeLayer(self, record, channels, current_group) elif ( Tag.SMART_OBJECT_LAYER_DATA1 in blocks or Tag.SMART_OBJECT_LAYER_DATA2 in blocks or Tag.PLACED_LAYER1 in blocks or Tag.PLACED_LAYER2 in blocks ): layer = SmartObjectLayer(self, record, channels, current_group) else: for key in adjustments.TYPES.keys(): if key in blocks: layer = adjustments.TYPES[key]( self, record, channels, current_group ) break # If nothing applies, this is either a shape or pixel layer. shape_condition = record.flags.pixel_data_irrelevant and ( Tag.VECTOR_ORIGINATION_DATA in blocks or Tag.VECTOR_MASK_SETTING1 in blocks or Tag.VECTOR_MASK_SETTING2 in blocks or Tag.VECTOR_STROKE_DATA in blocks or Tag.VECTOR_STROKE_CONTENT_DATA in blocks ) if isinstance(layer, (type(None), FillLayer)) and shape_condition: layer = ShapeLayer(self, record, channels, current_group) if layer is None: layer = PixelLayer(self, record, channels, current_group) assert layer is not None if not end_of_group: assert not isinstance(layer, PSDImage) current_group._layers.append(layer) self._compute_clipping_layers() def _update_record(self) -> None: """ Compiles the tree layer structure back into records and channels list recursively """ if not self._updated_layers: # Skip if nothing is changed. return layer_records, channel_image_data = _build_record_tree(self) # PSDImage.frompil doesn't create a LayerInfo attribute to LayerAndMaskInformation if not self._record.layer_and_mask_information.layer_info: self._record.layer_and_mask_information.layer_info = LayerInfo() if not self._record.layer_and_mask_information.global_layer_mask_info: self._record.layer_and_mask_information.global_layer_mask_info = ( GlobalLayerMaskInfo() ) if not self._record.layer_and_mask_information.tagged_blocks: self._record.layer_and_mask_information.tagged_blocks = TaggedBlocks() layer_info = self._record.layer_and_mask_information.layer_info layer_info.layer_records = layer_records layer_info.channel_image_data = channel_image_data layer_info.layer_count = len(layer_records)
def _build_record_tree( layer_group: GroupMixin, ) -> tuple[LayerRecords, ChannelImageData]: """ Builds the layer tree structure from records and channels list recursively """ layer_records = LayerRecords() channel_image_data = ChannelImageData() for layer in layer_group: if isinstance(layer, (Group, Artboard)): layer_records.append(layer._bounding_record) channel_image_data.append(layer._bounding_channels) tmp_layer_records, tmp_channel_image_data = _build_record_tree(layer) layer_records.extend(tmp_layer_records) channel_image_data.extend(tmp_channel_image_data) layer_records.append(layer._record) channel_image_data.append(layer._channels) return (layer_records, channel_image_data)