Blame for backend/app/api/object_removal.py

"""
object_removal.py — POST /object-removal/upload,
                    POST /object-removal/generate,
                    GET  /object-removal/result/{job_id}

Object-removal pipeline: identical to watermark removal except:
  - Storage:  storage/object-removal/{job_id}/steps/{n}/
  - job_type: "object"
  - Reuses app.state.inpaint_engine (same LaMa model)
"""

from __future__ import annotations

import io
import shutil
import traceback
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Literal

import cv2
import numpy as np
from fastapi import APIRouter, BackgroundTasks, File, HTTPException, Request, UploadFile
from PIL import Image, ImageOps
from pydantic import BaseModel

from app.config import STORAGE_DIR
from app.utils import job_folder_name
from app.jobs_store import (
    add_step,
    create_job,
    get_job,
    truncate_steps_after,
    update_job,
    update_step,
)
from app.models.schemas import Job, StepRecord
from app.services.blend_service import blend, check_output_quality
from app.services.inpaint_engine import BaseInpaintEngine
from app.services.mask_service import build_masks, save_mask

router = APIRouter()

_SUBDIR = "object-removal"


def _to_url(abs_path: str | None) -> str | None:
    """Convert an absolute path under STORAGE_DIR to a /storage/… URL."""
    if abs_path is None:
        return None
    try:
        rel = Path(abs_path).relative_to(STORAGE_DIR)
        return "/storage/" + rel.as_posix()
    except ValueError:
        return None


# ---------------------------------------------------------------------------
# Upload
# ---------------------------------------------------------------------------

@router.post("/upload")
async def obj_upload(file: UploadFile = File(...)) -> dict:
    """Accept an image upload for object removal, return job_id + dimensions."""
    if not file.content_type.startswith("image/"):
        raise HTTPException(
            status_code=415,
            detail=f"Expected an image file, got content-type '{file.content_type}'.",
        )

    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}")

    img_bgr: np.ndarray = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
    height, width = img_bgr.shape[:2]

    job_id     = str(uuid.uuid4())
    created_at = datetime.now(timezone.utc)
    folder     = job_folder_name(job_id, created_at)
    step0_dir  = STORAGE_DIR / _SUBDIR / folder / "steps" / "0"
    step0_dir.mkdir(parents=True, exist_ok=True)

    original_path = step0_dir / "original.jpg"
    ok = cv2.imwrite(str(original_path), img_bgr)
    if not ok:
        raise HTTPException(status_code=500, detail="Failed to write uploaded image to storage.")

    job = Job(id=job_id, current_step=0, job_type="object", created_at=created_at)
    create_job(job)

    step0 = StepRecord(
        step_number=0,
        original_image=str(original_path),
        result_image=str(original_path),
        status="completed",
    )
    add_step(job_id, step0)

    return {"job_id": job_id, "width": width, "height": height}


# ---------------------------------------------------------------------------
# Generate (background inpainting)
# ---------------------------------------------------------------------------

class ObjGenerateRequest(BaseModel):
    job_id:      str
    mask_source: Literal["manual", "ocr", "mixed"] = "manual"
    mask_json:   dict


def run_obj_inpaint_job(
    job_id:      str,
    folder_name: str,
    step_number: int,
    image:       np.ndarray,
    hard_mask:   np.ndarray,
    soft_mask:   np.ndarray,
    engine:      BaseInpaintEngine,
) -> None:
    """Run LaMa inpainting for an object-removal job step."""
    step_dir = STORAGE_DIR / _SUBDIR / folder_name / "steps" / str(step_number)
    try:
        lama_output: np.ndarray  = engine.inpaint(image, hard_mask)
        result_image: np.ndarray = blend(image, lama_output, soft_mask)

        passed, error_msg = check_output_quality(result_image, expected_shape=image.shape)
        if not passed:
            update_step(job_id, step_number, status="failed", error_message=error_msg)
            return

        result_path = step_dir / "result.jpg"
        ok = cv2.imwrite(str(result_path), result_image)
        if not ok:
            update_step(
                job_id, step_number,
                status="failed",
                error_message="cv2.imwrite failed writing result.jpg.",
            )
            return

        update_step(job_id, step_number, status="completed", result_image=str(result_path))

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


