lhk229/BGfilter
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add image_url input mode (service fetches the image)

Browse Source 
/remove-background now accepts either an uploaded `file` or an `image_url`
the service fetches itself (exactly one required). The fetch runs BEFORE
the compute gate so a slow download never holds a pool slot, and is
bounded by its own semaphore (2 x cpu_workers, hard cap 64) plus a shared
httpx client (connection reuse, matching max_connections) -- this restores
the throttling a direct upload has implicitly, so a burst of tiny URL
requests can't exhaust FDs/ephemeral ports or hammer the upstream.

Bounds: 25 MB body cap (413), 5 s connect / 15 s read timeout, <=3
redirects; MAX_IMAGE_PIXELS still applies post-decode. Deliberately NO
SSRF filtering (no private-IP/host allowlist) -- the endpoint trusts
authenticated callers; auth runs before any fetch. Do not expose to
untrusted clients without adding private-address/redirect filtering to
_fetch_image.

Adds httpx. New env knobs BGFILTER_FETCH_MAX_BYTES / BGFILTER_FETCH_TIMEOUT.
Docs (README/README_ZH/DEPLOY/DEPLOY_ZH) updated. _fetch_image unit-tested
locally (happy path, size cap ->413, bad scheme ->400, 404 ->400);
on-box end-to-end still pending (test box went offline).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
孔令弘
2026-07-14 14:33:34 +00:00
parent 07ba02fdde
commit 6506bdd643
6 changed files with 99 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app.py
+81 -4
View File
@@ -16,7 +16,9 @@ import time
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager
from pathlib import Path
from urllib.parse import urlparse
import httpx
import numpy as np
from fastapi import Depends, FastAPI, File, Form, Header, HTTPException, Response, UploadFile
from PIL import Image, ImageOps, UnidentifiedImageError
@@ -46,6 +48,16 @@ def _int_env(name: str, default: int) -> int:
MAX_IMAGE_PIXELS = _int_env("BGFILTER_MAX_IMAGE_PIXELS", 4_194_304)
# /remove-background accepts either an uploaded `file` or an `image_url` the service
# fetches itself. These bound the fetch so a huge or slow URL cannot hang a worker or
# blow RAM. NOTE: no SSRF hardening here (no host/IP allowlist) -- the endpoint is
# guarded by token auth and assumes trusted callers; do NOT expose it to untrusted
# clients without adding private-address / redirect filtering to _fetch_image.
FETCH_MAX_BYTES = _int_env("BGFILTER_FETCH_MAX_BYTES", 25 * 1024 * 1024)
FETCH_TIMEOUT_S = float(os.environ.get("BGFILTER_FETCH_TIMEOUT", "15"))
# Outbound fetch concurrency = 2 x cpu_workers, hard-capped here so a burst can never
# exhaust file descriptors / ephemeral ports (each fetch = one socket).
FETCH_CONCURRENCY_CAP = 64
# We enforce our own pixel budget from the image header before decoding, so disable
# Pillow's separate decompression-bomb ceiling (which would otherwise raise its own error).
Image.MAX_IMAGE_PIXELS = None
@@ -81,6 +93,11 @@ _manager: PipelineManager | None = None
# drains.
_executor: ThreadPoolExecutor | None = None
_sem: asyncio.Semaphore | None = None
# Outbound image_url fetches: a FIFO admission bound (restores the throttle a direct
# upload has implicitly) and a shared client for connection reuse. Both built in
# lifespan once the pool size is known. See _fetch_image / §image_url mode.
_fetch_sem: asyncio.Semaphore | None = None
_http_client: httpx.AsyncClient | None = None
def _build_settings():
@@ -106,13 +123,24 @@ def _build_settings():
@asynccontextmanager
async def lifespan(app: FastAPI):
global _manager, _executor, _sem
global _manager, _executor, _sem, _fetch_sem, _http_client
settings = _build_settings()
preload = os.environ.get("BGFILTER_PRELOAD", "1").strip().lower() not in ("0", "false", "no", "")
# Pool size: config default (server.cpu_workers), env override wins.
cpu_workers = max(1, _int_env("BGFILTER_CPU_WORKERS", settings.server.cpu_workers))
_executor = ThreadPoolExecutor(max_workers=cpu_workers, thread_name_prefix="bgf")
_sem = asyncio.Semaphore(cpu_workers)
# Outbound-fetch bound: 2x the pool (a small prefetch buffer), hard-capped. The
# shared client's max_connections matches so httpx never opens more sockets than
# the semaphore admits, and reuses connections across fetches.
fetch_conc = min(2 * cpu_workers, FETCH_CONCURRENCY_CAP)
_fetch_sem = asyncio.Semaphore(fetch_conc)
_http_client = httpx.AsyncClient(
follow_redirects=True,
max_redirects=3,
timeout=httpx.Timeout(connect=5.0, read=FETCH_TIMEOUT_S, write=10.0, pool=None),
limits=httpx.Limits(max_connections=fetch_conc),
)
logger.info(
"starting bgfilter service: device=%s default_seg=%s cross_check=%s preload=%s "
"cpu_workers=%d gpu_concurrency=%d max_pixels=%d auth=%s",
@@ -130,7 +158,10 @@ async def lifespan(app: FastAPI):
"BGFILTER_AUTH_TOKEN not set: /remove-background is OPEN (no app-level auth)"
)
_manager = PipelineManager(settings, preload=preload)
yield
try:
yield
finally:
await _http_client.aclose()
app = FastAPI(title="BGfilter", version=VERSION, lifespan=lifespan)
@@ -205,6 +236,34 @@ def _normalize_background_mode(value: str | None) -> str | None:
raise HTTPException(status_code=400, detail="invalid background_mode")
async def _fetch_image(url: str) -> bytes:
"""Download image bytes from a URL, bounded by size and time.
Robustness only -- no SSRF filtering (see FETCH_* note above): callers are
trusted via the auth token. Streams the body so an over-large response is
rejected without being fully buffered. Uses the shared client (connection reuse);
the caller holds _fetch_sem so outbound concurrency stays bounded.
"""
parsed = urlparse(url)
if parsed.scheme not in ("http", "https") or not parsed.hostname:
raise HTTPException(status_code=400, detail="image_url must be an http(s) URL")
try:
async with _http_client.stream("GET", url) as resp:
if resp.status_code != 200:
raise HTTPException(
status_code=400,
detail=f"image_url fetch failed (status {resp.status_code})",
)
buf = bytearray()
async for chunk in resp.aiter_bytes():
buf += chunk
if len(buf) > FETCH_MAX_BYTES:
raise HTTPException(status_code=413, detail="image_url body too large")
return bytes(buf)
except httpx.HTTPError:
raise HTTPException(status_code=400, detail="image_url fetch failed")
def _decode_image(data: bytes) -> np.ndarray:
if not data:
raise HTTPException(status_code=400, detail="empty image")
@@ -235,7 +294,8 @@ def _encode_png(rgb: np.ndarray, alpha: np.ndarray) -> bytes:
@app.post("/remove-background")
async def remove_background(
file: UploadFile = File(...),
file: UploadFile | None = File(None),
image_url: str | None = Form(None),
screen_color: str | None = Form(None),
seg_model: str | None = Form(None),
cross_check: str | None = Form(None),
@@ -245,6 +305,14 @@ async def remove_background(
if _manager is None:
raise HTTPException(status_code=503, detail="service not ready")
# Exactly one image source. (Auth already ran, so only authenticated callers can
# trigger an outbound fetch.)
url = image_url.strip() if image_url else None
if (file is None) == (not url):
raise HTTPException(
status_code=400, detail="provide exactly one of 'file' or 'image_url'"
)
normalized_color = _normalize_screen_color(screen_color)
normalized_seg = _normalize_seg_model(seg_model) or _manager.default_seg_model
normalized_cc = _normalize_cross_check(cross_check)
@@ -260,11 +328,20 @@ async def remove_background(
# Honors the lane lock (a locked cpu-fast box reports/runs "off" even on cross_check=on).
effective_cc = _manager.effective_cross_check(normalized_cc)
# Acquire the raw bytes BEFORE the compute gate: a slow fetch must not hold a
# pool/compute slot. The upload body is already buffered by Starlette (fast);
# a URL is fetched under its own looser bound (_fetch_sem).
if url is not None:
async with _fetch_sem:
raw = await _fetch_image(url)
else:
raw = await file.read()
# Bound in-flight requests to the pool size, then run on the pool. The GPU
# sections are serialized inside the pipeline (a shared GPU mutex); the CPU
# sections of concurrent requests overlap across the pool threads.
async with _sem:
rgb = _decode_image(await file.read())
rgb = _decode_image(raw)
start = time.perf_counter()
loop = asyncio.get_running_loop()
try: