# Copyright 2025 - Oumi## Licensed under the Apache License, Version 2.0 (the "License");# you may not use this file except in compliance with the License.# You may obtain a copy of the License at## http://www.apache.org/licenses/LICENSE-2.0## Unless required by applicable law or agreed to in writing, software# distributed under the License is distributed on an "AS IS" BASIS,# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.# See the License for the specific language governing permissions and# limitations under the License.importiofromimportlib.utilimportfind_specfrompathlibimportPathfromtypingimportFinal,Optional,UnionimportPIL.Imageimportrequestsfromoumi.utils.loggingimportlogger# For details on image modes, see# https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modesDEFAULT_IMAGE_MODE:Final[str]="RGB"_FILE_URL_PREFIX:Final[str]="file://"_DEFAULT_PDF_DPI:Final[int]=200
[docs]defcreate_png_bytes_from_image(pil_image:PIL.Image.Image)->bytes:"""Encodes PIL image into PNG format, and returns PNG image bytes. Args: pil_image: An input image. Returns: bytes: PNG bytes representation of the image. """try:output=io.BytesIO()pil_image.save(output,format="PNG")returnoutput.getvalue()exceptException:logger.error("Failed to convert an image to PNG bytes.")raise
[docs]defcreate_png_bytes_from_image_list(pil_images:list[PIL.Image.Image])->list[bytes]:"""Encodes PIL images into PNG format, and returns PNG image bytes. Args: pil_images: A list of input images. Returns: A list of PNG-encoded images. """return[create_png_bytes_from_image(image)forimageinpil_images]
[docs]defconvert_pil_image_mode(image:PIL.Image.Image,*,mode:Optional[str])->PIL.Image.Image:"""Converts a PIL image to the requested mode (if it's not in that mode already) . Args: image: An input image. mode: The requested image mode e.g., "RGB", "HSV", "RGBA", "P" (8-bit pixels, using a color palette). For details, see https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes Returns: An image in the requested mode . If an input image was already in the correct mode then return it for efficiency. Otherwise, a different image object is returned. """ifnotmodeorimage.mode==mode:# Return the original object for better performance.returnimageold_mode=image.modetry:returnimage.convert(mode)exceptExceptionase:raiseRuntimeError(f"Failed to convert an image from {old_mode} to {mode} mode!")frome
[docs]defload_pil_image_from_path(input_image_filepath:Union[str,Path],mode:str=DEFAULT_IMAGE_MODE)->PIL.Image.Image:"""Loads an image from a path. Args: input_image_filepath: A file path of an image. The image can be in any format supported by PIL. mode: The requested image mode e.g., "RGB", "HSV", "RGBA", "P" (8-bit pixels, using a color palette). For details, see https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes Returns: bytes: PNG bytes representation of the image. """ifnotinput_image_filepath:raiseValueError("Empty image file path.")ifisinstance(input_image_filepath,str)andinput_image_filepath.lower().startswith(_FILE_URL_PREFIX):input_image_filepath=input_image_filepath[len(_FILE_URL_PREFIX):]input_image_filepath=Path(input_image_filepath)ifnotinput_image_filepath.is_file():raiseValueError(f"Image path is not a file: {input_image_filepath}"ifinput_image_filepath.exists()elsef"Image path doesn't exist: {input_image_filepath}")try:pil_image=convert_pil_image_mode(PIL.Image.open(input_image_filepath),mode=mode)exceptException:logger.error(f"Failed to load an image from path: {input_image_filepath}")raisereturnpil_image
[docs]defload_pil_image_from_url(input_image_url:str,mode:str=DEFAULT_IMAGE_MODE)->PIL.Image.Image:"""Loads a PIL image from a URL. Args: input_image_url: An image URL. mode: The requested image mode e.g., "RGB", "HSV", "RGBA", "P" (8-bit pixels, using a color palette). For details, see https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes Returns: bytes: PNG bytes representation of the image. """ifnotinput_image_url:raiseValueError("Empty image URL!")try:response=requests.get(input_image_url,stream=True)response.raise_for_status()exceptrequests.exceptions.RequestException:logger.exception(f"Failed to download image: '{input_image_url}'")raisereturnload_pil_image_from_bytes(response.content,mode=mode)
[docs]defload_pil_image_from_bytes(image_bytes:Optional[bytes],mode:str=DEFAULT_IMAGE_MODE)->PIL.Image.Image:"""Loads an image from raw image bytes. Args: image_bytes: A input image bytes. Can be in any image format supported by PIL. mode: The requested image mode e.g., "RGB", "HSV", "RGBA", "P" (8-bit pixels, using a color palette). For details, see https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes Returns: PIL.Image.Image: PIL representation of the image. """ifimage_bytesisNoneorlen(image_bytes)==0:raiseValueError("No image bytes.")try:pil_image=convert_pil_image_mode(PIL.Image.open(io.BytesIO(image_bytes)),mode=mode)exceptException:logger.error(f"Failed to load an image from raw image bytes ({len(image_bytes)} bytes).")raisereturnpil_image
def_check_pdf2image_dependency():ifnotfind_spec("pdf2image"):raiseRuntimeError("Failed to find the required dependency package: 'pdf2image'. ""Run `pip install oumi[file_formats]`, and try again.")
[docs]defload_pdf_pages_from_path(input_pdf_filepath:Union[str,Path],*,dpi:int=_DEFAULT_PDF_DPI,mode:str=DEFAULT_IMAGE_MODE,)->list[PIL.Image.Image]:"""Loads PDF pages as PIL images from a path. Args: input_pdf_filepath: A file path of an PDF document. dpi: Resolution to use for PDF page images (dots per inch). mode: The requested image mode e.g., "RGB", "HSV", "RGBA", "P" (8-bit pixels, using a color palette). For details, see https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes Returns: PDF pages as PIL images (PIL.Image.Image). """ifnotinput_pdf_filepath:raiseValueError("Empty PDF file path.")ifisinstance(input_pdf_filepath,str)andinput_pdf_filepath.lower().startswith(_FILE_URL_PREFIX):input_pdf_filepath=input_pdf_filepath[len(_FILE_URL_PREFIX):]input_filepath=Path(input_pdf_filepath)ifnotinput_filepath.is_file():raiseValueError(f"PDF path is not a file: {input_filepath}"ifinput_filepath.exists()elsef"PDF path doesn't exist: {input_filepath}")_check_pdf2image_dependency()importpdf2image# pyright: ignore[reportMissingImports]page_images=pdf2image.convert_from_path(input_filepath,dpi=dpi)num_pages=len(page_images)forpage_idxinrange(num_pages):try:page_images[page_idx]=convert_pil_image_mode(page_images[page_idx],mode=mode)exceptException:logger.error("Failed to convert image mode for PDF page "f"{page_idx+1} of {num_pages}: {input_filepath}")raisereturnpage_images
[docs]defload_pdf_pages_from_bytes(pdf_bytes:Optional[bytes],*,dpi:int=_DEFAULT_PDF_DPI,mode:str=DEFAULT_IMAGE_MODE,)->list[PIL.Image.Image]:"""Loads PDF pages as PIL images from raw PDF file bytes. Args: pdf_bytes: PDF file content. dpi: Resolution to use for PDF page images (dots per inch). mode: The requested image mode e.g., "RGB", "HSV", "RGBA", "P" (8-bit pixels, using a color palette). For details, see https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes Returns: PDF pages as PIL images (PIL.Image.Image). """ifpdf_bytesisNoneorlen(pdf_bytes)==0:raiseValueError("No PDF bytes.")_check_pdf2image_dependency()importpdf2image# pyright: ignore[reportMissingImports]page_images=pdf2image.convert_from_bytes(pdf_bytes,dpi=dpi)num_pages=len(page_images)forpage_idxinrange(num_pages):try:page_images[page_idx]=convert_pil_image_mode(page_images[page_idx],mode=mode)exceptException:logger.error("Failed to convert image mode for PDF page "f"{page_idx+1} or {num_pages}")raisereturnpage_images
[docs]defload_pdf_pages_from_url(pdf_url:str,*,dpi:int=_DEFAULT_PDF_DPI,mode:str=DEFAULT_IMAGE_MODE)->list[PIL.Image.Image]:"""Loads PDF pages as PIL images from from PDF URL. Args: pdf_url: A PDF URL. dpi: Resolution to use for PDF page images (dots per inch). mode: The requested image mode e.g., "RGB", "HSV", "RGBA", "P" (8-bit pixels, using a color palette). For details, see https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes Returns: PDF pages as PIL images (PIL.Image.Image). """ifnotpdf_url:raiseValueError("Empty PDF URL!")try:response=requests.get(pdf_url,stream=True)response.raise_for_status()exceptrequests.exceptions.RequestException:logger.exception(f"Failed to download PDF: '{pdf_url}'")raisereturnload_pdf_pages_from_bytes(response.content,dpi=dpi,mode=mode)
[docs]defcreate_png_bytes_from_image_bytes(image_bytes:Optional[bytes],mode:str=DEFAULT_IMAGE_MODE)->bytes:"""Loads an image from raw image bytes, and converts to PNG image bytes. Args: image_bytes: A input image bytes. Can be in any image format supported by PIL. mode: The requested image mode e.g., "RGB", "HSV", "RGBA", "P" (8-bit pixels, using a color palette). For details, see https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes Returns: bytes: PNG bytes representation of the image. """pil_image=load_pil_image_from_bytes(image_bytes,mode=mode)returncreate_png_bytes_from_image(pil_image)
[docs]defload_image_png_bytes_from_path(input_image_filepath:Union[str,Path],mode:str=DEFAULT_IMAGE_MODE)->bytes:"""Loads an image from a path, converts it to PNG, and returns image bytes. Args: input_image_filepath: A file path of an image. The image can be in any format supported by PIL. mode: The requested image mode e.g., "RGB", "HSV", "RGBA", "P" (8-bit pixels, using a color palette). For details, see https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes Returns: bytes: PNG bytes representation of the image. """pil_image=load_pil_image_from_path(input_image_filepath,mode=mode)returncreate_png_bytes_from_image(pil_image)
[docs]defload_image_png_bytes_from_url(input_image_url:str,mode:str=DEFAULT_IMAGE_MODE)->bytes:"""Loads an image from a URL, converts it to PNG, and returns image bytes. Args: input_image_url: An image URL. mode: The requested image mode e.g., "RGB", "HSV", "RGBA", "P" (8-bit pixels, using a color palette). For details, see https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes Returns: bytes: PNG bytes representation of the image. """pil_image=load_pil_image_from_url(input_image_url,mode=mode)returncreate_png_bytes_from_image(pil_image)