backend/app/api/bgremove.py

"""
bgremove.py — POST /bgremove/upload  and  GET /bgremove/result/{job_id}

INDEPENDENT of the LaMa/watermark-removal pipeline, Real-ESRGAN upscale pipeline,
and object-removal pipeline.
Do not import from jobs_store.py, upscale_jobs_store.py, schemas.py, generate.py,
inpaint_engine.py, upscale_engine.py, or any other pipeline module.

Design
------
Background removal, like upscaling, has no user input after the image is
chosen — so upload and inference kick-off are combined into a single POST
/bgremove/upload endpoint.

Flow
----
1. POST /bgremove/upload (multipart file)
   a. Validate content-type.
   b. EXIF-correct and convert to RGB via Pillow.
   c. Assign UUID job_id; create storage/bgremove/{folder_name}/.
   d. Save EXIF-corrected image as original.jpg.
   e. Persist BgRemoveJob(status="pending") to bgremove_jobs.json.
   f. Immediately advance status to "processing" and enqueue BackgroundTask.
   g. Return {"job_id": job_id} immediately (inference runs in background).

2. GET /bgremove/result/{job_id}  (polling)
   Returns {"job_id", "status", "result_image_url", "error_message"}.
   result_image_url is a /storage/... URL served by the existing StaticFiles
   mount in main.py — no extra route needed.

3. run_bgremove_job (background worker)
   Runs after the POST response is sent; has no access to the Request object.
   Receives the engine as an explicit argument.
   Updates job status to "completed" or "failed"; never raises.

Storage layout (kept entirely separate from all other storage paths)
---------------------------------------------------------------------
  storage/
    bgremove/
      {folder_name}/
        original.jpg   — EXIF-corrected upload (JPEG)
        result.png     — RGBA background-removed output (PNG, for alpha channel)

These are served as:
  /storage/bgremove/{folder_name}/original.jpg
  /storage/bgremove/{folder_name}/result.png
via the existing app.mount("/storage", StaticFiles(...)) in main.py.

Why PNG for result?
-------------------
JPEG does not support an alpha channel.  The background removal output is RGBA
(foreground pixels are opaque/semi-transparent, background pixels are
transparent).  PNG losslessly encodes all four channels, so result.png
preserves the transparency that makes background removal useful.
"""

from __future__ import annotations

import io
import traceback
import uuid
from datetime import datetime, timezone
from pathlib import Path

from fastapi import APIRouter, BackgroundTasks, File, HTTPException, Request, UploadFile
from PIL import Image, ImageOps

from app.config import STORAGE_DIR
from app.utils import job_folder_name
from app.bgremove_jobs_store import BgRemoveJob, create_job, get_job, update_job
from app.services.bgremove_engine import BaseBackgroundRemovalEngine

router = APIRouter()

# Background-removal jobs live in a dedicated sub-directory of storage/ so they
# can never collide with watermark-removal, upscale, or object-removal folders.
_BGREMOVE_STORAGE = STORAGE_DIR / "bgremove"


# ---------------------------------------------------------------------------
# Background worker
# ---------------------------------------------------------------------------

def run_bgremove_job(
    job_id:        str,
    folder_name:   str,
    original_path: Path,
    engine:        BaseBackgroundRemovalEngine,
) -> None:
    """Run BiRefNet background removal and save result.png; update job status.

    Called by FastAPI BackgroundTasks after the /bgremove/upload response is sent.
    Has no access to the Request object — engine is passed as an explicit arg.
    Updates status to "completed" or "failed"; never raises.

    Pipeline
    --------
    original.jpg (RGB) → engine.remove_background() → RGBA PIL → result.png (PNG)

    The result is saved as PNG (not JPEG) because JPEG cannot encode the alpha
    channel that carries the transparency mask.
    """
    job_dir = _BGREMOVE_STORAGE / folder_name
    result_path = job_dir / "result.png"

    try:
        pil_img = Image.open(original_path).convert("RGB")

        # Run background removal — returns RGBA PIL Image
        rgba_img = engine.remove_background(pil_img)

        # Save as PNG to preserve the alpha channel
        rgba_img.save(str(result_path), format="PNG")

        # Build the URL path relative to the /storage mount.
        rel = result_path.relative_to(STORAGE_DIR)
        update_job(job_id, status="completed", result_image=f"/{rel.as_posix()}")

    except Exception:
        update_job(
            job_id,
            status="failed",
            error_message=traceback.format_exc(),
        )


# ---------------------------------------------------------------------------
# POST /bgremove/upload
# ---------------------------------------------------------------------------

@router.post("/upload")
async def bgremove_upload(
    file:             UploadFile,
    background_tasks: BackgroundTasks,
    request:          Request,
) -> dict:
    """Accept an image, EXIF-correct it, persist it, and kick off background removal.

    Form fields
    -----------
    file : image file (multipart)

    Returns ``{"job_id": <str>}`` immediately; background removal runs in the
    background.  Poll GET /bgremove/result/{job_id} for completion.
    """
    # --- Content-type guard --------------------------------------------------
    if not (file.content_type or "").startswith("image/"):
        raise HTTPException(
            status_code=415,
            detail=f"Expected an image file, got content-type '{file.content_type}'.",
        )

    # --- Read and EXIF-correct -----------------------------------------------
    raw_bytes = await file.read()
    if not raw_bytes:
        raise HTTPException(status_code=400, detail="Uploaded file is empty.")

    try:
        pil_img = Image.open(io.BytesIO(raw_bytes))
        pil_img = ImageOps.exif_transpose(pil_img)
        pil_img = pil_img.convert("RGB")
    except Exception as exc:
        raise HTTPException(status_code=422, detail=f"Could not decode image: {exc}")

    # --- Persist original image ----------------------------------------------
    job_id     = str(uuid.uuid4())
    created_at = datetime.now(timezone.utc)
    folder     = job_folder_name(job_id, created_at)
    job_dir    = _BGREMOVE_STORAGE / folder
    job_dir.mkdir(parents=True, exist_ok=True)

    original_path = job_dir / "original.jpg"
    pil_img.save(str(original_path), format="JPEG", quality=95)

    # --- Create job record ---------------------------------------------------
    orig_rel = original_path.relative_to(STORAGE_DIR)
    job = BgRemoveJob(
        id             = job_id,
        original_image = f"/{orig_rel.as_posix()}",
        status         = "processing",
        created_at     = created_at.isoformat(),
    )
    create_job(job)

    # --- Enqueue background removal ------------------------------------------
    engine: BaseBackgroundRemovalEngine = request.app.state.bgremove_engine

    background_tasks.add_task(
        run_bgremove_job,
        job_id        = job_id,
        folder_name   = folder,
        original_path = original_path,
        engine        = engine,
    )

    return {"job_id": job_id}


# ---------------------------------------------------------------------------
# GET /bgremove/result/{job_id}
# ---------------------------------------------------------------------------

@router.get("/result/{job_id}")
def bgremove_result(job_id: str) -> dict:
    """Return the current status of a background removal job.

    Response fields:
        job_id           : str
        status           : "pending" | "processing" | "completed" | "failed"
        result_image_url : str | null  — /storage/bgremove/{folder}/result.png
        error_message    : str | null
    """
    job = get_job(job_id)
    if job is None:
        raise HTTPException(
            status_code=404,
            detail=f"Background removal job '{job_id}' not found.",
        )

    return {
        "job_id":           job.id,
        "status":           job.status,
        "result_image_url": job.result_image,
        "error_message":    job.error_message,
    }