API: implement inpainting endpoints (upload image/mask, inpaint, download, logs); optional Bearer via API_TOKEN

README.md

@@ -13,14 +13,15 @@ hardware: cpu-basic
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
  
-## API (FastAPI)
-- Base URL: `https://logicgoinfotechspaces-hair-stable-new.hf.space`
-- Auth: Bearer token `logicgo@123`
+## Inpainting API (FastAPI)
+- Auth: optional Bearer via env `API_TOKEN`. If set, send `Authorization: Bearer <token>`.
 
 Endpoints:
 - `GET /health` → {"status":"healthy"}
-- `POST /upload` (form-data: image=file) → {"id":"<id>","filename":"name.png"}
-- `POST /get-hairswap` (JSON body with source_id, reference_id, …) → {"result":"output_xxx.png"}
+- `POST /upload-image` (form-data: image=file) → {"id":"<image_id>","filename":"name.png"}
+- `POST /upload-mask` (form-data: mask=file) → {"id":"<mask_id>","filename":"mask.png"}
+- `POST /inpaint` (JSON: {image_id, mask_id}) → {"result":"output_xxx.png"}
+- `POST /inpaint-multipart` (form-data: image=file, mask=file) → {"result":"output_xxx.png"}
 - `GET /download/{filename}` → image file
 - `GET /logs` → recent uploads/results
 

api/main.py

@@ -4,19 +4,26 @@ import shutil
 from datetime import datetime
 from typing import Dict, List, Optional
 
+import numpy as np
 from fastapi import FastAPI, UploadFile, File, HTTPException, Depends, Header
 from fastapi.responses import FileResponse, JSONResponse
 from pydantic import BaseModel
+from PIL import Image
 
-# Simple settings
-UPLOAD_DIR = os.path.join(os.getcwd(), "uploads")
-OUTPUT_DIR = os.path.join(os.getcwd(), "outputs")
-AUTH_TOKEN = "logicgo@123"  # Bearer token
+from src.core import process_inpaint
+
+# Directories
+ROOT = os.getcwd()
+UPLOAD_DIR = os.path.join(ROOT, "uploads")
+OUTPUT_DIR = os.path.join(ROOT, "outputs")
 
 os.makedirs(UPLOAD_DIR, exist_ok=True)
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-app = FastAPI(title="Hair Swap API", version="1.0.0")
+# Optional Bearer token: set env API_TOKEN to require auth; if not set, endpoints are open
+ENV_TOKEN = os.environ.get("API_TOKEN")
+
+app = FastAPI(title="Photo Object Removal API", version="1.0.0")
 
 # In-memory stores
 file_store: Dict[str, Dict[str, str]] = {}
@@ -24,20 +31,18 @@ logs: List[Dict[str, str]] = []
 
 
 def bearer_auth(authorization: Optional[str] = Header(default=None)) -> None:
+    if not ENV_TOKEN:
+        return
     if authorization is None or not authorization.lower().startswith("bearer "):
         raise HTTPException(status_code=401, detail="Unauthorized")
     token = authorization.split(" ", 1)[1]
-    if token != AUTH_TOKEN:
+    if token != ENV_TOKEN:
         raise HTTPException(status_code=403, detail="Forbidden")
 
 
-class HairSwapRequest(BaseModel):
-    source_id: str
-    reference_id: str
-    converter_scale: Optional[float] = 1.0
-    scale: Optional[float] = 1.0
-    guidance_scale: Optional[float] = 1.5
-    controlnet_conditioning_scale: Optional[float] = 1.0
+class InpaintRequest(BaseModel):
+    image_id: str
+    mask_id: str
 
 
 @app.get("/health")
@@ -45,38 +50,97 @@ def health() -> Dict[str, str]:
     return {"status": "healthy"}
 
 
-@app.post("/upload")
+@app.post("/upload-image")
 def upload_image(image: UploadFile = File(...), _: None = Depends(bearer_auth)) -> Dict[str, str]:
     ext = os.path.splitext(image.filename)[1] or ".png"
     file_id = str(uuid.uuid4())
     stored_name = f"{file_id}{ext}"
     stored_path = os.path.join(UPLOAD_DIR, stored_name)
-
     with open(stored_path, "wb") as f:
         shutil.copyfileobj(image.file, f)
-
     file_store[file_id] = {
+        "type": "image",
         "filename": image.filename,
         "stored_name": stored_name,
         "path": stored_path,
         "timestamp": datetime.utcnow().isoformat(),
     }
-    logs.append({"id": file_id, "filename": image.filename, "timestamp": datetime.utcnow().isoformat()})
+    logs.append({"id": file_id, "filename": image.filename, "type": "image", "timestamp": datetime.utcnow().isoformat()})
     return {"id": file_id, "filename": image.filename}
 
 
