Spaces:
Paused
Paused
| 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 | |
| def home(): | |
| return INDEX_HTML | |
| def health(): | |
| return jsonify({"status": "healthy", "service": "background-studio"}), 200 | |
| 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"))) | |