BiRefNet

https://github.com/zhengpeng7/birefnet
Admin
BiRefNet implements Bilateral Reference for High-Resolution Dichotomous Image Segmentation, offering...
Tokens:48,586
Snippets:504
Trust Score:8.3
License:MIT
Update:3 months ago
Show doc for...
Tokens:
Raw
### Environment Setup for BiRefNet

Source: https://github.com/zhengpeng7/birefnet/blob/main/README.md

Installs PyTorch and project dependencies. Ensure PyTorch version is 2.5.0 or later for faster training.

```shell
conda create -n birefnet python=3.11 -y && conda activate birefnet
pip install -r requirements.txt
```

--------------------------------

### ResBlk Example Usage

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/decoder-blocks.md

Demonstrates how to instantiate and use the ResBlk class. This example shows creating a block with channel expansion and processing a random feature map, verifying the output shape.

```python
from models.modules.decoder_blocks import ResBlk
import torch

# Create a residual decoder block with channel expansion
res_blk = ResBlk(in_channels=256, out_channels=128, inter_channels=64)

# Process feature map
x = torch.randn(4, 256, 32, 32)
output = res_blk(x)  # Shape: (4, 128, 32, 32)

# Verify residual connection preserves spatial dimensions
assert output.shape == (4, 128, 32, 32)
```

--------------------------------

### Batch Inference Command Line Examples

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/inference-api.md

Examples of how to run the inference script from the command line for different scenarios. Use --ckpt for a single checkpoint or --ckpt_folder for multiple.

```bash
python inference.py --ckpt ./ckpts/BiRefNet-epoch_100.pth --testsets DIS-VD
```

```bash
python inference.py --ckpt_folder ./ckpts/tmp --testsets DIS-VD+DIS-TE1
```

```bash
python inference.py --ckpt_folder ./ckpts/tmp --resolution 2048x2048
```

```bash
python inference.py --ckpt_folder ./ckpts/tmp --resolution None
```

--------------------------------

### Example Usage of build_backbone

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/backbones.md

Demonstrates how to instantiate different backbone architectures using the `build_backbone` function. Shows loading with and without pretrained weights, and integrating into the BiRefNet model.

```python
from models.backbones.build_backbone import build_backbone
import torch

# Load Swin-Large with pretrained weights
backbone = build_backbone('swin_v1_l', pretrained=True)

# Load ResNet50 without pretrained weights
backbone = build_backbone('resnet50', pretrained=False)

# Load DINO v3 base
backbone = build_backbone('dino_v3_b', pretrained=True)

# Use in BiRefNet
from models.birefnet import BiRefNet
model = BiRefNet(bb_pretrained=True)
```

--------------------------------

### Example Usage of MyData Dataset and DataLoader

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/dataset-utilities.md

Demonstrates creating training and evaluation datasets, accessing samples, and using a DataLoader for batch processing.

```python
from dataset import MyData
import torch

# Training dataset
train_dataset = MyData(
    datasets='DIS-TR',
    data_size=(1024, 1024),
    is_train=True
)
print(f"Training samples: {len(train_dataset)}")

# Get a sample
image, mask, class_label = train_dataset[0]
print(f"Image shape: {image.shape}")      # (3, 1024, 1024)
print(f"Mask shape: {mask.shape}")        # (1, 1024, 1024)

# Evaluation dataset
test_dataset = MyData(
    datasets='DIS-VD',
    data_size=(1024, 1024),
    is_train=False
)
image, mask, label_path = test_dataset[0]
print(f"Label path: {label_path}")

# DataLoader
train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=4,
    shuffle=True,
    num_workers=4,
    pin_memory=True
)

for images, masks, labels in train_loader:
    print(f"Batch - Images: {images.shape}, Masks: {masks.shape}")
    break
```

--------------------------------

### PixLoss Example Usage

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/loss-functions.md

Demonstrates how to instantiate and use the PixLoss class. It shows the expected format for predictions and ground truth tensors.

```python
from loss import PixLoss
import torch

pix_loss = PixLoss()
preds = [torch.randn(4, 1, 32, 32), torch.randn(4, 1, 64, 64), torch.randn(4, 1, 128, 128)]
gt = torch.randint(0, 2, (4, 1, 128, 128)).float()
loss = pix_loss(preds, gt)
```

--------------------------------

### StructureLoss Example Usage

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/loss-functions.md

Demonstrates how to instantiate and use the StructureLoss. Ensure 'loss' module is imported and tensors are correctly shaped.

```python
from loss import StructureLoss
import torch

struct_loss = StructureLoss()
pred_logits = torch.randn(4, 1, 256, 256)  # Before sigmoid
target = torch.randint(0, 2, (4, 1, 256, 256)).float()
loss = struct_loss(pred_logits, target)

```

--------------------------------

### Example Dynamic Size Configuration

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/configuration.md

Sets a dynamic resolution range for training or inference. This allows for variable input sizes within a specified range.

```python
config.dynamic_size = ((512-256, 2048+256), (512-256, 2048+256))  # 256-2304 range
```

--------------------------------

### Install ONNX and related packages

Source: https://github.com/zhengpeng7/birefnet/blob/main/tutorials/BiRefNet_pth2onnx.ipynb

Installs the ONNX, ONNX-Script, and ONNX Runtime GPU packages. These are essential for converting and running ONNX models.

```python
!pip install onnx onnxscript onnxruntime-gpu
```

--------------------------------

