utils/background.py

"""
Background removal using rembg and background scene selection.
"""

from PIL import Image
from rembg import remove
import os
import random

BACKGROUNDS_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "backgrounds")

SCENE_PREFERENCES = {
    "Dog": ["backyard", "park", "porch", "meadow"],
    "Cat": ["living_room", "couch", "garden", "porch"],
    "Other": ["park", "garden", "meadow", "backyard"],
}


def remove_background(image: Image.Image) -> Image.Image:
    return remove(image)


def get_available_backgrounds() -> list[str]:
    if not os.path.isdir(BACKGROUNDS_DIR):
        return []
    return [
        os.path.splitext(f)[0]
        for f in os.listdir(BACKGROUNDS_DIR)
        if f.endswith(".png")
    ]


def select_background(animal_type: str, custom_bg: Image.Image | None = None) -> Image.Image:
    if custom_bg is not None:
        return custom_bg.convert("RGB")

    available = get_available_backgrounds()
    if not available:
        return _generate_fallback_background()

    preferences = SCENE_PREFERENCES.get(animal_type, SCENE_PREFERENCES["Other"])
    candidates = [name for name in preferences if name in available]

    if not candidates:
        candidates = available

    chosen = random.choice(candidates)
    path = os.path.join(BACKGROUNDS_DIR, f"{chosen}.png")
    return Image.open(path).convert("RGB")


def _generate_fallback_background() -> Image.Image:
    from PIL import ImageDraw
    img = Image.new("RGB", (1024, 1024))
    draw = ImageDraw.Draw(img)
    for y in range(1024):
        t = y / 1024
        r = int(135 + (185 - 135) * t)
        g = int(195 + (220 - 195) * t)
        b = int(235 + (245 - 235) * t)
        draw.line([(0, y), (1024, y)], fill=(r, g, b))
    return img