103 lines
2.7 KiB
Markdown
103 lines
2.7 KiB
Markdown
# Background Removal Tool
|
|
|
|
A Rust-based background removal tool using ONNX models (BiRefNet) with GPU support.
|
|
|
|
## Features
|
|
|
|
- Download images from URLs
|
|
- Automatic ONNX model download and caching
|
|
- Background removal using BiRefNet-general-lite model
|
|
- Side-by-side comparison display (original vs processed)
|
|
- Checkered background visualization for transparency
|
|
- CLI with flexible options
|
|
|
|
## Building
|
|
|
|
### With Nix (Recommended)
|
|
|
|
```bash
|
|
# Enter the development environment
|
|
nix develop
|
|
|
|
# Build the project
|
|
cargo build --release
|
|
```
|
|
|
|
### Without Nix
|
|
|
|
If you don't use Nix, you'll need to ensure you have the following system dependencies:
|
|
- Rust toolchain
|
|
- OpenSSL development libraries (`libssl-dev` on Ubuntu/Debian)
|
|
- pkg-config
|
|
- X11 libraries (for window display)
|
|
- Vulkan loader
|
|
|
|
## Usage
|
|
|
|
Basic usage:
|
|
|
|
```bash
|
|
# Download and process an image
|
|
cargo run -- https://example.com/image.jpg
|
|
|
|
# Use custom model
|
|
cargo run -- https://example.com/image.jpg --model-path ./custom_model.onnx
|
|
|
|
# Offline mode (use only cached model)
|
|
cargo run -- https://example.com/image.jpg --offline
|
|
```
|
|
|
|
## How It Works
|
|
|
|
1. **Download**: Fetches the image from the provided URL
|
|
2. **Model Loading**: Auto-downloads BiRefNet model (cached in `~/.cache/remove_background/models/`)
|
|
3. **Preprocessing**: Resizes image to 1024x1024, normalizes pixels, converts to CHW format
|
|
4. **Inference**: Runs ONNX model to generate foreground mask
|
|
5. **Postprocessing**: Applies mask to original image, creates RGBA output
|
|
6. **Display**: Shows side-by-side comparison with checkered background
|
|
|
|
## Project Structure
|
|
|
|
```
|
|
src/
|
|
├── main.rs - CLI interface and orchestration
|
|
├── model.rs - Model download and ONNX session management
|
|
├── preprocessing.rs - Image preprocessing (resize, normalize, CHW conversion)
|
|
└── postprocessing.rs - Mask application and side-by-side display creation
|
|
```
|
|
|
|
## Dependencies
|
|
|
|
- **ort**: ONNX Runtime Rust bindings
|
|
- **ndarray**: N-dimensional array operations
|
|
- **image**: Image loading and manipulation
|
|
- **reqwest**: HTTP client for downloads
|
|
- **show-image**: Cross-platform window display
|
|
- **clap**: Command-line argument parsing
|
|
- **sha2**: Hash verification
|
|
|
|
## Future Enhancements
|
|
|
|
- [ ] CUDA execution provider support (GPU acceleration)
|
|
- [ ] Batch processing support
|
|
- [ ] Additional model options (U2-Net, ISNet, etc.)
|
|
- [ ] Save output to file
|
|
- [ ] Local file input support
|
|
- [ ] Progress bars for downloads
|
|
- [ ] WebAssembly support
|
|
|
|
## Controls
|
|
|
|
- **ESC**: Close the window and exit
|
|
|
|
## Model Information
|
|
|
|
Default model: **BiRefNet-general-lite**
|
|
- Size: ~50MB
|
|
- Input: 1024x1024 RGB
|
|
- Output: 1024x1024 mask
|
|
- Source: https://huggingface.co/ZhengPeng7/BiRefNet
|
|
|
|
## License
|
|
|
|
See LICENSE file for details.
|