### Launch Multi-GPU Training with Accelerate

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/training-api.md

This command launches distributed training using HuggingFace's Accelerate library. First, configure your hardware setup with 'accelerate config'. Accelerate simplifies multi-GPU training, automatic mixed precision, and gradient accumulation.

```bash
accelerate config  # Configure hardware setup once

accelerate launch \
    --multi_gpu \
    --num_processes=4 \
    train.py --use_accelerate --epochs 120
```

--------------------------------

### Fine-tuning BiRefNet Configuration and Training

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/training-api.md

This snippet shows how to prepare data, update configuration, adjust loss weights, load pretrained models, set learning rates, and start training for custom tasks.

```python
# 1. Prepare data in correct structure
# data_root_dir/CUSTOM_TASK/CUSTOM_DATASET/{im,gt}/

# 2. Update config
config.task = 'CUSTOM_TASK'
config.training_set = 'CUSTOM_DATASET'
config.testsets = 'CUSTOM_DATASET_VAL'

# 3. Adjust loss weights if needed (matting vs segmentation)
config.lambdas_pix_last['mae'] = 150  # For matting tasks

# 4. Load pretrained weights
model = BiRefNet(bb_pretrained=True)

# 5. Train with lower learning rate
config.lr = config.lr / 10

# 6. Run training
python train.py --epochs 50 --ckpt_dir ./ckpts/custom
```

--------------------------------

### Initialize and Use IoULoss

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/loss-functions.md

Demonstrates how to initialize the IoULoss class and compute IoU loss using predicted and target segmentation masks. Ensure PyTorch is installed and the loss module is accessible.

```python
from loss import IoULoss
import torch

iou_loss = IoULoss()
pred = torch.sigmoid(torch.randn(4, 1, 256, 256))
target = torch.randint(0, 2, (4, 1, 256, 256)).float()
loss = iou_loss(pred, target)
print(f"IoU Loss: {loss.item()}")
```

--------------------------------

### ASPP Module Initialization and Forward Pass Example

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/attention-modules.md

Demonstrates how to initialize and use the ASPP module for processing feature maps. Ensure the input tensor matches the expected dimensions and channels.

```python
from models.modules.aspp import ASPP
import torch

# Create ASPP module
aspp = ASPP(in_channels=256, out_channels=256, output_stride=16)

# Process feature map
x = torch.randn(4, 256, 32, 32)
output = aspp(x)  # Shape: (4, 256, 32, 32)
print(output.shape)
```

--------------------------------

### Multi-Dataset Training Setup

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/dataset-utilities.md

Combines multiple datasets into a single training set by specifying them as a concatenated string. Useful for augmenting training data.

```python
# Combine multiple datasets
train_dataset = MyData(
    datasets='DIS-TR+TR-HRSOD+TR-CAMO',  # Multiple datasets
    data_size=(1024, 1024),
    is_train=True
)
```

--------------------------------

### Load BiRefNet Model from Hugging Face

Source: https://github.com/zhengpeng7/birefnet/blob/main/tutorials/BiRefNet_inference_video.ipynb

Loads the BiRefNet model and its weights from the Hugging Face Hub. This is a convenient option for quick setup.

```python
# Imports
from PIL import Image
import torch
from torchvision import transforms
from IPython.display import display

sys.path.insert(0, "../")
from models.birefnet import BiRefNet


# Load Model
# Option 2 and Option 3 is better for local running -- we can modify codes locally.

# # # Option 1: loading BiRefNet with weights:
# from transformers import AutoModelForImageSegmentation
# birefnet = AutoModelForImageSegmentation.from_pretrained('zhengpeng7/BiRefNet', trust_remote_code=True)

# Option-2: loading weights with BiReNet codes:
model = [
        'zhengpeng7/BiRefNet',
        'zhengpeng7/BiRefNet-portrait',
        'zhengpeng7/BiRefNet-legacy', 'zhengpeng7/BiRefNet-DIS5K-TR_TEs', 'zhengpeng7/BiRefNet-DIS5K', 'zhengpeng7/BiRefNet-HRSOD', 'zhengpeng7/BiRefNet-COD',
        'zhengpeng7/BiRefNet_lite',     # Modify the `bb` in `config.py` to `swin_v1_tiny`.
    ][0]
birefnet = BiRefNet.from_pretrained(
    model
)
model_name = model.split('/')[-1]

# # Option-3: Loading model and weights from local disk:
# from utils import check_state_dict

# birefnet = BiRefNet(bb_pretrained=False)
# state_dict = torch.load('../BiRefNet-general-epoch_244.pth', map_location='cpu', weights_only=True)
# state_dict = check_state_dict(state_dict)
# birefnet.load_state_dict(state_dict)

device = 'cuda' if torch.cuda.is_available() else 'cpu'

torch.set_float32_matmul_precision(['high', 'highest'][0])

birefnet.to(device)
birefnet.eval()
print('BiRefNet is ready to use.')

# Input Data
transform_image = transforms.Compose([
    transforms.Resize((1024, 1024)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
```

--------------------------------

### Example Usage of Inference Function

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/inference-api.md

Demonstrates how to load a BiRefNet model, create a test dataset and data loader, and then run batch inference using the `inference` function. Ensure the model checkpoint and dataset paths are correct.

