Spaces:

hackerbhai
/

RemBG_API

Paused

File size: 21,362 Bytes

import hmac
import io
import os
import time
from threading import Lock

from dotenv import load_dotenv
from flask import Flask, jsonify, request, send_file
from PIL import Image, ImageColor, ImageOps, UnidentifiedImageError
from rembg import new_session, remove

load_dotenv()

app = Flask(__name__)
app.config["MAX_CONTENT_LENGTH"] = 12 * 1024 * 1024  # 12 MB upload limit

MODELS = {
    "isnet-general-use": "Quality · ISNet General",
    "u2net": "Balanced · U2Net",
    "u2netp": "Fast · U2NetP",
    "u2net_human_seg": "People · Human Segmentation",
    "isnet-anime": "Anime · ISNet Anime",
}
FORMATS = {"png", "jpg", "webp"}
SESSIONS = {}
SESSION_LOCK = Lock()
RATE_LOCK = Lock()
RATE_BUCKETS = {}
RATE_LIMIT = 8
RATE_WINDOW = 60

INDEX_HTML = r"""
<!doctype html>
<html lang="en">
<head>
  <meta charset="utf-8">
  <meta name="viewport" content="width=device-width, initial-scale=1">
  <title>Background Studio</title>
  <style>
    :root { color-scheme: dark; --bg:#090d18; --panel:#121a2e; --panel2:#18233d; --line:rgba(148,163,184,.2); --text:#f7f9ff; --muted:#95a5c2; --cyan:#43ddc6; --blue:#8298ff; --pink:#ff8eaa; --shadow:0 26px 80px rgba(0,0,0,.32); }
    * { box-sizing:border-box; }
    body { min-height:100vh; margin:0; color:var(--text); font:15px/1.55 Inter,ui-sans-serif,system-ui,-apple-system,"Segoe UI",sans-serif; background:radial-gradient(circle at 12% 0%,rgba(130,152,255,.24),transparent 34rem),radial-gradient(circle at 94% 17%,rgba(67,221,198,.13),transparent 30rem),var(--bg); }
    button,input,select { font:inherit; }
    button { cursor:pointer; }
    .shell { width:min(1180px,calc(100% - 34px)); margin:0 auto; padding:30px 0 58px; }
    .topbar { display:flex; align-items:center; justify-content:space-between; gap:20px; }
    .brand { display:flex; align-items:center; gap:11px; font-weight:850; letter-spacing:.05em; }
    .brand-mark { display:grid; place-items:center; width:39px; height:39px; border-radius:13px; color:#06111b; background:linear-gradient(135deg,var(--cyan),var(--blue)); box-shadow:0 9px 28px rgba(67,221,198,.2); }
    .brand small { display:block; color:var(--muted); font-weight:500; letter-spacing:0; }
    .system { display:inline-flex; align-items:center; gap:8px; color:#b7c6df; font-size:13px; }
    .dot { width:8px; height:8px; border-radius:50%; background:var(--cyan); box-shadow:0 0 13px var(--cyan); }
    .hero { max-width:770px; margin:72px 0 34px; }
    .eyebrow { margin:0 0 13px; color:var(--cyan); font-size:12px; font-weight:850; letter-spacing:.15em; text-transform:uppercase; }
    h1 { margin:0; font-size:clamp(2.7rem,7vw,5.5rem); line-height:.93; letter-spacing:-.075em; }
    .hero p { max-width:650px; margin:22px 0 0; color:var(--muted); font-size:18px; }
    .studio { display:grid; grid-template-columns:380px minmax(0,1fr); gap:18px; align-items:start; }
    .card { border:1px solid var(--line); border-radius:23px; background:rgba(18,26,46,.82); box-shadow:var(--shadow); backdrop-filter:blur(16px); }
    .controls { padding:22px; position:sticky; top:18px; }
    .card-head { display:flex; align-items:center; justify-content:space-between; gap:12px; margin-bottom:20px; font-weight:800; }
    .pill { padding:5px 9px; border:1px solid rgba(67,221,198,.25); border-radius:999px; color:var(--cyan); background:rgba(67,221,198,.08); font-size:11px; font-weight:800; }
    .dropzone { display:grid; place-items:center; min-height:166px; padding:22px; border:1px dashed rgba(130,152,255,.55); border-radius:17px; color:var(--muted); background:rgba(8,13,28,.45); text-align:center; transition:.18s; }
    .dropzone.drag { border-color:var(--cyan); background:rgba(67,221,198,.08); transform:scale(1.01); }
    .upload-icon { display:grid; place-items:center; width:44px; height:44px; margin-bottom:10px; border-radius:14px; color:var(--cyan); background:rgba(67,221,198,.1); font-size:21px; }
    .dropzone strong { display:block; color:var(--text); }
    .dropzone span { display:block; margin-top:4px; font-size:12px; }
    .file-name { overflow:hidden; max-width:100%; margin-top:10px; color:#b8c8ff; text-overflow:ellipsis; white-space:nowrap; }
    .browse { display:inline-block; margin-top:13px; padding:8px 12px; border:1px solid var(--line); border-radius:9px; color:var(--text); background:rgba(130,152,255,.12); cursor:pointer; font-size:12px; font-weight:750; }
    #file { display:none; }
    label { display:block; margin:18px 0 8px; color:#c9d5ec; font-size:11px; font-weight:800; letter-spacing:.09em; text-transform:uppercase; }
    select,input[type=password],input[type=color] { width:100%; padding:11px 12px; border:1px solid var(--line); border-radius:11px; outline:0; color:var(--text); background:rgba(7,12,25,.62); }
    select:focus,input:focus { border-color:var(--blue); box-shadow:0 0 0 4px rgba(130,152,255,.13); }
    input[type=color] { height:42px; padding:5px; }
    .row { display:grid; grid-template-columns:1fr 1fr; gap:10px; }
    .check { display:flex; align-items:center; gap:9px; margin-top:17px; color:#c5d2e9; font-size:13px; }
    .check input { width:16px; height:16px; accent-color:var(--cyan); }
    details { margin-top:14px; color:var(--muted); font-size:12px; }
    summary { cursor:pointer; color:#b8c8ff; font-weight:750; }
    .advanced { display:grid; grid-template-columns:1fr 1fr 1fr; gap:9px; margin-top:10px; }
    .advanced input { padding:8px; font-size:12px; }
    .advanced label { margin:0 0 5px; font-size:10px; }
    .run { width:100%; margin-top:22px; padding:14px 16px; border:0; border-radius:12px; color:#06111b; background:linear-gradient(100deg,var(--cyan),#91a7ff); box-shadow:0 12px 30px rgba(67,221,198,.15); font-weight:850; }
    .run:disabled { opacity:.55; cursor:wait; }
    .status { min-height:21px; margin:13px 0 0; color:var(--muted); font-size:12px; }
    .status.success { color:var(--cyan); }
    .status.error { color:var(--pink); }
    .result { min-width:0; padding:22px; }
    .result-head { display:flex; align-items:center; justify-content:space-between; gap:14px; margin-bottom:18px; }
    .result-head h2 { margin:0; font-size:19px; }
    .result-head span { color:var(--muted); font-size:12px; }
    .empty { display:grid; place-items:center; min-height:500px; border:1px dashed var(--line); border-radius:17px; color:#687995; text-align:center; }
    .empty-icon { margin-bottom:9px; font-size:32px; opacity:.7; }
    .compare-wrap { display:none; }
    .compare-wrap.show { display:block; }
    .compare { position:relative; overflow:hidden; min-height:430px; border:1px solid var(--line); border-radius:17px; background-color:#edf2fa; background-image:linear-gradient(45deg,#dfe6f2 25%,transparent 25%),linear-gradient(-45deg,#dfe6f2 25%,transparent 25%),linear-gradient(45deg,transparent 75%,#dfe6f2 75%),linear-gradient(-45deg,transparent 75%,#dfe6f2 75%); background-position:0 0,0 11px,11px -11px,-11px 0; background-size:22px 22px; }
    .compare img { position:absolute; inset:0; width:100%; height:100%; object-fit:contain; }
    .after-clip { position:absolute; inset:0; clip-path:inset(0 50% 0 0); }
    .divider { position:absolute; top:0; bottom:0; left:50%; width:2px; background:white; box-shadow:0 0 8px rgba(0,0,0,.4); pointer-events:none; }
    .tag { position:absolute; top:13px; z-index:2; padding:5px 8px; border-radius:7px; color:white; background:rgba(5,10,22,.65); font-size:11px; font-weight:800; letter-spacing:.06em; text-transform:uppercase; }
    .tag.before { left:13px; }
    .tag.after { right:13px; }
    .range { width:100%; margin:18px 0 0; accent-color:var(--blue); cursor:ew-resize; }
    .result-actions { display:flex; align-items:center; justify-content:space-between; gap:12px; margin-top:15px; color:var(--muted); font-size:12px; }
    .download { display:none; padding:8px 11px; border:1px solid rgba(67,221,198,.28); border-radius:9px; color:var(--cyan); text-decoration:none; font-weight:800; }
    .download.show { display:inline-block; }
    .hint { margin-top:16px; color:#6e7e9b; font-size:12px; }
    @media (max-width:850px) { .studio{grid-template-columns:1fr;} .controls{position:static;} }
    @media (max-width:540px) { .shell{width:min(100% - 22px,1180px);padding-top:20px;} .system{display:none;} .hero{margin-top:55px;} .row,.advanced{grid-template-columns:1fr;} .compare{min-height:300px;} }
  </style>
</head>
<body>
  <main class="shell">
    <header class="topbar">
      <div class="brand"><span class="brand-mark">✦</span><div>BACKGROUND STUDIO<small>Advanced cutout workspace</small></div></div>
      <div class="system"><span class="dot"></span> CPU processing ready</div>
    </header>
    <section class="hero">
      <p class="eyebrow">Clean edges. Better cutouts.</p>
      <h1>Make the subject stand out.</h1>
      <p>Upload an image, choose the right segmentation model, fine-tune the edges, and download a transparent result.</p>
    </section>
    <section class="studio">
      <aside class="card controls">
        <div class="card-head"><span>Removal settings</span><span class="pill">ADVANCED</span></div>
        <form id="form">
          <div class="dropzone" id="dropzone">
            <div class="upload-icon">↑</div>
            <strong>Drop an image here</strong>
            <span>PNG, JPG, WEBP · up to 12 MB</span>
            <label class="browse" for="file">Choose image</label>
            <input id="file" type="file" accept="image/png,image/jpeg,image/webp">
            <div class="file-name" id="file-name">No file selected</div>
          </div>
          <label for="model">Segmentation model</label>
          <select id="model">
            <option value="isnet-general-use">Quality · ISNet General</option>
            <option value="u2net">Balanced · U2Net</option>
            <option value="u2netp">Fast · U2NetP</option>
            <option value="u2net_human_seg">People · Human Segmentation</option>
            <option value="isnet-anime">Anime · ISNet Anime</option>
          </select>
          <div class="row">
            <div><label for="format">Output</label><select id="format"><option value="png">PNG · transparent</option><option value="webp">WEBP · smaller</option><option value="jpg">JPG · solid background</option></select></div>
            <div><label for="background">JPG background</label><input id="background" type="color" value="#ffffff" title="Used when JPG output is selected"></div>
          </div>
          <label class="check"><input id="matting" type="checkbox"> Enable alpha matting for finer edges</label>
          <details><summary>Fine-tune alpha matting</summary><div class="advanced"><div><label for="foreground">Foreground</label><input id="foreground" type="number" min="0" max="255" value="240"></div><div><label for="threshold">Background</label><input id="threshold" type="number" min="0" max="255" value="10"></div><div><label for="erode">Erode</label><input id="erode" type="number" min="0" max="40" value="10"></div></div></details>
          <label for="api-key">API key <span style="color:var(--muted);font-weight:500;text-transform:none;letter-spacing:0">(optional)</span></label>
          <input id="api-key" type="password" autocomplete="off" placeholder="Only needed if the Space secret is enabled">
          <button class="run" id="run" type="submit">Remove background →</button>
          <div class="status" id="status">Choose an image to begin.</div>
        </form>
      </aside>
      <section class="card result">
        <div class="result-head"><h2>Preview</h2><span id="meta">Before and after comparison</span></div>
        <div class="empty" id="empty"><div><div class="empty-icon">◌</div>Upload an image and your result will appear here.</div></div>
        <div class="compare-wrap" id="compare-wrap">
          <div class="compare" id="compare"><span class="tag before">Original</span><span class="tag after">Cutout</span><img id="before" alt="Original image"><div class="after-clip" id="after-clip"><img id="after" alt="Background removed image"></div><div class="divider" id="divider"></div></div>
          <input class="range" id="range" type="range" min="0" max="100" value="50" aria-label="Compare original and cutout">
          <div class="result-actions"><span id="details">Transparent PNG ready</span><a class="download" id="download" download="background-removed.png">Download result ↓</a></div>
        </div>
        <p class="hint">Tip: use alpha matting when hair, fur, glass, or fine object edges need extra refinement.</p>
      </section>
    </section>
  </main>
  <script>
    const form = document.getElementById('form');
    const fileInput = document.getElementById('file');
    const dropzone = document.getElementById('dropzone');
    const fileName = document.getElementById('file-name');
    const run = document.getElementById('run');
    const status = document.getElementById('status');
    const empty = document.getElementById('empty');
    const wrap = document.getElementById('compare-wrap');
    const before = document.getElementById('before');
    const after = document.getElementById('after');
    const clip = document.getElementById('after-clip');
    const divider = document.getElementById('divider');
    const range = document.getElementById('range');
    const download = document.getElementById('download');
    const details = document.getElementById('details');
    const meta = document.getElementById('meta');
    let originalUrl = null;
    let resultUrl = null;
    function setFile(file) { if (!file) return; fileInput.files = (() => { const dt = new DataTransfer(); dt.items.add(file); return dt.files; })(); fileName.textContent = file.name; status.textContent = 'Ready to process.'; status.className = 'status'; if (originalUrl) URL.revokeObjectURL(originalUrl); originalUrl = URL.createObjectURL(file); before.src = originalUrl; }
    fileInput.addEventListener('change', () => setFile(fileInput.files[0]));
    ['dragenter','dragover'].forEach((event) => dropzone.addEventListener(event, (e) => { e.preventDefault(); dropzone.classList.add('drag'); }));
    ['dragleave','drop'].forEach((event) => dropzone.addEventListener(event, (e) => { e.preventDefault(); dropzone.classList.remove('drag'); }));
    dropzone.addEventListener('drop', (e) => setFile(e.dataTransfer.files[0]));
    function updateCompare() { const value = Number(range.value); clip.style.clipPath = `inset(0 ${100 - value}% 0 0)`; divider.style.left = `${value}%`; }
    range.addEventListener('input', updateCompare);
    form.addEventListener('submit', async (event) => {
      event.preventDefault();
      const file = fileInput.files[0];
      if (!file) { status.textContent = 'Choose an image first.'; status.className = 'status error'; return; }
      run.disabled = true; status.className = 'status'; status.textContent = 'Processing… the first model run may take longer.'; meta.textContent = 'Working on your cutout';
      const outputFormat = document.getElementById('format').value;
      const body = new FormData(); body.append('image', file); body.append('model', document.getElementById('model').value); body.append('format', outputFormat); body.append('background', outputFormat === 'jpg' ? document.getElementById('background').value : 'transparent'); body.append('alpha_matting', document.getElementById('matting').checked ? 'true' : 'false'); body.append('foreground_threshold', document.getElementById('foreground').value); body.append('background_threshold', document.getElementById('threshold').value); body.append('erode_size', document.getElementById('erode').value);
      const apiKey = document.getElementById('api-key').value.trim();
      try {
        const headers = apiKey ? { 'X-API-Key': apiKey } : {};
        const response = await fetch('/remove-bg', { method:'POST', headers, body });
        if (!response.ok) { const data = await response.json().catch(() => ({})); throw new Error(data.error || `Request failed: ${response.status}`); }
        const blob = await response.blob();
        if (resultUrl) URL.revokeObjectURL(resultUrl); resultUrl = URL.createObjectURL(blob); after.src = resultUrl; download.href = resultUrl; download.classList.add('show'); empty.style.display = 'none'; wrap.classList.add('show'); range.value = 50; updateCompare();
        const model = document.getElementById('model').value; const format = outputFormat.toUpperCase(); details.textContent = `${model} · ${format} ready`; meta.textContent = 'Drag the slider to compare'; status.className = 'status success'; status.textContent = 'Background removed successfully.';
      } catch (error) { status.className = 'status error'; status.textContent = error.message; }
      finally { run.disabled = false; }
    });
  </script>
</body>
</html>
"""