-@app.post("/get-hairswap")
-def get_hairswap(req: HairSwapRequest, _: None = Depends(bearer_auth)) -> Dict[str, str]:
-    if req.source_id not in file_store:
-        raise HTTPException(status_code=404, detail="source_id not found")
-    if req.reference_id not in file_store:
-        raise HTTPException(status_code=404, detail="reference_id not found")
+@app.post("/upload-mask")
+def upload_mask(mask: UploadFile = File(...), _: None = Depends(bearer_auth)) -> Dict[str, str]:
+    ext = os.path.splitext(mask.filename)[1] or ".png"
+    file_id = str(uuid.uuid4())
+    stored_name = f"{file_id}{ext}"
+    stored_path = os.path.join(UPLOAD_DIR, stored_name)
+    with open(stored_path, "wb") as f:
+        shutil.copyfileobj(mask.file, f)
+    file_store[file_id] = {
+        "type": "mask",
+        "filename": mask.filename,
+        "stored_name": stored_name,
+        "path": stored_path,
+        "timestamp": datetime.utcnow().isoformat(),
+    }
+    logs.append({"id": file_id, "filename": mask.filename, "type": "mask", "timestamp": datetime.utcnow().isoformat()})
+    return {"id": file_id, "filename": mask.filename}
+
+
+def _load_rgba_image(path: str) -> Image.Image:
+    img = Image.open(path)
+    return img.convert("RGBA")
+
+
+def _load_rgba_mask_from_image(img: Image.Image) -> np.ndarray:
+    # Expected by process_inpaint: RGBA where alpha=0 for drawn (to remove), 255 elsewhere
+    if img.mode != "RGBA":
+        # If no alpha, treat non-black/white>0 as masked areas
+        gray = img.convert("L")
+        arr = np.array(gray)
+        alpha = np.where(arr > 0, 0, 255).astype(np.uint8)
+        rgba = np.zeros((img.height, img.width, 4), dtype=np.uint8)
+        rgba[:, :, 3] = alpha
+        return rgba
+    return np.array(img)
+
+
+@app.post("/inpaint")
+def inpaint(req: InpaintRequest, _: None = Depends(bearer_auth)) -> Dict[str, str]:
+    if req.image_id not in file_store or file_store[req.image_id]["type"] != "image":
+        raise HTTPException(status_code=404, detail="image_id not found")
+    if req.mask_id not in file_store or file_store[req.mask_id]["type"] != "mask":
+        raise HTTPException(status_code=404, detail="mask_id not found")
+
+    img_rgba = _load_rgba_image(file_store[req.image_id]["path"])
+    mask_img = Image.open(file_store[req.mask_id]["path"])  # may be RGB/gray/RGBA
+    mask_rgba = _load_rgba_mask_from_image(mask_img)
+
+    result = process_inpaint(np.array(img_rgba), mask_rgba)
+    result_name = f"output_{uuid.uuid4().hex}.png"
+    result_path = os.path.join(OUTPUT_DIR, result_name)
+    Image.fromarray(result).save(result_path)
+
+    logs.append({"result": result_name, "timestamp": datetime.utcnow().isoformat()})
+    return {"result": result_name}
+
+
+@app.post("/inpaint-multipart")
+def inpaint_multipart(
+    image: UploadFile = File(...),
+    mask: UploadFile = File(...),
+    _: None = Depends(bearer_auth),
+) -> Dict[str, str]:
+    # Load in-memory
+    img = Image.open(image.file).convert("RGBA")
+    m = Image.open(mask.file)
+    mask_rgba = _load_rgba_mask_from_image(m)
 
-    # Placeholder: In lieu of a real hair-swap model, return the source image as the result
-    src_path = file_store[req.source_id]["path"]
+    result = process_inpaint(np.array(img), mask_rgba)
     result_name = f"output_{uuid.uuid4().hex}.png"
     result_path = os.path.join(OUTPUT_DIR, result_name)
-    shutil.copyfile(src_path, result_path)
+    Image.fromarray(result).save(result_path)
 
     logs.append({"result": result_name, "timestamp": datetime.utcnow().isoformat()})
     return {"result": result_name}
@@ -84,14 +148,9 @@ def get_hairswap(req: HairSwapRequest, _: None = Depends(bearer_auth)) -> Dict[s
 
 @app.get("/download/{filename}")
 def download_file(filename: str, _: None = Depends(bearer_auth)):
-    # Serve from outputs first, then uploads as a fallback
     path = os.path.join(OUTPUT_DIR, filename)
     if not os.path.isfile(path):
-        alt = os.path.join(UPLOAD_DIR, filename)
-        if os.path.isfile(alt):
-            path = alt
-        else:
-            raise HTTPException(status_code=404, detail="file not found")
+        raise HTTPException(status_code=404, detail="file not found")
     return FileResponse(path)