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""" Background Studio

Clean edges. Better cutouts.

Make the subject stand out.

Upload an image, choose the right segmentation model, fine-tune the edges, and download a transparent result.

Preview

Before and after comparison

◌

Upload an image and your result will appear here.

OriginalCutout Original image

Transparent PNG readyDownload result ↓

Tip: use alpha matting when hair, fur, glass, or fine object edges need extra refinement.

""" 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")))