Source code for SDF.data_model.image_data

import base64
from io import BytesIO
from xml.etree.ElementTree import Element

from PIL import Image

from SDF.data_model._helper_functions import pop_element_attribute, pop_element_text, element_is_empty
from SDF.data_model.abstract import Data


class ImageData(Data):
    """Wraps a PIL Image. Represents an SDF <data> element of type 'img'."""

    def __init__(self, data: Image.Image):
        if not isinstance(data, Image.Image):
            raise TypeError(f"Expected a PIL Image, got {type(data)}")
        self.__image = data

    @property
    def data(self) -> Image.Image:
        return self.__image

    @property
    def type_for_xml(self) -> str:
        return "img"

    def to_xml_element(self) -> Element:
        element = Element("data", attrib={"encoding": "base64", "type": "image/png"})

        element.text = ImageData.image_to_base64(self.data)
        return element

    @classmethod
    def from_xml_element(cls, element: Element) -> "ImageData":
        if element.tag != "data":
            raise ValueError(f"Expected a <data> element, got {element.tag}")

        # pop valid, but irrelevant attributes
        for key in ("dtype", "encoding", "type"):
            if key in element.attrib:
                pop_element_attribute(element, key)

        base64_str = pop_element_text(element).strip()

        if not element_is_empty(element):
            raise ValueError("Element is not empty")

        return cls(ImageData.base64_to_image(base64_str))

    def __repr__(self):
        return f"{self.__class__.__name__}({self.data!r})"

    @staticmethod
    def image_to_base64(image: Image.Image) -> str:
        """Return base64-encoded PNG representation of `image`"""
        png_buffer = BytesIO()
        image.save(png_buffer, format="png")
        return base64.encodebytes(png_buffer.getvalue()).decode("ascii")

    @staticmethod
    def base64_to_image(base64_str: str) -> Image.Image:
        """Read base64-encoded PNG image"""
        img_buffer = BytesIO()
        img_buffer.write(base64.decodebytes(base64_str.strip().encode("ascii")))
        return Image.open(img_buffer, formats=["PNG"])

    def __eq__(self, other):
        if isinstance(other, ImageData) and self.data.size == other.data.size:
            # compare resolutions (fast), then compare pixel values (slow)
            return (self.data.size == other.data.size and
                    self.image_to_base64(self.__image) == self.image_to_base64(other.__image))
        return False