Spaces:
Runtime error
Runtime error
# processing.py | |
import os | |
import random | |
import textwrap | |
import logging | |
from typing import List, Optional | |
from PIL import Image, ImageDraw, ImageFont, ImageOps, ImageFilter, ImageEnhance, UnidentifiedImageError | |
from config import Config | |
from utils import add_noise_to_image | |
logger = logging.getLogger(__name__) | |
def apply_template(user_image_path: str, caption: str, template_path: str) -> Optional[str]: | |
""" | |
Applies user image and caption to a predefined template. | |
Assumes template_path points to a PNG with a transparent "window" for the user image, | |
and other graphics are the foreground. | |
""" | |
output_filename = f"result_{os.path.basename(template_path).split('.')[0]}_{random.randint(10000, 99999)}.jpg" | |
output_path = os.path.join(Config.OUTPUT_DIR, output_filename) | |
logger.debug(f"Applying foreground template '{os.path.basename(template_path)}' over user image '{os.path.basename(user_image_path)}'. Output: {output_path}") | |
try: | |
with Image.open(user_image_path).convert("RGBA") as user_image_orig, \ | |
Image.open(template_path).convert("RGBA") as foreground_graphic_template: | |
logger.debug(f"User image size: {user_image_orig.size}, Foreground template size: {foreground_graphic_template.size}") | |
# 1. Resize user's image to fit the placeholder size | |
logger.debug(f"Resizing user image to placeholder size {Config.PLACEHOLDER_SIZE}") | |
user_image_resized = ImageOps.fit(user_image_orig, Config.PLACEHOLDER_SIZE, Image.Resampling.LANCZOS) | |
logger.debug(f"User image resized to: {user_image_resized.size}") | |
# 2. Create a new transparent base canvas of the final template size | |
final_base_image = Image.new("RGBA", Config.TEMPLATE_SIZE, (0, 0, 0, 0)) | |
logger.debug(f"Created base canvas of size {Config.TEMPLATE_SIZE}") | |
# 3. Paste the user's resized image onto this base canvas at the placeholder position | |
logger.debug(f"Pasting resized user image onto base canvas at {Config.PLACEHOLDER_POSITION}") | |
final_base_image.paste(user_image_resized, Config.PLACEHOLDER_POSITION, user_image_resized if user_image_resized.mode == 'RGBA' else None) | |
# 4. Ensure your Canva template (foreground_graphic_template) is the correct final size | |
if foreground_graphic_template.size != Config.TEMPLATE_SIZE: | |
logger.warning(f"Foreground template size {foreground_graphic_template.size} " | |
f"differs from expected {Config.TEMPLATE_SIZE}. Resizing it.") | |
foreground_graphic_template = foreground_graphic_template.resize(Config.TEMPLATE_SIZE, Image.Resampling.LANCZOS) | |
# 5. Composite the Canva template ON TOP of the base image | |
logger.debug("Alpha compositing user image base with foreground Canva template.") | |
combined_image = Image.alpha_composite(final_base_image, foreground_graphic_template) | |
logger.debug(f"Combined image size: {combined_image.size}, mode: {combined_image.mode}") | |
# --- Add Caption Text (on top of everything) --- | |
logger.debug("Adding caption text to the combined image") | |
draw = ImageDraw.Draw(combined_image) | |
try: | |
font_size = Config.MAX_FONT_SIZE // 2 | |
font = ImageFont.truetype(Config.FONT_PATH, font_size) | |
logger.debug(f"Using font '{Config.FONT_PATH}' size {font_size}") | |
except IOError: | |
logger.warning(f"Failed loading font '{Config.FONT_PATH}'. Using Pillow's default.") | |
font = ImageFont.load_default() | |
wrapped_text = textwrap.fill(caption, width=Config.MAX_CAPTION_WIDTH) | |
# --- Adjust Text Position based on Config --- | |
text_bbox = draw.textbbox((0,0), wrapped_text, font=font, align="center") | |
text_width = text_bbox[2] - text_bbox[0] | |
text_height = text_bbox[3] - text_bbox[1] | |
text_x = (Config.TEMPLATE_SIZE[0] - text_width) // 2 # Centered horizontally on whole template | |
text_y = Config.TEXT_AREA_Y_START + (Config.TEXT_AREA_HEIGHT - text_height) // 2 # Centered vertically in defined text area | |
logger.debug(f"Calculated text position: ({text_x}, {text_y}) for text area starting at Y={Config.TEXT_AREA_Y_START} with height {Config.TEXT_AREA_HEIGHT}") | |
shadow_offset = 1 | |
draw.text((text_x + shadow_offset, text_y + shadow_offset), wrapped_text, font=font, fill="black", align="center") | |
draw.text((text_x, text_y), wrapped_text, font=font, fill="white", align="center") | |
final_output_image = combined_image.convert("RGB") | |
logger.debug(f"Saving final image to {output_path} with quality {Config.JPEG_QUALITY}") | |
final_output_image.save(output_path, "JPEG", quality=Config.JPEG_QUALITY) | |
logger.info(f"Generated final image: {output_path}") | |
return output_path | |
except FileNotFoundError: | |
logger.error(f"Template or user image not found. Template: '{template_path}', User Image: '{user_image_path}'") | |
except UnidentifiedImageError: | |
logger.error(f"Could not identify image file. Template: '{template_path}', User Image: '{user_image_path}'") | |
except Exception as e: | |
logger.error(f"Error applying template '{os.path.basename(template_path)}': {e}", exc_info=True) | |
return None | |
def create_auto_template(user_image_path: str, caption: str, variant: int) -> Optional[str]: | |
output_filename = f"auto_template_{variant}_{random.randint(10000, 99999)}.jpg" | |
output_path = os.path.join(Config.OUTPUT_DIR, output_filename) | |
logger.debug(f"Creating auto template variant {variant} -> {output_path}") | |
try: | |
bg_color = (random.randint(0, 150), random.randint(0, 150), random.randint(0, 150)) | |
bg = Image.new('RGB', Config.AUTO_TEMPLATE_SIZE, color=bg_color) | |
with Image.open(user_image_path) as user_img_orig: | |
user_img = user_img_orig.copy() | |
user_img = ImageOps.fit(user_img, Config.AUTO_USER_IMAGE_SIZE, Image.Resampling.LANCZOS) | |
effect_choice = random.choice(['blur', 'noise', 'color', 'contrast', 'sharpness', 'none']) | |
if effect_choice == 'blur': user_img = user_img.filter(ImageFilter.GaussianBlur(random.uniform(0.5, 1.8))) | |
elif effect_choice == 'noise': user_img = add_noise_to_image(user_img, Config.NOISE_INTENSITY) | |
elif effect_choice == 'color': user_img = ImageEnhance.Color(user_img).enhance(random.uniform(0.3, 1.7)) | |
elif effect_choice == 'contrast': user_img = ImageEnhance.Contrast(user_img).enhance(random.uniform(0.7, 1.4)) | |
elif effect_choice == 'sharpness': user_img = ImageEnhance.Sharpness(user_img).enhance(random.uniform(1.1, 2.0)) | |
border_color = (random.randint(180, 255), random.randint(180, 255), random.randint(180, 255)) | |
border_width = random.randint(8, 20) | |
user_img = ImageOps.expand(user_img, border=border_width, fill=border_color) | |
paste_x = (bg.width - user_img.width) // 2 | |
paste_y = (bg.height - user_img.height) // 3 | |
bg.paste(user_img, (paste_x, paste_y)) | |
draw = ImageDraw.Draw(bg) | |
try: | |
font_size = random.randint(Config.MIN_FONT_SIZE, Config.MAX_FONT_SIZE) | |
font = ImageFont.truetype(Config.FONT_PATH, font_size) | |
except IOError: | |
logger.warning(f"Failed loading font '{Config.FONT_PATH}'. Using default.") | |
font = ImageFont.load_default() | |
wrapped_text = textwrap.fill(caption, width=Config.MAX_CAPTION_WIDTH) | |
text_bbox = draw.textbbox((0, 0), wrapped_text, font=font, align="center") | |
text_x = (bg.width - (text_bbox[2] - text_bbox[0])) // 2 | |
text_y = paste_y + user_img.height + 30 # Position text below the image | |
text_color = (random.randint(200, 255), random.randint(200, 255), random.randint(200, 255)) | |
stroke_color = (random.randint(0, 50), random.randint(0, 50), random.randint(0, 50)) | |
draw.text((text_x, text_y), wrapped_text, font=font, fill=text_color, stroke_width=Config.TEXT_STROKE_WIDTH, stroke_fill=stroke_color, align="center") | |
bg.save(output_path, "JPEG", quality=Config.JPEG_QUALITY) | |
logger.info(f"Generated auto-template image: {output_path}") | |
return output_path | |
except FileNotFoundError: logger.error(f"User image not found: '{user_image_path}'") | |
except UnidentifiedImageError: logger.error(f"Unidentified user image: '{user_image_path}'") | |
except Exception as e: logger.error(f"Error creating auto-template {variant}: {e}", exc_info=True) | |
return None | |
def load_predefined_templates() -> List[str]: | |
templates = [] | |
template_dir = Config.PREDEFINED_TEMPLATES_DIR | |
logger.debug(f"Searching templates in: {os.path.abspath(template_dir)}") | |
try: | |
if not os.path.isdir(template_dir): | |
logger.warning(f"Templates directory not found: '{template_dir}'") | |
return [] | |
files = os.listdir(template_dir) | |
for file in files: | |
if file.lower().endswith(('.png', '.jpg', '.jpeg')): | |
full_path = os.path.join(template_dir, file) | |
if os.path.isfile(full_path): templates.append(full_path) | |
if not templates: logger.warning(f"No valid template files found in '{template_dir}'.") | |
else: logger.info(f"Loaded {len(templates)} templates.") | |
except Exception as e: logger.error(f"Error loading templates: {e}", exc_info=True) | |
return templates | |
def process_images_sync(user_image_path: str, caption: str) -> List[str]: | |
logger.info("Starting sync image processing task...") | |
generated_image_paths: List[str] = [] | |
predefined_templates = load_predefined_templates() | |
if predefined_templates: | |
for template_path in predefined_templates: | |
result_path = apply_template(user_image_path, caption, template_path) | |
if result_path: generated_image_paths.append(result_path) | |
for i in range(Config.AUTO_TEMPLATES_COUNT): | |
result_path = create_auto_template(user_image_path, caption, i) | |
if result_path: generated_image_paths.append(result_path) | |
logger.info(f"Sync processing finished. Generated {len(generated_image_paths)} images.") | |
return generated_image_paths | |