""" 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": }`` 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, }