utils/image_helpers.py

"""
Image manipulation helpers for resizing, positioning, and compositing
the animal subject onto background scenes.
"""

from PIL import Image, ImageFilter, ImageEnhance, ImageDraw
import numpy as np


def resize_to_fill(img: Image.Image, target_size: tuple[int, int]) -> Image.Image:
    tw, th = target_size
    iw, ih = img.size
    scale = max(tw / iw, th / ih)
    new_w = int(iw * scale)
    new_h = int(ih * scale)
    img = img.resize((new_w, new_h), Image.LANCZOS)
    left = (new_w - tw) // 2
    top = (new_h - th) // 2
    return img.crop((left, top, left + tw, top + th))


def resize_to_fit(img: Image.Image, max_size: tuple[int, int]) -> Image.Image:
    mw, mh = max_size
    iw, ih = img.size
    scale = min(mw / iw, mh / ih)
    new_w = int(iw * scale)
    new_h = int(ih * scale)
    return img.resize((new_w, new_h), Image.LANCZOS)


def feather_bottom_edge(subject: Image.Image, fade_height: int = 60) -> Image.Image:
    if subject.mode != "RGBA":
        subject = subject.convert("RGBA")

    arr = np.array(subject)
    h = arr.shape[0]
    fade_start = max(0, h - fade_height)

    for y in range(fade_start, h):
        t = (y - fade_start) / fade_height
        alpha_mult = 1.0 - (t * t)
        arr[y, :, 3] = (arr[y, :, 3] * alpha_mult).astype(np.uint8)

    return Image.fromarray(arr)


def add_soft_glow(subject: Image.Image, radius: int = 8, opacity: int = 30) -> Image.Image:
    if subject.mode != "RGBA":
        subject = subject.convert("RGBA")

    glow = subject.copy()
    glow_arr = np.array(glow)
    glow_arr[:, :, :3] = 255
    glow_arr[:, :, 3] = np.minimum(glow_arr[:, :, 3], opacity)
    glow = Image.fromarray(glow_arr)
    glow = glow.filter(ImageFilter.GaussianBlur(radius=radius))

    result = Image.new("RGBA", subject.size, (0, 0, 0, 0))
    result = Image.alpha_composite(result, glow)
    result = Image.alpha_composite(result, subject)
    return result


def color_match_subject(subject: Image.Image, background: Image.Image) -> Image.Image:
    if subject.mode != "RGBA":
        subject = subject.convert("RGBA")

    bg_arr = np.array(background.convert("RGB")).astype(float)
    avg_color = bg_arr.mean(axis=(0, 1))

    sub_arr = np.array(subject).astype(float)
    blend_strength = 0.08
    sub_arr[:, :, 0] = sub_arr[:, :, 0] * (1 - blend_strength) + avg_color[0] * blend_strength
    sub_arr[:, :, 1] = sub_arr[:, :, 1] * (1 - blend_strength) + avg_color[1] * blend_strength
    sub_arr[:, :, 2] = sub_arr[:, :, 2] * (1 - blend_strength) + avg_color[2] * blend_strength

    return Image.fromarray(np.clip(sub_arr, 0, 255).astype(np.uint8))


def composite_animal_on_background(
    subject_rgba: Image.Image,
    background: Image.Image,
    output_size: tuple[int, int] = (1024, 1024),
    subject_scale: float = 0.75,
) -> Image.Image:
    bg = resize_to_fill(background.convert("RGB"), output_size)
    bg = bg.convert("RGBA")

    max_subject_h = int(output_size[1] * subject_scale)
    max_subject_w = int(output_size[0] * 0.90)
    subject = resize_to_fit(subject_rgba.convert("RGBA"), (max_subject_w, max_subject_h))

    subject = color_match_subject(subject, bg)
    subject = feather_bottom_edge(subject, fade_height=int(subject.size[1] * 0.12))
    subject = add_soft_glow(subject, radius=10, opacity=20)

    sw, sh = subject.size
    x = (output_size[0] - sw) // 2
    y = output_size[1] - sh + int(sh * 0.03)

    bg.paste(subject, (x, y), subject)

    enhancer = ImageEnhance.Color(bg)
    bg = enhancer.enhance(1.03)
    enhancer = ImageEnhance.Brightness(bg)
    bg = enhancer.enhance(1.01)

    return bg.convert("RGB")