gpt4free

1
import re
2
from io import BytesIO
3
import base64
4
from .typing import ImageType, Union
5
from PIL import Image
6

7
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'webp', 'svg'}
8

9
def to_image(image: ImageType, is_svg: bool = False) -> Image.Image:
10
    """
11
    Converts the input image to a PIL Image object.
12

13
    Args:
14
        image (Union[str, bytes, Image.Image]): The input image.
15

16
    Returns:
17
        Image.Image: The converted PIL Image object.
18
    """
19
    if is_svg:
20
        try:
21
            import cairosvg
22
        except ImportError:
23
            raise RuntimeError('Install "cairosvg" package for open svg images')
24
        if not isinstance(image, bytes):
25
            image = image.read()
26
        buffer = BytesIO()
27
        cairosvg.svg2png(image, write_to=buffer)
28
        image = Image.open(buffer)
29
    if isinstance(image, str):
30
        is_data_uri_an_image(image)
31
        image = extract_data_uri(image)
32
    if isinstance(image, bytes):
33
        is_accepted_format(image)
34
        image = Image.open(BytesIO(image))
35
    elif not isinstance(image, Image.Image):
36
        image = Image.open(image)
37
        copy = image.copy()
38
        copy.format = image.format
39
        image = copy
40
    return image
41

42
def is_allowed_extension(filename: str) -> bool:
43
    """
44
    Checks if the given filename has an allowed extension.
45

46
    Args:
47
        filename (str): The filename to check.
48

49
    Returns:
50
        bool: True if the extension is allowed, False otherwise.
51
    """
52
    return '.' in filename and \
53
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
54

55
def is_data_uri_an_image(data_uri: str) -> bool:
56
    """
57
    Checks if the given data URI represents an image.
58

59
    Args:
60
        data_uri (str): The data URI to check.
61

62
    Raises:
63
        ValueError: If the data URI is invalid or the image format is not allowed.
64
    """
65
    # Check if the data URI starts with 'data:image' and contains an image format (e.g., jpeg, png, gif)
66
    if not re.match(r'data:image/(\w+);base64,', data_uri):
67
        raise ValueError("Invalid data URI image.")
68
    # Extract the image format from the data URI
69
    image_format = re.match(r'data:image/(\w+);base64,', data_uri).group(1)
70
    # Check if the image format is one of the allowed formats (jpg, jpeg, png, gif)
71
    if image_format.lower() not in ALLOWED_EXTENSIONS:
72
        raise ValueError("Invalid image format (from mime file type).")
73

74
def is_accepted_format(binary_data: bytes) -> bool:
75
    """
76
    Checks if the given binary data represents an image with an accepted format.
77

78
    Args:
79
        binary_data (bytes): The binary data to check.
80

81
    Raises:
82
        ValueError: If the image format is not allowed.
83
    """
84
    if binary_data.startswith(b'\xFF\xD8\xFF'):
85
        pass # It's a JPEG image
86
    elif binary_data.startswith(b'\x89PNG\r\n\x1a\n'):
87
        pass # It's a PNG image
88
    elif binary_data.startswith(b'GIF87a') or binary_data.startswith(b'GIF89a'):
89
        pass # It's a GIF image
90
    elif binary_data.startswith(b'\x89JFIF') or binary_data.startswith(b'JFIF\x00'):
91
        pass # It's a JPEG image
92
    elif binary_data.startswith(b'\xFF\xD8'):
93
        pass # It's a JPEG image
94
    elif binary_data.startswith(b'RIFF') and binary_data[8:12] == b'WEBP':
95
        pass # It's a WebP image
96
    else:
97
        raise ValueError("Invalid image format (from magic code).")
98

99
def extract_data_uri(data_uri: str) -> bytes:
100
    """
101
    Extracts the binary data from the given data URI.
102

103
    Args:
104
        data_uri (str): The data URI.
105

106
    Returns:
107
        bytes: The extracted binary data.
108
    """