def _client_key():
    forwarded = request.headers.get("X-Forwarded-For", "")
    return forwarded.split(",")[0].strip() or request.remote_addr or "unknown"


def _within_rate_limit():
    now = time.time()
    key = _client_key()
    with RATE_LOCK:
        recent = [stamp for stamp in RATE_BUCKETS.get(key, []) if now - stamp < RATE_WINDOW]
        if len(recent) >= RATE_LIMIT:
            RATE_BUCKETS[key] = recent
            return False
        recent.append(now)
        RATE_BUCKETS[key] = recent
        return True


def _check_api_key():
    configured = os.getenv("API_KEY", "").strip()
    if not configured:
        return True
    supplied = request.headers.get("X-API-Key", "")
    return bool(supplied) and hmac.compare_digest(supplied, configured)


def _int_field(name, default, low, high):
    try:
        value = int(request.form.get(name, default))
    except (TypeError, ValueError):
        raise ValueError(f"{name} must be a number")
    return max(low, min(high, value))


def _get_session(model):
    if model not in MODELS:
        raise ValueError("Unknown segmentation model")
    with SESSION_LOCK:
        if model not in SESSIONS:
            SESSIONS[model] = new_session(model)
        return SESSIONS[model]


def _composite(image, color):
    try:
        rgb = ImageColor.getrgb(color)
    except ValueError as exc:
        raise ValueError("Invalid background color") from exc
    background = Image.new("RGBA", image.size, rgb + (255,))
    background.alpha_composite(image.convert("RGBA"))
    return background


