class LazyLoadingImage:
"""
A class representing an image that can be lazily loaded from various sources.
This class supports loading an image from a filepath, URL, a PIL Image object,
or a base64 encoded string. The image is only loaded into memory when it's
accessed, not at the time of object creation. If multiple sources are provided,
an error is raised. The class also provides functionality to convert the image
to a base64 string and to access it as a PIL Image object.
Attributes:
_lazy_filepath (str): Path to the image file, if provided.
_lazy_url (str): URL of the image, if provided.
_img (Image.Image): PIL Image object, if provided.
Methods:
_load_img: Lazily loads the image from the specified source.
as_base64: Returns the image encoded as a base64 string.
as_pillow: Returns the image as a PIL Image object.
save_image_as_base64: Static method to convert a PIL Image to a base64 string.
load_image_from_base64: Static method to load an image from a base64 string.
__get_pydantic_core_schema__: Class method for Pydantic schema generation.
"""
def __init__(
self,
*,
filepath: str | None = None,
url: str | None = None,
img: "Image.Image | None" = None,
b64: str | None = None,
):
if not filepath and not url and not img and not b64:
msg = "You must specify a url or filepath or img or base64 string"
raise ValueError(msg)
if sum([bool(filepath), bool(url), bool(img), bool(b64)]) > 1:
raise ValueError("You cannot multiple input methods")
# validate file exists
if filepath and not os.path.exists(filepath):
msg = f"File does not exist: {filepath}"
raise FileNotFoundError(msg)
# validate url is valid url
if url:
from urllib3.exceptions import LocationParseError
from urllib3.util import parse_url
try:
parsed_url = parse_url(url)
except LocationParseError:
raise InvalidUrlError(f"Invalid url: {url}") # noqa
if parsed_url.scheme not in {"http", "https"} or not parsed_url.host:
msg = f"Invalid url: {url}"
raise InvalidUrlError(msg)
if b64:
img = self.load_image_from_base64(b64)
self._lazy_filepath = filepath
self._lazy_url = url
self._img = img
def __getattr__(self, key):
if key == "_img":
# http://nedbatchelder.com/blog/201010/surprising_getattr_recursion.html
raise AttributeError()
self._load_img()
return getattr(self._img, key)
def __setstate__(self, state):
self.__dict__.update(state)
def __getstate__(self):
return self.__dict__
def _load_img(self):
if self._img is None:
from PIL import Image, ImageOps
if self._lazy_filepath:
self._img = Image.open(self._lazy_filepath)
logger.debug(
f"Loaded input 🖼 of size {self._img.size} from {self._lazy_filepath}"
)
elif self._lazy_url:
import requests
self._img = Image.open(
BytesIO(
requests.get(self._lazy_url, stream=True, timeout=60).content
)
)
logger.debug(
f"Loaded input 🖼 of size {self._img.size} from {self._lazy_url}"
)
else:
raise ValueError("You must specify a url or filepath")
# fix orientation
self._img = ImageOps.exif_transpose(self._img)
@classmethod
def __get_pydantic_core_schema__(
cls, source_type: Any, handler: GetCoreSchemaHandler
) -> core_schema.CoreSchema:
def validate(value: Any) -> "LazyLoadingImage":
from PIL import Image, UnidentifiedImageError
if isinstance(value, cls):
return value
if isinstance(value, Image.Image):
return cls(img=value)
if isinstance(value, str):
if "." in value[:1000]:
try:
return cls(filepath=value)
except FileNotFoundError as e:
raise ValueError(str(e)) # noqa
try:
return cls(b64=value)
except UnidentifiedImageError:
msg = "base64 string was not recognized as a valid image type"
raise ValueError(msg) # noqa
if isinstance(value, dict):
return cls(**value)
msg = "Image value must be either a LazyLoadingImage, PIL.Image.Image or a Base64 string"
raise ValueError(msg)
def handle_b64(value: Any) -> "LazyLoadingImage":
if isinstance(value, str):
return cls(b64=value)
msg = "Image value must be either a LazyLoadingImage, PIL.Image.Image or a Base64 string"
raise ValueError(msg)
return core_schema.json_or_python_schema(
json_schema=core_schema.chain_schema(
[
core_schema.str_schema(),
core_schema.no_info_before_validator_function(
handle_b64, core_schema.any_schema()
),
]
),
python_schema=core_schema.no_info_before_validator_function(
validate, core_schema.any_schema()
),
serialization=core_schema.plain_serializer_function_ser_schema(str),
)
@staticmethod
def save_image_as_base64(image: "Image.Image") -> str:
buffered = io.BytesIO()
image.save(buffered, format="PNG")
img_bytes = buffered.getvalue()
return base64.b64encode(img_bytes).decode()
@staticmethod
def load_image_from_base64(image_str: str) -> "Image.Image":
from PIL import Image
img_bytes = base64.b64decode(image_str)
return Image.open(io.BytesIO(img_bytes))
def as_base64(self):
self._load_img()
return self.save_image_as_base64(self._img)
def as_pillow(self):
self._load_img()
return self._img
def __str__(self):
return self.as_base64()
def __repr__(self):
"""human readable representation.
shows filepath or url if available.
"""
try:
return f"<LazyLoadingImage filepath={self._lazy_filepath} url={self._lazy_url}>"
except Exception as e: # noqa
return f"<LazyLoadingImage RENDER EXCEPTION*{e}*>"