Spaces:
Sleeping
Sleeping
| """ | |
| 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 | |
| # --------------------------------------------------------------------------- | |
| 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()) | |
| 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 | |
| # --------------------------------------------------------------------------- | |
| 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, | |
| } | |