@app.get("/")
def home():
    return INDEX_HTML


@app.get("/health")
def health():
    return jsonify({"status": "healthy", "service": "background-studio"}), 200


@app.post("/remove-bg")
def remove_background():
    if not _within_rate_limit():
        return jsonify({"error": "Rate limit reached. Please wait a minute."}), 429
    if not _check_api_key():
        return jsonify({"error": "Invalid or missing API key"}), 401
    if "image" not in request.files:
        return jsonify({"error": "Choose an image file"}), 400

    uploaded = request.files["image"]
    if not uploaded.filename:
        return jsonify({"error": "No image selected"}), 400

    model = request.form.get("model", "isnet-general-use")
    output_format = request.form.get("format", "png").lower()
    background = request.form.get("background", "#ffffff")
    alpha_matting = request.form.get("alpha_matting", "false").lower() == "true"
    if output_format not in FORMATS:
        return jsonify({"error": "Unsupported output format"}), 400

    try:
        image = ImageOps.exif_transpose(Image.open(uploaded.stream))
        image.load()
        if image.width * image.height > 25_000_000:
            return jsonify({"error": "Image dimensions are too large"}), 400
        image = image.convert("RGBA")
        session = _get_session(model)
        output = remove(
            image,
            session=session,
            alpha_matting=alpha_matting,
            alpha_matting_foreground_threshold=_int_field("foreground_threshold", 240, 0, 255),
            alpha_matting_background_threshold=_int_field("background_threshold", 10, 0, 255),
            alpha_matting_erode_size=_int_field("erode_size", 10, 0, 40),
        ).convert("RGBA")

        if output_format == "jpg" or background.lower() != "transparent":
            output = _composite(output, background if background.lower() != "transparent" else "#ffffff")
        if output_format == "jpg":
            output = output.convert("RGB")

        result = io.BytesIO()
        if output_format == "jpg":
            output.save(result, format="JPEG", quality=95, optimize=True)
        elif output_format == "webp":
            output.save(result, format="WEBP", quality=95, method=6)
        else:
            output.save(result, format="PNG", optimize=True)
        result.seek(0)
        return send_file(result, mimetype=f"image/{'jpeg' if output_format == 'jpg' else output_format}", as_attachment=True, download_name=f"background-removed.{output_format}")
    except (UnidentifiedImageError, OSError):
        return jsonify({"error": "The uploaded file is not a readable image"}), 400
    except ValueError as exc:
        return jsonify({"error": str(exc)}), 400
    except Exception:
        app.logger.exception("Background removal failed")
        return jsonify({"error": "Background removal failed. Try a smaller image or another model."}), 500


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=int(os.getenv("PORT", "7860")))