```python
from inference import inference
from dataset import MyData
import torch
from models.birefnet import BiRefNet

# Load model
model = BiRefNet(bb_pretrained=False)
model.load_state_dict(torch.load('checkpoint.pth'))
model = model.cuda()

# Create test dataset and loader
test_dataset = MyData(datasets='DIS-VD', data_size=(1024, 1024), is_train=False)
test_loader = torch.utils.data.DataLoader(
    test_dataset,
    batch_size=8,
    shuffle=False,
    num_workers=4,
    pin_memory=True
)

# Run inference
inference(
    model=model,
    data_loader_test=test_loader,
    pred_root='predictions',
    method='BiRefNet-epoch_100',
    testset='DIS-VD',
    device=0
)
```

--------------------------------

### Channel Adaptation for Swin-Large Backbone

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/lateral-blocks.md

Example of initializing lateral blocks when using a Swin-Large backbone. The input channels of the lateral blocks are derived from the backbone's output channels, and the output channels match the decoder's expected dimensions.

```python
# Swin-Large encoder outputs channels: [1536, 768, 384, 192]
# Decoder expects: [768, 384, 192, 96]

lat4 = BasicLatBlk(in_channels=768, out_channels=384)   # Skip from level 3
lat3 = BasicLatBlk(in_channels=384, out_channels=192)   # Skip from level 2
lat2 = BasicLatBlk(in_channels=192, out_channels=96)    # Skip from level 1
```

--------------------------------

### Complete BiRefNet Inference Pipeline

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/inference-api.md

Provides a step-by-step guide to performing inference using a BiRefNet model. It covers loading the model from a checkpoint, loading and preprocessing an input image, running the inference, and post-processing the output to save it as an image or a binary mask.

```python
import torch
import cv2
from PIL import Image
from torchvision import transforms
from models.birefnet import BiRefNet
from utils import check_state_dict, save_tensor_img

# 1. Load model
model = BiRefNet(bb_pretrained=False)
checkpoint = torch.load('model.pth', map_location='cpu', weights_only=True)
checkpoint = check_state_dict(checkpoint)
model.load_state_dict(checkpoint)
model = model.cuda()
model.eval()

# 2. Load image
image_path = 'input.jpg'
image = Image.open(image_path).convert('RGB')
original_size = image.size  # Save for later

# 3. Preprocess
image = image.resize((1024, 1024))
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
x = transform(image).unsqueeze(0).cuda()

# 4. Inference
with torch.no_grad():
    output = model(x)

# 5. Post-process
pred = torch.sigmoid(output).cpu()
pred = torch.nn.functional.interpolate(
    pred,
    size=original_size[::-1],  # PIL uses (W, H), interpolate uses (H, W)
    mode='bilinear',
    align_corners=True
)

# 6. Save
save_tensor_img(pred, 'output.png')

# 7. Optional: binary mask
binary_mask = (pred > 0.5).float()
save_tensor_img(binary_mask, 'mask_binary.png')
```

--------------------------------

### Instantiate and Use BasicDecBlk

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/decoder-blocks.md

Demonstrates how to create an instance of BasicDecBlk and process an input feature map, showing the resulting output shape.

```python
from models.modules.decoder_blocks import BasicDecBlk
import torch

# Create a decoder block
dec_blk = BasicDecBlk(in_channels=256, out_channels=128, inter_channels=64)

# Process feature map
x = torch.randn(4, 256, 32, 32)
output = dec_blk(x)  # Shape: (4, 128, 32, 32)
print(output.shape)
```

--------------------------------

### Initialize PyTorch Environment and Load Weights

Source: https://github.com/zhengpeng7/birefnet/blob/main/tutorials/BiRefNet_pth2onnx.ipynb

Sets up the PyTorch environment by defining the weights file and selecting the appropriate device (CUDA or CPU). It imports necessary libraries for model definition and state dictionary checking.

```python
import torch


weights_file = 'BiRefNet_dynamic-general-epoch_174.pth'  # https://github.com/ZhengPeng7/BiRefNet/releases/download/v1/BiRefNet_dynamic-general-epoch_174.pth
device = 'cuda' if torch.cuda.is_available() else 'cpu'
```

--------------------------------

### Mixed Precision Inference with PyTorch

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/inference-api.md

Utilize mixed-precision inference on CUDA devices for improved performance. This example demonstrates using `torch.amp.autocast` with `torch.float16`.

```python
import torch
from contextlib import nullcontext

model = BiRefNet(bb_pretrained=True).cuda()
model.eval()

# FP16 inference
dtype = torch.float16
autocast_ctx = torch.amp.autocast(device_type='cuda', dtype=dtype)

x = torch.randn(4, 3, 1024, 1024, dtype=dtype).cuda()
with autocast_ctx, torch.no_grad():
    output = model(x)

pred = torch.sigmoid(output)
print(pred.shape)  # (4, 1, 1024, 1024)
```

--------------------------------

### Instantiate and Use BasicLatBlk

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/lateral-blocks.md

Demonstrates how to create an instance of BasicLatBlk and use it to transform an encoder feature map. The output shape is printed to verify the transformation.

```python
from models.modules.lateral_blocks import BasicLatBlk
import torch

# Create lateral block to transform encoder features
lat_blk = BasicLatBlk(in_channels=384, out_channels=256)

# Process encoder feature map
x = torch.randn(4, 384, 32, 32)
output = lat_blk(x)  # Shape: (4, 256, 32, 32)
print(output.shape)
```

--------------------------------

### Typical Integration of ASPP or ASPPDeformable

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/attention-modules.md

Shows how to conditionally instantiate either ASPP or ASPPDeformable based on the configuration and how to apply the chosen module in the forward pass.

```python
# In decoder block initialization
if config.dec_att == 'ASPP':
    self.dec_att = ASPP(in_channels=inter_channels)
elif config.dec_att == 'ASPPDeformable':
    self.dec_att = ASPPDeformable(in_channels=inter_channels)

# In decoder block forward
if hasattr(self, 'dec_att'):
    x = self.dec_att(x)

```

--------------------------------

### Configure Lightweight Model for Edge Devices

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/configuration.md

Set up configuration for a lightweight model suitable for edge devices, using a tiny backbone and fp16 mixed precision. Compilation may not be beneficial.

```python
config.bb = 'swin_v1_t'  # Tiny Swin
config.batch_size = 16
config.mixed_precision = 'fp16'
config.compile = False  # May not help on edge devices
```

--------------------------------

### Loss Configuration Schema

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/types-and-schemas.md

Defines the structure for loss configuration dictionaries, specifying weights for various pixel-level loss components. Examples for matting and segmentation tasks are provided.

```python
lambda_dict = {
    # Pixel-level losses
    'bce': float,              # Binary cross-entropy weight
    'iou': float,              # IoU loss weight
    'iou_patch': float,        # Patch-wise IoU weight
    'mae': float,              # L1/MAE loss weight
    'mse': float,              # Mean squared error weight
    'ssim': float,             # Structural similarity weight
    'triplet': float,          # Triplet loss weight
    'reg': float,              # Threshold regularization weight
    'cnt': float,              # Contour loss weight
    'structure': float,        # Structure loss weight (boundary emphasis)
}
```

```python
# Example for matting task
matting_lambdas = {
    'bce': 30,
    'mae': 100,    # Higher for alpha accuracy
    'ssim': 10,
    'others': 0    # Disabled
}
```

```python
# Example for segmentation
segmentation_lambdas = {
    'bce': 30,
    'iou': 0.5,
    'mae': 100,
    'ssim': 10,
    'others': 0
}
```

--------------------------------

### Access BiRefNet Configuration

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/configuration.md

Demonstrates how to import and instantiate the global Config object for use in models, training, and data loading. The Config instance is a singleton and is reused throughout training.

```python
from config import Config

config = Config()

# In models
model = BiRefNet(bb_pretrained=True)  # Uses global config

# In training
from loss import PixLoss
pix_loss = PixLoss()  # Reads config.lambdas_pix_last

# In data loading
from dataset import MyData
dataset = MyData(config.training_set, config.size, is_train=True)
```

--------------------------------

### Configure Matting with Dynamic Resolutions

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/configuration.md

Set up configuration for a Matting task with dynamic input resolutions and adjust loss weights for alpha accuracy.

```python
config.task = 'Matting'
config.dynamic_size = ((512-256, 2048+256), (512-256, 2048+256))
config.bb = 'swin_v1_l'
config.lambdas_pix_last['mae'] = 150  # Higher L1 emphasis for alpha
```

--------------------------------

### Configure DIS5K Task with High-Resolution Inputs

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/configuration.md

Set up configuration for a DIS5K task with specified input size, backbone, batch size, and mixed precision.

```python
from config import Config

config = Config()
config.task = 'DIS5K'
config.size = (1024, 1024)
config.bb = 'swin_v1_l'
config.batch_size = 8
config.mixed_precision = 'bf16'
```

--------------------------------

### MyData Constructor

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/dataset-utilities.md

Initializes the MyData dataset. Specify dataset names, target size, and training mode.

```python
def __init__(self, datasets, data_size, is_train=True):
```

--------------------------------

### Initialize Config Class

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/configuration.md

Initializes the configuration parameters with default values. This constructor requires no arguments.

```python
def __init__(self):

```

--------------------------------

### Preloading All Images Configuration

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/dataset-utilities.md

Sets the configuration to load all images into CPU memory. This speeds up training by eliminating disk I/O but requires significant RAM.

```python
config.load_all = True
```

--------------------------------

### Resume Training from Checkpoint

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/training-api.md

To resume training, provide the path to the checkpoint file and the desired total number of epochs.

```bash
python train.py --resume ./ckpts/tmp/BiRefNet-general-epoch_50.pth --epochs 150
```

--------------------------------

### Training Configuration Settings

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/training-api.md

Configure key training parameters such as optimizer, learning rate, decay schedule, precision, and loss weighting in `config.py`.

```python
# Optimization
config.optimizer = 'AdamW'           # Adam optimizer
config.lr = 1e-4 * sqrt(batch_size / 4)  # Scaled learning rate
config.lr_decay_epochs = [1e5]       # LR decay schedule
config.lr_decay_rate = 0.5           # Decay multiplier

# Precision & Compilation
config.compile = True                # Use torch.compile
config.mixed_precision = 'bf16'      # fp16, bf16, no, fp8
config.batch_size = 8                # Batch size per GPU

# Loss Weighting
config.lambdas_pix_last = {
    'bce': 30,
    'iou': 0.5,
    'mae': 100,
    'ssim': 10,
    # ... other losses
}
```

--------------------------------

### Initialize and Use AverageMeter Utility

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/training-api.md

Initializes an AverageMeter utility from the 'utils' module to track running averages of loss during training. It updates the meter with loss values and batch size, and prints the average loss.

```python
from utils import AverageMeter

loss_meter = AverageMeter()
loss_meter.update(loss_value.item(), batch_size)

print(f"Average loss: {loss_meter.avg:.4f}")
```

--------------------------------

### Single GPU Training

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/training-api.md

Use this command for training on a single GPU. Specify the number of epochs and the checkpoint directory.

```bash
python train.py --epochs 120 --ckpt_dir ./ckpts/tmp
```

--------------------------------

### Initialize BiRefNet Model

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/birefnet-core-model.md

Instantiate the BiRefNet model. Set `bb_pretrained` to True to load pretrained backbone weights.

```python
model = BiRefNet(bb_pretrained=True)
model.eval()
x = torch.randn(2, 3, 1024, 1024)  # Batch of 2 images
with torch.no_grad():
    features, class_preds = model.forward_enc(x)
# features is a tuple of 4 feature maps, class_preds is None in eval mode
```

--------------------------------

### Load Pretrained Model and Fine-tune

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/INDEX.md

Load a pretrained BiRefNet model and fine-tune it on a new dataset. Ensure the state dict is compatible before loading.

```python
from models.birefnet import BiRefNet
import torch

model = BiRefNet(bb_pretrained=True)
checkpoint = torch.load('best_checkpoint.pth', weights_only=True)
from utils import check_state_dict
checkpoint = check_state_dict(checkpoint)
model.load_state_dict(checkpoint)

# Fine-tune with lower learning rate
model = model.cuda()
model.train()
```

--------------------------------

### Configure Research with DINO Foundation Model

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/configuration.md

Set up configuration for research using a DINO foundation model, enabling features like frozen backbone, multi-scale supervision, and specific decoder attention.

```python
config.bb = 'dino_v3_l'
config.freeze_bb = True  # Auto-set for DINO
config.ms_supervision = True
config.out_ref = True
config.dec_att = 'ASPPDeformable'  # Best performance
```

--------------------------------

### Launch Multi-GPU Training with DDP

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/training-api.md

Use this command to launch distributed training with PyTorch's native Distributed Data Parallel (DDP). Specify the number of GPUs and the master port. This method synchronizes gradients across GPUs for effective distributed training.

```bash
python -m torch.distributed.launch \
    --nproc_per_node=4 \
    --master_port=29500 \
    train.py --dist True --epochs 120
```

--------------------------------

### BasicLatBlk Constructor

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/lateral-blocks.md

Defines the parameters for the BasicLatBlk, including input/output channels, kernel size, stride, and padding.

```python
def __init__(self, in_channels=64, out_channels=64, ks=1, s=1, p=0):
    pass
```

--------------------------------

### Load and Preprocess Image for Inference

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/inference-api.md

Demonstrates loading an image using either OpenCV (cv2) or PIL, converting it to RGB, resizing it, and converting it to a normalized tensor suitable for PyTorch models. This is a prerequisite for feeding images into the model.

```python
import cv2
from PIL import Image
import torch
from torchvision import transforms

# Method 1: Using cv2
image = cv2.imread('image.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# Method 2: Using PIL
image = Image.open('image.jpg').convert('RGB')

# Resize
image = image.resize((1024, 1024))

# Convert to tensor with normalization
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
x = transform(image).unsqueeze(0)  # Add batch dimension
```

--------------------------------

### Load BiRefNet Model and Set Up Device

Source: https://github.com/zhengpeng7/birefnet/blob/main/tutorials/BiRefNet_inference.ipynb

Loads the BiRefNet model using pre-trained weights and configures the device (GPU or CPU) for inference. It also sets the precision for matrix multiplication.

```python
# Imports
from PIL import Image
import torch
from torchvision import transforms
from IPython.display import display

sys.path.insert(0, "../")
from models.birefnet import BiRefNet


# Load Model
# Option 2 and Option 3 is better for local running -- we can modify codes locally.

# # # Option 1: loading BiRefNet with weights:
# from transformers import AutoModelForImageSegmentation
# birefnet = AutoModelForImageSegmentation.from_pretrained('zhengpeng7/BiRefNet', trust_remote_code=True)

# Option-2: loading weights with BiReNet codes:
birefnet = BiRefNet.from_pretrained(
    [
        'zhengpeng7/BiRefNet',
        'zhengpeng7/BiRefNet-portrait',
        'zhengpeng7/BiRefNet-legacy', 'zhengpeng7/BiRefNet-DIS5K-TR_TEs', 'zhengpeng7/BiRefNet-DIS5K', 'zhengpeng7/BiRefNet-HRSOD', 'zhengpeng7/BiRefNet-COD',
        'zhengpeng7/BiRefNet_lite',     # Modify the `bb` in `config.py` to `swin_v1_tiny`.
    ][0]
)

# # Option-3: Loading model and weights from local disk:
# from utils import check_state_dict

# birefnet = BiRefNet(bb_pretrained=False)
# state_dict = torch.load('../BiRefNet-general-epoch_244.pth', map_location='cpu', weights_only=True)
# state_dict = check_state_dict(state_dict)
# birefnet.load_state_dict(state_dict)

device = 'cuda' if torch.cuda.is_available() else 'cpu'

torch.set_float32_matmul_precision(['high', 'highest'][0])

birefnet.to(device)
birefnet.eval()
print('BiRefNet is ready to use.')

# Input Data
transform_image = transforms.Compose([
    transforms.Resize((1024, 1024)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
```

--------------------------------

### ResBlk Constructor

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/decoder-blocks.md

