tsystems/visual_document_retrieval

1from PIL import Image2import numpy as np3import base644import io5from io import BytesIO6from PIL import Image, ImageFile7from pdf2image import convert_from_path8import tempfile9from multiprocessing import Pool10import os11from loguru import logger12import uuid13 14from typing import Any, List, Tuple, Type, Literal, Optional, Union, Dict15 16def encode_image(image_path):17  with open(image_path, "rb") as image_file:18    return base64.b64encode(image_file.read()).decode('utf-8')19 20def load_image_from_base64(image):21    return Image.open(BytesIO(base64.b64decode(image)))22 23def pil_image_to_base64(image: Image) -> str:24    """25    Convert a PIL Image object to its base64 representation.26 27    Args:28        image (Image): The PIL Image object to be converted.29 30    Returns:31        str: The base64 representation of the image.32    """33 34    # Create a bytes buffer35    buffer = io.BytesIO()36 37    # Save the image to the buffer38    image.save(buffer, format="PNG")39 40    # Get the bytes from the buffer41    img_bytes = buffer.getvalue()42 43    # Convert the bytes to base6444    img_base64 = base64.b64encode(img_bytes).decode("utf-8")45 46    return img_base6447 48def scale_image(image: Image.Image, new_height: int = 1024) -> Image.Image:49    """50    Scale an image to a new height while maintaining the aspect ratio.51    """52    width, height = image.size53    aspect_ratio = width / height54    new_width = int(new_height * aspect_ratio)55 56    scaled_image = image.resize((new_width, new_height))57 58    return scaled_image59 60def unflatten_array(flat_list, vector_size=128):61    return np.array(flat_list).reshape(-1, vector_size)62 63def get_image_embedding(image_list: list[Image], openai_client, model: str, flatten: bool = False) -> list:64    """65    Get the embedding of an image.66 67    Args:68        image (Image): The image to be embedded.69 70    Returns:71        list[list[float]] if flatten, 72        else: list[list[list[float]]] with shape = (number of images (m), number of vector for each text (n), vector dim = 128)73    """74    if not isinstance(image_list, list):75        image_list = [image_list]76 77    input_base64_list = [f"data:image/png;base64,{pil_image_to_base64(image)}" for image in image_list]78    # Get the embedding of the image79    embedding = openai_client.embeddings.create(80        input=input_base64_list,81        model=model,82        extra_body={83            "modality": "image",84            "encoding_format":"float" if not flatten else "base64",85        },86    )87 88    result = []89    for embed in embedding.data:90        result.append(embed.embedding) # embed.embedding is a list[float] in case of flatten, else: list[list[float]]91    return result92 93def get_text_embedding(texts: list[str], openai_client, model: str, flatten: bool = False) -> list:94    """95    Get the embedding of a text.96 97    Args:98        text (str): The text to be embedded.99 100    Returns:101        list[list[float]] if flatten, 102        else: list[list[list[float]]] with shape = (number of texts (m), number of vector for each text (n), vector dim = 128)103    """104    if not isinstance(texts, list):105        texts = [texts]106 107    # Get the embedding of the text108    embedding = openai_client.embeddings.create(109        input=texts,110        model=model,111        extra_body={112            "encoding_format":"float" if not flatten else "base64",113        },114    )115 116    result = []117    for embed in embedding.data:118        result.append(embed.embedding) # embed.embedding is a list[float] in case of flatten, else: list[list[float]]119    return result120 121def load_images(image_paths):122    """123    Load images from a list of paths and return a list of PIL image objects.124 125    Args:126        image_paths (list): List of image paths.127 128    Returns:129        list: List of PIL image objects.130    """131    images = []132    for path in image_paths:133        try:134            img = Image.open(path)135            images.append(img)136        except Exception as e:137            logger.error(f"Error loading image at path {path}: {str(e)}")138    return images139    140 141def process_pdf(pdf_path: str, output_folder: str, thread_count=1):142    result_image_paths = []143 144    with tempfile.TemporaryDirectory() as temp_dir:145        images = convert_from_path(pdf_path, dpi=200, output_folder=temp_dir, thread_count=thread_count)146 147    # for page_num, image in enumerate(images):148    #     image_filename = f"{str(uuid.uuid4())}.png"149    #     image_path = os.path.join(output_folder, image_filename)150    #     image.save(image_path, "PNG")151    #     result_image_paths.append(image_path)152    153    # del images154    # return result_image_paths155    return images156 157 158def pdf_folder_to_images(pdf_folder: str, output_folder: str, process_count: int = 2):159    try:160        if process_count is None:161            process_count = os.cpu_count()162 163        pdf_files = [os.path.join(pdf_folder, f) for f in os.listdir(pdf_folder)164                     if f.lower().endswith('.pdf')]165        166        # Create a list of tuples containing (pdf_file, output_folder)167        args = [(pdf_file, output_folder) for pdf_file in pdf_files]168        169        with Pool(process_count) as pool:170            all_images = pool.starmap(process_pdf, args)171        172        result = [img for sublist in all_images for img in sublist]173 174        logger.debug(f"Number of pdfs processed: {len(all_images)} - Number of images: {len(result)}")175        return result176    except Exception as e:177        logger.exception(f"Error during processing pdf: {e}")178 179 180