$ show danielgatis/rembg
Rembg is a production-grade tool for removing image backgrounds, usable as a Python library, CLI, HTTP server, or Docker container. It wraps multiple ONNX-based segmentation models (U2Net, BiRefNet, SAM, BRIA-RMBG, etc.) and provides several cutout strategies: naive, alpha matting, color decontamination, and ViTMatte refinement. The architecture is modular: a session abstraction handles model download and inference, while the `remove()` function in `rembg/bg.py` orchestrates preprocessing, mask prediction, and post-processing. The CLI offers subcommands for single files (`i`), batch processing (`p`), HTTP serving (`s`), binary stream input (`b`), model download (`d`), and legacy migration (`m`). The HTTP server uses FastAPI with a Gradio UI and includes security measures like DNS rebinding protection. The codebase is well-tested, with synthetic fixtures for color decontamination and integration tests using perceptual hashing. It demonstrates strong engineering practices: idempotent migration, session caching, and careful handling of edge cases like oversized uploads and failed alpha matting convergence.
The LLM identified 5 files worth highlighting for reuse:
Implements ViTMatte-based alpha refinement as a drop-in replacement for closed-form alpha matting. The module handles model variant selection, download via pooch, ONNX session caching, trimap generation from a coarse mask, and fixed-resolution inference with bilinear resizing. It is self-contained (only depends on BaseSession for path utilities) and demonstrates non-obvious skill in managing fixed-input-size models and fallback logic. Could be extracted into a standalone matting library with minimal changes.
A well-crafted CLI command for migrating model files from a legacy flat directory layout to a per-model directory structure. Features dry-run mode, idempotent copy (originals preserved unless --delete-source), orphan cleanup, and size verification before deletion. The pattern of safe, reversible data migration with progress reporting is reusable for any project that needs to restructure file storage.
Processes raw RGB24 pixel streams from stdin (e.g., from FFmpeg) into background-removed PNG images. Uses asyncio for efficient I/O, handles variable-length streams, and supports printf-style output filenames. This pattern is directly reusable for any pipeline that needs to process raw video frames through a machine learning model.
Implements a full HTTP server with FastAPI and optional Gradio UI. Includes security-critical URL validation that resolves hostnames and rejects private/internal IPs, with a custom aiohttp resolver to prevent DNS rebinding attacks. Also features session caching keyed by model name and extras, thread limiting, and CORS configuration. The URL validation and resolver pattern is reusable for any service that fetches remote images.
Contains a synthetic fixture generator (`make_fringed_image`) that composites semi-transparent strands over a saturated background to produce a known ground-truth for testing color decontamination. The test verifies that decontamination recovers the true foreground color within 10 units, while naive cutout leaves a large error. This pattern of building controlled synthetic test data for image processing algorithms is reusable across any project that needs to validate color unmixing or alpha matting.