Initializes a ResBlk instance. Configure input, output, and intermediate channel dimensions. Defaults are provided for common use cases.

```python
def __init__(self, in_channels=64, out_channels=None, inter_channels=64):
    pass
```

--------------------------------

### Available Optimizers

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/types-and-schemas.md

Lists the supported optimization algorithms for training the model.

```python
optimizers = ['Adam', 'AdamW']
```

--------------------------------

### Navigate to Project Directory

Source: https://github.com/zhengpeng7/birefnet/blob/main/tutorials/BiRefNet_pth2onnx.ipynb

Changes the current directory to the parent directory. This is often a prerequisite for running subsequent scripts or accessing project files.

```bash
cd ..
```

--------------------------------

### Initialize and Use Logger Utility

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/training-api.md

Initializes a logger utility from the 'utils' module to write informational messages to a log file. Log messages include epoch, loss, and learning rate.

```python
from utils import Logger

logger = Logger('training.log')
logger.info(f'Epoch {epoch}: loss={loss:.4f}, lr={optimizer.param_groups[0]["lr"]:.6f}')
logger.info('Training completed')
logger.close()
```

--------------------------------

### Multi-GPU Accelerate Training

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/training-api.md

Utilize Hugging Face Accelerate for multi-GPU training, supporting mixed precision like FP16/BF16. Specify the number of processes.

```bash
accelerate launch --multi_gpu --num_processes=4 train.py --use_accelerate --epochs 120
```

--------------------------------

### Initializing and Running Evaluation Metrics

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/inference-api.md

Use the `Evaluator` class from `evaluation.metrics` to compute performance metrics after inference. Configure it with ground truth and prediction roots, methods, datasets, and desired metrics.

```python
# From eval_existingOnes.py or evaluation/metrics.py
from evaluation.metrics import Evaluator

evaluator = Evaluator(
    gt_root='ground_truth',
    pred_root='predictions',
    methods=['BiRefNet-epoch_100'],
    datasets=['DIS-VD'],
    metrics=['Sm', 'wFm', 'HCE']
)

results = evaluator.evaluate()
```

--------------------------------

### MyData __getitem__ Method

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/dataset-utilities.md

Retrieves a sample (image, mask, label) by index. Returns class label for training and file path for evaluation.

```python
def __getitem__(self, index: int) -> tuple:
```

--------------------------------

### Configuration Class

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/INDEX.md

Details the `Config` class for managing all training and model settings.

```python
Config() (config.py)
├── __init__()
└── Attributes: All training/model settings
```

--------------------------------

### Load BiRefNet Model and State Dictionary

Source: https://github.com/zhengpeng7/birefnet/blob/main/tutorials/BiRefNet_pth2onnx.ipynb

Initializes the BiRefNet model, loads the state dictionary from the specified weights file, checks its integrity, and moves the model to the selected device. It also sets the float32 matmul precision for performance.

```python
from utils import check_state_dict
from models.birefnet import BiRefNet


birefnet = BiRefNet(bb_pretrained=False)
state_dict = torch.load('./{}'.format(weights_file), map_location=device, weights_only=True)
state_dict = check_state_dict(state_dict)
birefnet.load_state_dict(state_dict)

torch.set_float32_matmul_precision(['high', 'highest'][0])

birefnet.to(device)
_ = birefnet.eval()
```

--------------------------------

### Standard On-Demand Image Loading Configuration

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/dataset-utilities.md

Sets the configuration to load images on-demand from disk. This is the default behavior, offering memory efficiency suitable for larger datasets at the cost of slower training due to disk I/O.

```python
config.load_all = False
```

--------------------------------

### BiRefNet Training Script

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/training-api.md

This script initializes the BiRefNet model, sets up data loaders, defines loss functions, and executes the training loop. It supports mixed-precision training and includes logging and checkpoint saving.

```python
import torch
import os
from config import Config
from models.birefnet import BiRefNet
from loss import PixLoss, ClsLoss
from dataset import MyData, custom_collate_fn
from utils import Logger, AverageMeter, set_seed, save_checkpoint

# Configuration
config = Config()
set_seed(42)

# Create directories
os.makedirs('ckpts/tmp', exist_ok=True)

# Initialize logger
logger = Logger('ckpts/tmp/log.txt')

# Create datasets
train_dataset = MyData(
    datasets=config.training_set,
    data_size=None if config.dynamic_size else config.size,
    is_train=True
)

train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=config.batch_size,
    shuffle=True,
    num_workers=4,
    pin_memory=True,
    collate_fn=custom_collate_fn if config.dynamic_size else None
)

# Create model and optimizer
model = BiRefNet(bb_pretrained=True).cuda()
if config.compile:
    model = torch.compile(model)

optimizer = torch.optim.AdamW(
    [p for p in model.parameters() if p.requires_grad],
    lr=config.lr,
    weight_decay=1e-2
)

scheduler = torch.optim.lr_scheduler.MultiStepLR(
    optimizer,
    milestones=[40, 80],
    gamma=0.5
)

# Loss functions
pix_loss_fn = PixLoss()
cls_loss_fn = ClsLoss() if config.auxiliary_classification else None

# Training
num_epochs = 120
mixed_dtype = torch.float16 if config.mixed_precision == 'fp16' else None
autocast_ctx = torch.amp.autocast(device_type='cuda', dtype=mixed_dtype) if mixed_dtype else torch.no_grad().__enter__()

for epoch in range(1, num_epochs + 1):
    model.train()
    loss_meter = AverageMeter()
    
    for batch_idx, (images, masks, class_labels) in enumerate(train_loader):
        images = images.cuda()
        masks = masks.cuda()
        
        # Forward
        with autocast_ctx:
            scaled_preds, class_preds = model(images)
            loss = pix_loss_fn(scaled_preds, masks)
            
            if cls_loss_fn is not None:
                loss = loss + cls_loss_fn(class_preds, class_labels)
        
        # Backward
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        loss_meter.update(loss.item(), images.size(0))
        
        if batch_idx % 50 == 0:
            logger.info(f'Epoch {epoch}, Batch {batch_idx}, Loss: {loss_meter.avg:.6f}')
    
    scheduler.step()
    
    # Save checkpoint
    if epoch % 10 == 0:
        save_checkpoint(
            {'state_dict': model.state_dict()},
            'ckpts/tmp',
            f'BiRefNet-{config.task}-epoch_{epoch}.pth'
        )
        logger.info(f'Checkpoint saved: epoch {epoch}')

logger.close()
```

