aphilippov/python-server-api

1import base642import mimetypes3import re4from io import BytesIO5from typing import Any, Dict, List, Optional, Tuple, Union6 7import requests8from PIL import Image9 10 11def get_image_data(image_file: str, use_b64=True) -> bytes:12    if image_file.startswith("http://") or image_file.startswith("https://"):13        response = requests.get(image_file)14        content = response.content15    elif re.match(r"data:image/(?:png|jpeg);base64,", image_file):16        return re.sub(r"data:image/(?:png|jpeg);base64,", "", image_file)17    else:18        image = Image.open(image_file).convert("RGB")19        buffered = BytesIO()20        image.save(buffered, format="PNG")21        content = buffered.getvalue()22 23    if use_b64:24        return base64.b64encode(content).decode("utf-8")25    else:26        return content27 28 29def llava_formatter(prompt: str, order_image_tokens: bool = False) -> Tuple[str, List[str]]:30    """31    Formats the input prompt by replacing image tags and returns the new prompt along with image locations.32 33    Parameters:34        - prompt (str): The input string that may contain image tags like <img ...>.35        - order_image_tokens (bool, optional): Whether to order the image tokens with numbers.36            It will be useful for GPT-4V. Defaults to False.37 38    Returns:39        - Tuple[str, List[str]]: A tuple containing the formatted string and a list of images (loaded in b64 format).40    """41 42    # Initialize variables43    new_prompt = prompt44    image_locations = []45    images = []46    image_count = 047 48    # Regular expression pattern for matching <img ...> tags49    img_tag_pattern = re.compile(r"<img ([^>]+)>")50 51    # Find all image tags52    for match in img_tag_pattern.finditer(prompt):53        image_location = match.group(1)54 55        try:56            img_data = get_image_data(image_location)57        except Exception as e:58            # Remove the token59            print(f"Warning! Unable to load image from {image_location}, because of {e}")60            new_prompt = new_prompt.replace(match.group(0), "", 1)61            continue62 63        image_locations.append(image_location)64        images.append(img_data)65 66        # Increment the image count and replace the tag in the prompt67        new_token = f"<image {image_count}>" if order_image_tokens else "<image>"68 69        new_prompt = new_prompt.replace(match.group(0), new_token, 1)70        image_count += 171 72    return new_prompt, images73 74 75def convert_base64_to_data_uri(base64_image):76    def _get_mime_type_from_data_uri(base64_image):77        # Decode the base64 string78        image_data = base64.b64decode(base64_image)79        # Check the first few bytes for known signatures80        if image_data.startswith(b"\xff\xd8\xff"):81            return "image/jpeg"82        elif image_data.startswith(b"\x89PNG\r\n\x1a\n"):83            return "image/png"84        elif image_data.startswith(b"GIF87a") or image_data.startswith(b"GIF89a"):85            return "image/gif"86        elif image_data.startswith(b"RIFF") and image_data[8:12] == b"WEBP":87            return "image/webp"88        return "image/jpeg"  # use jpeg for unknown formats, best guess.89 90    mime_type = _get_mime_type_from_data_uri(base64_image)91    data_uri = f"data:{mime_type};base64,{base64_image}"92    return data_uri93 94 95def gpt4v_formatter(prompt: str) -> List[Union[str, dict]]:96    """97    Formats the input prompt by replacing image tags and returns a list of text and images.98 99    Parameters:100        - prompt (str): The input string that may contain image tags like <img ...>.101 102    Returns:103        - List[Union[str, dict]]: A list of alternating text and image dictionary items.104    """105    output = []106    last_index = 0107    image_count = 0108 109    # Regular expression pattern for matching <img ...> tags110    img_tag_pattern = re.compile(r"<img ([^>]+)>")111 112    # Find all image tags113    for match in img_tag_pattern.finditer(prompt):114        image_location = match.group(1)115 116        try:117            img_data = get_image_data(image_location)118        except Exception as e:119            # Warning and skip this token120            print(f"Warning! Unable to load image from {image_location}, because {e}")121            continue122 123        # Add text before this image tag to output list124        output.append({"type": "text", "text": prompt[last_index : match.start()]})125 126        # Add image data to output list127        output.append({"type": "image_url", "image_url": {"url": convert_base64_to_data_uri(img_data)}})128 129        last_index = match.end()130        image_count += 1131 132    # Add remaining text to output list133    output.append({"type": "text", "text": prompt[last_index:]})134    return output135 136 137def extract_img_paths(paragraph: str) -> list:138    """139    Extract image paths (URLs or local paths) from a text paragraph.140 141    Parameters:142        paragraph (str): The input text paragraph.143 144    Returns:145        list: A list of extracted image paths.146    """147    # Regular expression to match image URLs and file paths148    img_path_pattern = re.compile(149        r"\b(?:http[s]?://\S+\.(?:jpg|jpeg|png|gif|bmp)|\S+\.(?:jpg|jpeg|png|gif|bmp))\b", re.IGNORECASE150    )151 152    # Find all matches in the paragraph153    img_paths = re.findall(img_path_pattern, paragraph)154    return img_paths155 156 157def _to_pil(data: str) -> Image.Image:158    """159    Converts a base64 encoded image data string to a PIL Image object.160 161    This function first decodes the base64 encoded string to bytes, then creates a BytesIO object from the bytes,162    and finally creates and returns a PIL Image object from the BytesIO object.163 164    Parameters:165        data (str): The base64 encoded image data string.166 167    Returns:168        Image.Image: The PIL Image object created from the input data.169    """170    return Image.open(BytesIO(base64.b64decode(data)))171