@router.post("/generate")
async def obj_generate(
    body: ObjGenerateRequest,
    background_tasks: BackgroundTasks,
    request: Request,
) -> dict:
    """Create new step, enqueue background inpainting for object removal."""
    job_id = body.job_id
    job = get_job(job_id)
    if job is None:
        raise HTTPException(status_code=404, detail=f"Job '{job_id}' not found.")

    folder_name = job_folder_name(job.id, job.created_at)

    current_step_num = job.current_step
    current_step = next(
        (s for s in job.steps if s.step_number == current_step_num), None
    )
    if current_step is None or current_step.result_image is None:
        raise HTTPException(
            status_code=400,
            detail=f"Step {current_step_num} has no result image to build from.",
        )

    base_path = Path(current_step.result_image)
    if not base_path.exists():
        raise HTTPException(
            status_code=500,
            detail=f"Base image missing from storage: {base_path}",
        )

    image: np.ndarray = cv2.imread(str(base_path))
    if image is None:
        raise HTTPException(
            status_code=500,
            detail=f"cv2.imread could not decode: {base_path}",
        )

    try:
        hard_mask, soft_mask = build_masks(body.mask_json)
    except ValueError as exc:
        raise HTTPException(status_code=422, detail=str(exc))

    new_step_num = current_step_num + 1
    if any(s.step_number >= new_step_num for s in job.steps):
        removed = truncate_steps_after(job_id, keep_through=current_step_num)
        for s in removed:
            orphan = STORAGE_DIR / _SUBDIR / folder_name / "steps" / str(s["step_number"])
            if orphan.exists():
                shutil.rmtree(orphan)

    step_dir = STORAGE_DIR / _SUBDIR / folder_name / "steps" / str(new_step_num)
    step_dir.mkdir(parents=True, exist_ok=True)

    step_original = step_dir / "original.jpg"
    shutil.copy2(str(base_path), str(step_original))

    hard_mask_path = step_dir / "mask_hard.png"
    soft_mask_path = step_dir / "mask_soft.png"
    save_mask(hard_mask, str(hard_mask_path))
    save_mask(soft_mask, str(soft_mask_path))

    new_step = StepRecord(
        step_number=new_step_num,
        original_image=str(step_original),
        mask_image=str(hard_mask_path),
        mask_source=body.mask_source,
        status="processing",
    )
    add_step(job_id, new_step)
    update_job(job_id, current_step=new_step_num)

    engine: BaseInpaintEngine = request.app.state.inpaint_engine
    background_tasks.add_task(
        run_obj_inpaint_job,
        job_id      = job_id,
        folder_name = folder_name,
        step_number = new_step_num,
        image       = image,
        hard_mask   = hard_mask,
        soft_mask   = soft_mask,
        engine      = engine,
    )

    return {"job_id": job_id, "step_number": new_step_num, "status": "processing"}


# ---------------------------------------------------------------------------
# Result
# ---------------------------------------------------------------------------

@router.get("/result/{job_id}")
def obj_get_result(job_id: str) -> dict:
    """Return current object-removal job state including full step history."""
    job = get_job(job_id)
    if job is None:
        raise HTTPException(status_code=404, detail=f"Job '{job_id}' not found.")

    current = next((s for s in job.steps if s.step_number == job.current_step), None)
    step0   = next((s for s in job.steps if s.step_number == 0), None)

    steps_summary = [
        {
            "step_number":        s.step_number,
            "status":             s.status,
            "original_image_url": _to_url(s.original_image),
            "result_image_url":   _to_url(s.result_image),
            "error_message":      s.error_message,
        }
        for s in sorted(job.steps, key=lambda s: s.step_number)
    ]

    return {
        "job_id":             job.id,
        "current_step":       job.current_step,
        "status":             current.status        if current else "unknown",
        "error_message":      current.error_message if current else None,
        "original_image_url": _to_url(step0.original_image) if step0 else None,
        "result_image_url":   _to_url(current.result_image) if current else None,
        "steps":              steps_summary,
    }