attribution-2steps-method / app /utils /image_processing.py
thanhnt-cf's picture
update demo app
9645c29
import base64
import io
from PIL import Image
import requests
def get_image_data(image_path):
with open(image_path, "rb") as f:
image_data = base64.b64encode(f.read()).decode("utf-8")
return image_data
def get_data_format(image_path):
image_format = image_path.split(".")[-1]
if image_format == "jpg":
image_format = "jpeg"
return
def get_image_base64_and_type(image_url: str, max_dimension: int = 2048) -> tuple[str | None, str | None]:
try:
# --- 1. Download the image ---
response = requests.get(image_url, stream=True, timeout=20) # Added timeout
response.raise_for_status() # Raise an exception for bad status codes (4xx or 5xx)
# Check content type
content_type = response.headers.get('content-type')
allowed_types = ['image/png', 'image/jpeg', 'image/webp', 'image/gif']
if not content_type or content_type not in allowed_types:
raise ValueError(f"Unsupported image type: {content_type}. Expected one of {allowed_types}.")
# --- 2. Open the image using Pillow ---
image_data = response.content
img = Image.open(io.BytesIO(image_data))
# Check if the image is animated (GIF)
if img.format == 'GIF' and getattr(img, 'is_animated', False):
raise ValueError("Animated GIFs are not supported.")
# --- 3. Check dimensions and resize if necessary ---
width, height = img.size
longest_dim = max(width, height)
if longest_dim > max_dimension:
# print(f"Image dimensions ({width}x{height}) exceed max dimension ({max_dimension}). Resizing...")
if width > height:
# Width is the longest dimension
new_width = max_dimension
new_height = int(height * (max_dimension / width))
else:
# Height is the longest or they are equal
new_height = max_dimension
new_width = int(width * (max_dimension / height))
# Resize the image - Use Resampling.LANCZOS for high-quality downscaling
img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
# print(f"Image resized to: {img.size}"
width, height = img.size
shortest_dim = min(width, height)
if shortest_dim > 768:
if width < height:
new_width = 768
new_height = int(height * (768 / width))
else:
new_height = 768
new_width = int(width * (768 / height))
img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
# --- 4. Save the image to a byte buffer ---
# We need to save the potentially modified image back to bytes
buffer = io.BytesIO()
# Save with the JPG format. Handle potential format issues.
try:
img_format = 'JPEG'
img.save(buffer, format=img_format, quality=100)
except Exception as save_err:
try:
# Fallback to PNG if original format saving fails
img_format = 'PNG'
img.save(buffer, format=img_format)
except Exception as png_save_err:
raise Exception(f"Failed to save image in PNG format. Error: {png_save_err}")
image_bytes = buffer.getvalue()
# --- 5. Encode the image bytes to base64 ---
base64_encoded_image = base64.b64encode(image_bytes).decode('utf-8')
return base64_encoded_image, img_format
except Exception as e:
raise ValueError(f"Invalid image URL: {e}")