--------------------------------

### PixLoss Constructor and Configuration

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/loss-functions.md

Initializes the PixLoss class, which aggregates pixel-level losses. Component losses are initialized based on the config.lambdas_pix_last dictionary.

```python
config.lambdas_pix_last = {
    'bce': 30 * 1,          # Binary cross-entropy
    'iou': 0.5 * 1,         # IoU loss
    'iou_patch': 0.5 * 0,   # Patch IoU loss
    'mae': 100 * 1,         # L1/MAE loss
    'mse': 30 * 0,          # Mean squared error
    'triplet': 3 * 0,       # Triplet loss
    'reg': 100 * 0,         # Threshold regularization
    'ssim': 10 * 1,         # SSIM loss
    'cnt': 5 * 0,           # Contour loss
    'structure': 5 * 0,     # Structure loss
}
```

--------------------------------

### Typical PyTorch Training Loop

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/training-api.md

This snippet outlines a standard PyTorch training loop. It includes model training, validation, checkpointing, and learning rate scheduling. Ensure you have your data loaders (train_loader, val_loader), model, optimizer, loss functions, and scheduler defined.

```python
for epoch in range(start_epoch, num_epochs):
    # Training
    model.train()
    for batch_idx, (images, masks, class_labels) in enumerate(train_loader):
        images = images.to(device)
        masks = masks.to(device)
        
        # Forward pass
        with torch.amp.autocast(dtype=torch.bfloat16):
            scaled_preds, class_preds = model(images)
        
        # Compute losses
        pix_loss = pix_loss_fn(scaled_preds, masks)
        cls_loss = cls_loss_fn(class_preds, class_labels) if class_preds else 0
        
        loss = pix_loss + cls_loss
        
        # Backward
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    # Validation
    model.eval()
    with torch.no_grad():
        for images, masks, _ in val_loader:
            # Evaluation
            pass
    
    # Checkpoint
    if epoch % save_freq == 0:
        torch.save(model.state_dict(), f'model_epoch_{epoch}.pth')
    
    # Learning rate decay
    scheduler.step()
```

--------------------------------

### BasicDecBlk Constructor

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/decoder-blocks.md

Defines the parameters for initializing a basic decoder block, including input, output, and intermediate channel dimensions.

```python
def __init__(self, in_channels=64, out_channels=64, inter_channels=64):
```

--------------------------------

### Multi-GPU DDP Training

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/training-api.md

Launch distributed training using PyTorch's DDP. Specify the number of processes per node (GPUs) and enable distributed training.

```bash
python -m torch.distributed.launch --nproc_per_node=2 train.py --dist True --epochs 120
```

--------------------------------

### Train, Test, and Evaluate BiRefNet

Source: https://github.com/zhengpeng7/birefnet/blob/main/README.md

Executes the training, testing, and evaluation pipeline. Specify the run name and the number of GPUs for training and testing.

```shell
# Train & Test & Evaluation
./train_test.sh RUN_NAME GPU_NUMBERS_FOR_TRAINING GPU_NUMBERS_FOR_TEST
# Example: ./train_test.sh tmp-proj 0,1,2,3,4,5,6,7 0
```

--------------------------------

### Configure BF16 Mixed Precision Training

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/training-api.md

Set the mixed precision configuration to 'bf16' and enable compilation. This can be launched using Accelerate with the specified mixed precision flag.

```python
config.mixed_precision = 'bf16'
config.compile = True

# Automatic with Accelerate
accelerate launch --mixed_precision=bf16 train.py --use_accelerate
```

--------------------------------

### Load and Resize Image from Path

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/dataset-utilities.md

This function loads an image from a given file path and optionally resizes it to a specified dimension. It supports loading in RGB or grayscale format. Ensure the 'utils' module is imported.

```python
from utils import path_to_image

image = path_to_image('/path/to/image.jpg', size=(1024, 1024), color_type='rgb')
mask = path_to_image('/path/to/mask.png', size=(1024, 1024), color_type='gray')
```

--------------------------------

### SSIMLoss Constructor

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/loss-functions.md

Initializes the SSIMLoss class. Allows customization of window size and averaging behavior for SSIM computation.

```python
def __init__(self, window_size=11, size_average=True):

```

--------------------------------

### Load Video Paths for Inference

Source: https://github.com/zhengpeng7/birefnet/blob/main/tutorials/BiRefNet_inference_video.ipynb

