Spaces:
Build error
Build error
File size: 3,452 Bytes
5ce8318 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
import base64
from PIL import Image
from io import BytesIO
from fastapi import HTTPException
def base64_to_image(base64_str: str) -> Image.Image:
"""Convert base64 string to PIL Image.
Args:
base64_str: Base64 encoded image string
Returns:
PIL.Image: Decoded image
Raises:
HTTPException: If base64 string is invalid
"""
try:
# Handle frontend base64 format (data:image/jpeg;base64,{base64_data})
if "," in base64_str:
base64_str = base64_str.split(",", 1)[1]
image_data = base64.b64decode(base64_str)
image = Image.open(BytesIO(image_data))
# Convert RGBA to RGB if necessary
if image.mode in ('RGBA', 'LA'):
background = Image.new('RGB', image.size, (255, 255, 255))
if image.mode == 'RGBA':
background.paste(image, mask=image.split()[3]) # 3 is the alpha channel
else:
background.paste(image, mask=image.split()[1]) # 1 is the alpha channel
image = background
elif image.mode != 'RGB':
image = image.convert('RGB')
return image
except Exception as e:
print(f"Base64 decoding error: {str(e)}")
raise HTTPException(status_code=400, detail=f"Invalid Base64 image: {str(e)}")
def image_to_base64(image: Image.Image, format: str = "JPEG") -> str:
"""Convert PIL Image to base64 string.
Args:
image: PIL Image object
format: Output format (JPEG, PNG, etc.)
Returns:
str: Base64 encoded image string
"""
try:
# Convert RGBA to RGB if saving as JPEG
if format.upper() == "JPEG" and image.mode in ('RGBA', 'LA'):
background = Image.new('RGB', image.size, (255, 255, 255))
if image.mode == 'RGBA':
background.paste(image, mask=image.split()[3])
else:
background.paste(image, mask=image.split()[1])
image = background
elif format.upper() == "JPEG" and image.mode != 'RGB':
image = image.convert('RGB')
buffered = BytesIO()
image.save(buffered, format=format)
return base64.b64encode(buffered.getvalue()).decode("utf-8")
except Exception as e:
print(f"Error converting image to base64: {str(e)}")
# Try JPEG as fallback
if format.upper() != "JPEG":
return image_to_base64(image, format="JPEG")
raise
def is_image_file(filename: str) -> bool:
"""Check if a filename has a valid image extension.
Args:
filename: Name of the file to check
Returns:
bool: True if file has valid image extension
"""
valid_extensions = (".png", ".jpg", ".jpeg", ".bmp", ".gif", ".tiff", ".webp")
return filename.lower().endswith(valid_extensions)
def get_image_format(filename: str) -> str:
"""Get the format to use for saving an image based on its filename.
Args:
filename: Name of the file
Returns:
str: Format to use (JPEG, PNG, etc.)
"""
ext = filename.lower().split('.')[-1]
if ext in ('jpg', 'jpeg'):
return 'JPEG'
elif ext == 'png':
return 'PNG'
elif ext == 'webp':
return 'WEBP'
elif ext == 'gif':
return 'GIF'
else:
return 'JPEG' # Default to JPEG |