app.py

"""
Shelter Buddy: Adoption Booster
A Gradio app for animal shelters to create beautiful adoption flyers.
"""

import os
os.environ.setdefault("NUMBA_DISABLE_JIT", "1")

from dotenv import load_dotenv
load_dotenv()

import gradio as gr
from PIL import Image
import logging

from utils.background import remove_background, select_background
from utils.image_helpers import composite_animal_on_background
from utils.bio_generator import generate_bio

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

CHARACTERISTICS = [
    "Good with Kids",
    "Good with other pets",
    "High Energy",
    "Couch Potato",
    "Senior",
    "Puppy/Kitten",
    "House Trained",
    "Special Needs",
    "Smart",
    "Cuddly",
]

OUTPUT_SIZE = (1024, 1024)


def process_animal(
    original_photo: Image.Image | None,
    name: str,
    animal_type: str,
    characteristics: list[str],
    custom_background: Image.Image | None,
) -> tuple[Image.Image | None, str]:
    """
    Main pipeline: remove bg, composite onto scene, generate bio.
    Returns (new_profile_photo, adoption_bio).
    """
    if original_photo is None:
        gr.Warning("Please upload a photo of the animal first.")
        return None, ""

    if not animal_type:
        animal_type = "Dog"

    gr.Info("Removing background...")
    try:
        subject_rgba = remove_background(original_photo)
    except Exception as e:
        logger.error(f"Background removal failed: {e}")
        gr.Warning("Background removal failed. Using original image.")
        subject_rgba = original_photo.convert("RGBA")

    gr.Info("Compositing onto new background...")
    try:
        background = select_background(animal_type, custom_background)
        new_photo = composite_animal_on_background(
            subject_rgba, background, output_size=OUTPUT_SIZE
        )
    except Exception as e:
        logger.error(f"Compositing failed: {e}")
        gr.Warning("Image compositing failed.")
        new_photo = original_photo

    gr.Info("Writing adoption bio...")
    try:
        bio = generate_bio(name, animal_type, characteristics)
    except Exception as e:
        logger.error(f"Bio generation failed: {e}")
        bio = "Could not generate bio. Please try again."

    return new_photo, bio


CUSTOM_CSS = """
.main-title {
    text-align: center;
    margin-bottom: 0.5rem;
}
.main-title h1 {
    font-size: 2rem;
    font-weight: 700;
}
.subtitle {
    text-align: center;
    opacity: 0.8;
    margin-bottom: 1.5rem;
    font-size: 1rem;
}
.output-bio textarea {
    min-height: 250px !important;
    font-size: 1rem !important;
    line-height: 1.6 !important;
}
footer {
    display: none !important;
}
"""


THEME = gr.themes.Soft(
    primary_hue="violet",
    secondary_hue="violet",
    neutral_hue="slate",
    font=gr.themes.GoogleFont("Inter"),
)


def build_app() -> gr.Blocks:
    with gr.Blocks(title="Shelter Buddy: Adoption Booster") as app:

        gr.HTML(
            '<div class="main-title"><h1>🐾 Shelter Buddy: Adoption Booster</h1></div>'
            '<div class="subtitle">Upload a shelter photo to give it a warm makeover and generate a loving bio.</div>'
        )

        with gr.Row(equal_height=False):
            with gr.Column(scale=1):
                original_photo = gr.Image(
                    label="Original Photo",
                    type="pil",
                    sources=["upload", "clipboard"],
                    height=300,
                )

                name_input = gr.Textbox(
                    label="Name",
                    placeholder="e.g. Barnaby",
                    max_lines=1,
                )

                animal_type = gr.Radio(
                    choices=["Dog", "Cat", "Other"],
                    label="Animal Type",
                    value="Dog",
                )

                characteristics = gr.CheckboxGroup(
                    choices=CHARACTERISTICS,
                    label="Characteristics",
                )

                custom_bg = gr.Image(
                    label="Custom Background (Optional)",
                    type="pil",
                    sources=["upload"],
                    height=180,
                )

            with gr.Column(scale=1):
                new_photo = gr.Image(
                    label="New Profile Photo",
                    type="pil",
                    interactive=False,
                    height=400,
                )

                bio_output = gr.Textbox(
                    label="Adoption Bio",
                    lines=10,
                    interactive=False,
                    elem_classes=["output-bio"],
                )

        transform_btn = gr.Button(
            "✨ Transform & Write Bio",
            variant="primary",
            size="lg",
        )

        transform_btn.click(
            fn=process_animal,
            inputs=[original_photo, name_input, animal_type, characteristics, custom_bg],
            outputs=[new_photo, bio_output],
        )

    return app


if __name__ == "__main__":
    app = build_app()
    app.launch(server_name="0.0.0.0", server_port=7860, theme=THEME, css=CUSTOM_CSS)