Spaces:
Running
Running
| from PIL import Image, ImageOps | |
| import io | |
| def parse_aspect_ratio(aspect_ratio_str: str) -> float: | |
| """ | |
| Parse aspect ratio string like "16:9" or "9:16" to float. | |
| """ | |
| parts = aspect_ratio_str.split(':') | |
| if len(parts) == 2: | |
| return float(parts[0]) / float(parts[1]) | |
| return 16.0 / 9.0 # fallback | |
| def _parse_hex_color(fill_color: str) -> tuple: | |
| """ | |
| Safely parse hex color like "#RRGGBB" to an RGB tuple for PIL. | |
| Falls back to black if parsing fails. | |
| """ | |
| if not fill_color: | |
| return (0, 0, 0) | |
| s = fill_color.strip() | |
| try: | |
| if s.startswith("#"): | |
| s = s[1:] | |
| if len(s) == 6: | |
| r = int(s[0:2], 16) | |
| g = int(s[2:4], 16) | |
| b = int(s[4:6], 16) | |
| return (r, g, b) | |
| except Exception: | |
| pass | |
| return (0, 0, 0) | |
| def process_image(image_bytes: bytes, mode: str, fill_color: str = "#000000", target_aspect_ratio: str = "16:9") -> bytes: | |
| """ | |
| Processes an image to fit the specified aspect ratio. | |
| Args: | |
| image_bytes: The original image in bytes. | |
| mode: The processing mode ('stretch_to_fill', 'compress_to_fit', 'fill'). | |
| fill_color: The hex color code for the 'fill' mode background. | |
| target_aspect_ratio: The target aspect ratio string like "16:9" or "9:16". | |
| Returns: | |
| The processed image in bytes (JPEG format). | |
| """ | |
| if not mode or mode == 'none': | |
| return image_bytes | |
| try: | |
| target_ratio = parse_aspect_ratio(target_aspect_ratio) | |
| image = Image.open(io.BytesIO(image_bytes)) | |
| # Normalize EXIF orientation to ensure width/height reflect actual display | |
| try: | |
| image = ImageOps.exif_transpose(image) | |
| except Exception: | |
| pass | |
| original_width, original_height = image.size | |
| original_aspect_ratio = original_width / original_height | |
| # If aspect ratio is already correct, no processing needed | |
| if abs(original_aspect_ratio - target_ratio) < 1e-6: | |
| return image_bytes | |
| if mode == 'stretch_to_fill': | |
| if original_aspect_ratio > target_ratio: # Wider than target | |
| # Height is relatively too short, stretch it | |
| new_width = original_width | |
| new_height = int(new_width / target_ratio) | |
| processed_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS) | |
| else: # Taller than target | |
| # Width is relatively too short, stretch it | |
| new_height = original_height | |
| new_width = int(new_height * target_ratio) | |
| processed_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS) | |
| elif mode == 'compress_to_fit': | |
| if original_aspect_ratio > target_ratio: # Wider than target | |
| # Width is relatively too long, compress it | |
| new_height = original_height | |
| new_width = int(new_height * target_ratio) | |
| processed_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS) | |
| else: # Taller than target | |
| # Height is relatively too long, compress it | |
| new_width = original_width | |
| new_height = int(new_width / target_ratio) | |
| processed_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS) | |
| elif mode == 'fill': | |
| # Pad with color to meet target aspect ratio without scaling the source | |
| if original_aspect_ratio > target_ratio: | |
| # Wider than target -> pad top/bottom | |
| new_width = original_width | |
| new_height = int(round(new_width / target_ratio)) | |
| paste_x, paste_y = 0, (new_height - original_height) // 2 | |
| else: | |
| # Taller than target -> pad left/right | |
| new_height = original_height | |
| new_width = int(round(new_height * target_ratio)) | |
| paste_x, paste_y = (new_width - original_width) // 2, 0 | |
| # Use parsed RGB to avoid mode issues and ensure consistent color | |
| fill_rgb = _parse_hex_color(fill_color or "#000000") | |
| background = Image.new('RGB', (new_width, new_height), fill_rgb) | |
| # Preserve alpha if present so the background color shows through transparent regions | |
| if image.mode in ('RGBA', 'LA') or (image.mode == 'P' and 'transparency' in image.info): | |
| if image.mode != 'RGBA': | |
| image = image.convert('RGBA') | |
| background.paste(image, (paste_x, paste_y), image) | |
| else: | |
| background.paste(image, (paste_x, paste_y)) | |
| processed_image = background | |
| else: | |
| # If mode is unknown, return original image | |
| return image_bytes | |
| # Save the processed image to a byte buffer | |
| byte_arr = io.BytesIO() | |
| # Ensure image is in RGB mode before saving as JPEG | |
| if processed_image.mode != 'RGB': | |
| processed_image = processed_image.convert('RGB') | |
| processed_image.save(byte_arr, format='JPEG') | |
| return byte_arr.getvalue() | |
| except Exception as e: | |
| print(f"Error processing image: {e}") | |
| # In case of error, return the original image bytes | |
| return image_bytes |