This snippet loads and sorts video file paths from a specified directory. It filters for common video extensions and sorts them numerically based on their filenames.

```python
import os
from glob import glob

video_src_paths = sorted([f for f in glob('../videos_todo/*') if os.path.splitext(f)[1] in ('.mp4', '.avi')], key=lambda x: int(x.split('/')[-1].split('.')[0]))
print('video_src_paths:', video_src_paths)
```

--------------------------------

### build_backbone Function

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/backbones.md

Factory function for instantiating and loading pretrained backbone networks.

```APIDOC
## build_backbone Function

**Module**: `models.backbones.build_backbone`

Factory function for instantiating and loading pretrained backbone networks.

### Function Signature

```python
def build_backbone(bb_name, pretrained=True, params_settings='')
```

**Parameters**:

| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| bb_name | str | — | Backbone architecture name |
| pretrained | bool | True | Whether to load ImageNet pretrained weights |
| params_settings | str | '' | Additional parameters for backbone initialization (used for models loaded from string eval) |

**Returns**: Instantiated backbone nn.Module

### Supported Backbones

#### CNN-based Backbones

**VGG16**:
- Name: `'vgg16'`
- Output channels: [512, 512, 256, 128]
- Multi-scale features: Yes (4 pyramid levels)
- Pretrained source: torchvision VGG16_Weights.DEFAULT

**VGG16 with Batch Normalization**:
- Name: `'vgg16bn'`
- Output channels: [512, 512, 256, 128]
- Multi-scale features: Yes (4 pyramid levels)
- Pretrained source: torchvision VGG16_BN_Weights.DEFAULT

**ResNet50**:
- Name: `'resnet50'`
- Output channels: [2048, 1024, 512, 256]
- Multi-scale features: Yes (4 pyramid levels)
- Pretrained source: torchvision ResNet50_Weights.DEFAULT

#### Transformer-based Backbones

**Swin Transformer v1 - Tiny**:
- Name: `'swin_v1_t'`
- Output channels: [768, 384, 192, 96]
- Pretrained weights: Auto-downloaded from config.weights['swin_v1_t']

**Swin Transformer v1 - Small**:
- Name: `'swin_v1_s'`
- Output channels: [768, 384, 192, 96]

**Swin Transformer v1 - Base**:
- Name: `'swin_v1_b'`
- Output channels: [1024, 512, 256, 128]

**Swin Transformer v1 - Large**:
- Name: `'swin_v1_l'`
- Output channels: [1536, 768, 384, 192]

**Pyramid Vision Transformer v2 - Base0**:
- Name: `'pvt_v2_b0'`
- Output channels: [256, 160, 64, 32]

**Pyramid Vision Transformer v2 - Base1**:
- Name: `'pvt_v2_b1'`
- Output channels: [512, 320, 128, 64]

**Pyramid Vision Transformer v2 - Base2**:
- Name: `'pvt_v2_b2'`
- Output channels: [512, 320, 128, 64]

**Pyramid Vision Transformer v2 - Base5**:
- Name: `'pvt_v2_b5'`
- Output channels: [512, 320, 128, 64]

#### Vision Foundation Models

**DINO v3 Small**:
- Name: `'dino_v3_s'`
- Output channels: [384, 384, 384, 384]
- Notes: Requires pyramid neck in decoder; frozen backbone

**DINO v3 Small Plus**:
- Name: `'dino_v3_s_plus'`
- Output channels: [384, 384, 384, 384]

**DINO v3 Base**:
- Name: `'dino_v3_b'`
- Output channels: [768, 768, 768, 768]

**DINO v3 Large**:
- Name: `'dino_v3_l'`
- Output channels: [1024, 1024, 1024, 1024]

**DINO v3 H-Plus**:
- Name: `'dino_v3_h_plus'`
- Output channels: [1280, 1280, 1280, 1280]

**DINO v3 7B**:
- Name: `'dino_v3_7b'`
- Output channels: [4096, 4096, 4096, 4096]

### Example Usage

```python
from models.backbones.build_backbone import build_backbone
import torch

# Load Swin-Large with pretrained weights
backbone = build_backbone('swin_v1_l', pretrained=True)

# Load ResNet50 without pretrained weights
backbone = build_backbone('resnet50', pretrained=False)

# Load DINO v3 base
backbone = build_backbone('dino_v3_b', pretrained=True)

# Use in BiRefNet
from models.birefnet import BiRefNet
model = BiRefNet(bb_pretrained=True)
```
```

--------------------------------

### File System Path Configuration

Source: https://github.com/zhengpeng7/birefnet/blob/main/_autodocs/types-and-schemas.md

Defines key directory variables used for organizing data, weights, code, and experiment outputs. These paths are typically configured in a 'config.py' file.

```python
# Configured in config.py
paths = {
    'sys_home_dir': '/path/to/root',           # sys_home_dir setting
    'data_root_dir': f'{sys_home_dir}/datasets/dis',
    'weights_dir': f'{sys_home_dir}/weights/cv',
    'code_dir': f'{sys_home_dir}/codes/dis/BiRefNet',
    'task_dir': f'{data_root_dir}/{task}',
    'dataset_dir': f'{task_dir}/{dataset_name}',
    'image_dir': f'{dataset_dir}/im',
    'gt_dir': f'{dataset_dir}/gt',
    'ckpt_dir': './ckpts/{experiment_name}',
    'pred_dir': './predictions/{method}/{testset}',
}
```
Tokens:
Raw