109
    data = data_uri.split(",")[1]
110
    data = base64.b64decode(data)
111
    return data
112

113
def get_orientation(image: Image.Image) -> int:
114
    """
115
    Gets the orientation of the given image.
116

117
    Args:
118
        image (Image.Image): The image.
119

120
    Returns:
121
        int: The orientation value.
122
    """
123
    exif_data = image.getexif() if hasattr(image, 'getexif') else image._getexif()
124
    if exif_data is not None:
125
        orientation = exif_data.get(274) # 274 corresponds to the orientation tag in EXIF
126
        if orientation is not None:
127
            return orientation
128

129
def process_image(img: Image.Image, new_width: int, new_height: int) -> Image.Image:
130
    """
131
    Processes the given image by adjusting its orientation and resizing it.
132

133
    Args:
134
        img (Image.Image): The image to process.
135
        new_width (int): The new width of the image.
136
        new_height (int): The new height of the image.
137

138
    Returns:
139
        Image.Image: The processed image.
140
    """
141
    orientation = get_orientation(img)
142
    if orientation:
143
        if orientation > 4:
144
            img = img.transpose(Image.FLIP_LEFT_RIGHT)
145
        if orientation in [3, 4]:
146
            img = img.transpose(Image.ROTATE_180)
147
        if orientation in [5, 6]:
148
            img = img.transpose(Image.ROTATE_270)
149
        if orientation in [7, 8]:
150
            img = img.transpose(Image.ROTATE_90)
151
    img.thumbnail((new_width, new_height))
152
    return img
153

154
def to_base64(image: Image.Image, compression_rate: float) -> str:
155
    """
156
    Converts the given image to a base64-encoded string.
157

158
    Args:
159
        image (Image.Image): The image to convert.
160
        compression_rate (float): The compression rate (0.0 to 1.0).
161

162
    Returns:
163
        str: The base64-encoded image.
164
    """
165
    output_buffer = BytesIO()
166
    if image.mode != "RGB":
167
        image = image.convert('RGB')
168
    image.save(output_buffer, format="JPEG", quality=int(compression_rate * 100))
169
    return base64.b64encode(output_buffer.getvalue()).decode()
170

171
def format_images_markdown(images, alt: str, preview: str="{image}?w=200&h=200") -> str:
172
    """
173
    Formats the given images as a markdown string.
174

175
    Args:
176
        images: The images to format.
177
        alt (str): The alt for the images.
178
        preview (str, optional): The preview URL format. Defaults to "{image}?w=200&h=200".
179

180
    Returns:
181
        str: The formatted markdown string.
182
    """
183
    if isinstance(images, str):
184
        images = f"[![{alt}]({preview.replace('{image}', images)})]({images})"
185
    else:
186
        images = [f"[![#{idx+1} {alt}]({preview.replace('{image}', image)})]({image})" for idx, image in enumerate(images)]
187
        images = "\n".join(images)
188
    start_flag = "<!-- generated images start -->\n"
189
    end_flag = "<!-- generated images end -->\n"
190
    return f"\n{start_flag}{images}\n{end_flag}\n"
191

192
def to_bytes(image: Image.Image) -> bytes:
193
    """
194
    Converts the given image to bytes.
195

196
    Args:
197
        image (Image.Image): The image to convert.
198

199
    Returns:
200
        bytes: The image as bytes.
201
    """
202
    bytes_io = BytesIO()
203
    image.save(bytes_io, image.format)
204
    image.seek(0)
205
    return bytes_io.getvalue()
206

207
class ImageResponse():
208
    def __init__(
209
        self,
210
        images: Union[str, list],
211
        alt: str,
212
        options: dict = {}
213
    ):
214
        self.images = images
215
        self.alt = alt
216
        self.options = options
217
        
218
    def __str__(self) -> str:
219
        return format_images_markdown(self.images, self.alt)
220
    
221
    def get(self, key: str):
222
        return self.options.get(key)