U2-Net

https://github.com/xuebinqin/u-2-net
Admin
U2-Net is a deep learning model for salient object detection, featuring a nested U-structure for...
Tokens:41,353
Snippets:240
Trust Score:9.2
License:Apache-2.0
Update:1 month ago
Show doc for...
Tokens:
Raw
### Usage Example

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/model-u2net.md

Example of how to instantiate U2NETP, load pretrained weights, and perform inference.

```APIDOC
## Usage Example

### Description
Demonstrates loading the U2NETP model, loading state dict, and running inference on a random input.

### Code
```python
import torch
from model import U2NETP

net = U2NETP(3, 1)  # select 'u2netp' in u2net_test.py to use this
net.load_state_dict(torch.load('saved_models/u2netp/u2netp.pth', map_location='cpu'))
net.eval()

d0, d1, d2, d3, d4, d5, d6 = net(torch.rand(1, 3, 320, 320))
saliency = d0[:, 0, :, :]  # [1, 320, 320], values in (0, 1)
```
```

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

### SalObjDataset Usage Example

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/data-loader.md

Demonstrates initializing the dataset for inference without labels and iterating through a DataLoader.

```python
from torch.utils.data import DataLoader
from torchvision import transforms
from data_loader import SalObjDataset, RescaleT, ToTensorLab

ds = SalObjDataset(
    img_name_list=['a.jpg', 'b.jpg'],
    lbl_name_list=[],                                   # inference: no labels
    transform=transforms.Compose([RescaleT(320), ToTensorLab(flag=0)]),
)
loader = DataLoader(ds, batch_size=1, shuffle=False, num_workers=1)

for batch in loader:
    x = batch['image'].type(torch.FloatTensor)          # [N, 3, 320, 320]
    label = batch['label']                              # [N, 1, 320, 320] (zeros here)
```

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

### Complete portrait pipeline example

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/image-processing-helpers.md

Demonstrates the full workflow from loading the model and detecting a face to generating and saving the portrait map.

```python
import cv2, torch
from model import U2NET
from u2net_portrait_demo import detect_single_face, crop_face, inference

img = cv2.imread('my_photo.jpg')
net = U2NET(3, 1)
net.load_state_dict(torch.load('saved_models/u2net_portrait/u2net_portrait.pth'))
net.eval()

cascade = cv2.CascadeClassifier('saved_models/face_detection_cv2/haarcascade_frontalface_default.xml')
face = detect_single_face(cascade, img)   # None if not found -> whole image used
roi = crop_face(img, face)                # 512x512 BGR patch

portrait_map = inference(net, roi)        # numpy [512, 512] in [0, 1], 1 = portrait
cv2.imwrite('out.png', (portrait_map * 255).astype('uint8'))
```

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

### U2NET Usage Example

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/model-u2net.md

Demonstrates initializing the model, loading pretrained weights, and performing inference on a random input tensor.

```python
import torch
from model import U2NET  # re-exported by model/__init__.py

net = U2NET(3, 1)

# load pretrained SOD weights (176.3 MB)
state = torch.load('saved_models/u2net/u2net.pth', map_location='cpu')
net.load_state_dict(state)
net.eval()

# normalized input, float32, [N, C, H, W]
x = torch.rand(1, 3, 320, 320)

d0, d1, d2, d3, d4, d5, d6 = net(x)
pred = d0[:, 0, :, :]      # fused map; script u2net_test.py uses d1[:,0,:,:]
print(pred.shape)          # torch.Size([1, 320, 320])
```

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

### Compose transforms example with RescaleT and ToTensorLab

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/data-loader.md

Shows how to compose RescaleT and ToTensorLab transforms using torchvision.transforms.Compose. This is the typical setup for creating a transform pipeline for the dataset.

```python
from torchvision import transforms
from data_loader import RescaleT, ToTensorLab, SalObjDataset

transform = transforms.Compose([RescaleT(320), ToTensorLab(flag=0)])
```

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

### Training transformation chain example

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/data-loader.md

Example of a torchvision transformation pipeline used in u2net_train.py.

```python
from torchvision import transforms
from data_loader import RescaleT, RandomCrop, ToTensorLab

train_transform = transforms.Compose([
    RescaleT(320),     # square 320x320
    RandomCrop(288),   # crop to 288x288
    ToTensorLab(flag=0),
])
```

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

### Instantiating U-2-Net RSU Stages

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/model-u2net.md

Example showing how to instantiate the different RSU stages and perform a forward pass with a random input tensor.

```python
import torch
from model.u2net import RSU7, RSU6, RSU5, RSU4

# exact instantiation used by U2NET stage1..stage4
stage1 = RSU7(in_ch=3,  mid_ch=32,  out_ch=64)
stage2 = RSU6(in_ch=64, mid_ch=32,  out_ch=128)
stage3 = RSU5(in_ch=128, mid_ch=64, out_ch=256)
stage4 = RSU4(in_ch=256, mid_ch=128, out_ch=512)

y = stage1(torch.rand(1, 3, 320, 320))  # -> [1, 64, 320, 320]
```

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

### U2NETP usage example

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/model-u2net.md

Shows how to instantiate U2NETP, load its pretrained weights, and run inference on a random input tensor. The output is a tuple of saliency maps, with d0 being the fused output.

```python
import torch
from model import U2NETP

net = U2NETP(3, 1)  # select 'u2netp' in u2net_test.py to use this
net.load_state_dict(torch.load('saved_models/u2netp/u2netp.pth', map_location='cpu'))
net.eval()

d0, d1, d2, d3, d4, d5, d6 = net(torch.rand(1, 3, 320, 320))
saliency = d0[:, 0, :, :]  # [1, 320, 320], values in (0, 1)
```

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

### U2NET Error Handling Example

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/model-u2net.md

Shows how to catch a RuntimeError when loading a state dictionary that does not match the model architecture.

```python
import torch
from model import U2NET

net = U2NET(3, 1)
try:
    net.load_state_dict(torch.load('saved_models/u2netp/u2netp.pth', map_location='cpu'))
except RuntimeError as e:
    print('state dict mismatch:', str(e)[:200])
```

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

### U2NET

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/model-u2net.md

Documentation for the U2NET class, including constructor parameters, forward method, usage example, and error handling.

```APIDOC
## Class `U2NET(nn.Module)`

### Description
Full-size U²-Net (≈173.6 MB weights). Six-stage encoder/decoder composite of the RSU blocks above, with six side outputs and one fused output. Its `forward` returns a **7-tuple**.

### Constructor

`U2NET(in_ch=3, out_ch=1)`

#### Parameters
- **in_ch** (int) - Optional, default `3` - Number of input image channels (3 = RGB)
- **out_ch** (int) - Optional, default `1` - Number of output channels per side map (1 = single saliency/probability map)

### Method: `forward(self, x)`

#### Description
Performs forward pass through the network. `x` is a `torch.Tensor` of shape `[N, in_ch, H, W]`. Returns a 7-tuple of `torch.Tensor`s, each produced by `F.sigmoid(...)` (values in `(0, 1)`).

#### Parameters
- **x** (torch.Tensor) - Required - Input tensor of shape `[N, in_ch, H, W]`

#### Returns
A 7-tuple of `torch.Tensor`s, each of shape `[N, out_ch, H, W]`:
- **d0** - Fused saliency map: `outconv(cat(d1..d6, 1))` — final prediction
- **d1** - Side output of the shallowest decoder stage (`side1(hx1d)`), full input resolution
- **d2** - Side output of decoder stage 2, upsampled to `d1` size
- **d3** - Side output of decoder stage 3, upsampled to `d1` size
- **d4** - Side output of decoder stage 4, upsampled to `d1` size
- **d5** - Side output of decoder stage 5, upsampled to `d1` size
- **d6** - Side output of the encoder bottleneck (`side6(hx6)`), upsampled to `d1` size

### Usage Example
```python
import torch
from model import U2NET  # re-exported by model/__init__.py

net = U2NET(3, 1)

# load pretrained SOD weights (176.3 MB)
state = torch.load('saved_models/u2net/u2net.pth', map_location='cpu')
net.load_state_dict(state)
net.eval()

# normalized input, float32, [N, C, H, W]
x = torch.rand(1, 3, 320, 320)

d0, d1, d2, d3, d4, d5, d6 = net(x)
pred = d0[:, 0, :, :]      # fused map; script u2net_test.py uses d1[:,0,:,:]
print(pred.shape)          # torch.Size([1, 320, 320])
```

### Error Handling Example
`load_state_dict` raises `torch.nn.modules.module.RuntimeError` (an ordinary `RuntimeError`) when the state-dict keys do not match the architecture — e.g. loading the small model's weights (`u2netp.pth`) into `U2NET`, or loading CUDA checkpoints without `map_location`:

```python
import torch
from model import U2NET

net = U2NET(3, 1)
try:
    net.load_state_dict(torch.load('saved_models/u2netp/u2netp.pth', map_location='cpu'))
except RuntimeError as e:
    print('state dict mismatch:', str(e)[:200])
```
```

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

### REBNCONV Usage Example

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/model-u2net.md

Demonstrates initializing the REBNCONV block and applying it to a random tensor, including a dilated variant.

```python
import torch
from model.u2net import REBNCONV

conv = REBNCONV(in_ch=3, out_ch=3, dirate=1)
out = conv(torch.rand(1, 3, 64, 64))  # -> [1, 3, 64, 64]

# dilated variant used by RSU4F: 3x3 conv with dilation 8
conv_d8 = REBNCONV(in_ch=64, out_ch=64, dirate=8)
```

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

### U2NETP class definition (partial)

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/model-u2net.md

Shows the beginning of the U2NETP class definition, including the constructor and the first stage setup. The full class is in model/u2net.py at line 423.

```python
class U2NETP(nn.Module):

    def __init__(self,in_ch=3,out_ch=1):
        super(U2NETP,self).__init__()

        self.stage1 = RSU7(in_ch,16,64)
        # ...
```

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

### Error handling example: verify output ordering

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/model-u2net-refactor.md

Error handling example: because forward returns a list (not a tuple), unpacking with the classic 7-arg call of model/u2net.py silently swaps d1/d6 without raising an exception. This snippet verifies output ordering explicitly by asserting the list length and that all maps are at input resolution.

```python
import torch
from model.u2net_refactor import U2NET_full

net = U2NET_full().eval()
maps = net(torch.rand(1, 3, 320, 320))
assert len(maps) == 7
assert maps[0].shape[-2:] == maps[1].shape[-2:]  # all maps at input resolution
```

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

### muti_bce_loss_fusion usage example

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/image-processing-helpers.md

Demonstrates how to invoke the loss function with seven sigmoid maps and ground-truth labels, followed by a backward pass.

```python
import torch
from u2net_train import muti_bce_loss_fusion  # WARNING: importing u2net_train starts training

maps = [torch.sigmoid(torch.randn(4, 1, 288, 288)) for _ in range(7)]
labels = torch.randint(0, 2, (4, 1, 288, 288)).float()

loss0, loss = muti_bce_loss_fusion(*maps, labels_v=labels)
loss.backward()
```

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

### Argparse setup for u2net_portrait_composite.py

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/cli-scripts.md

Defines the -s and -a command-line flags for sigma and alpha. Note that no parser defaults are set, so omitting a flag results in None and float(None) raises a TypeError.

```python
parser = argparse.ArgumentParser(description="image and portrait composite")
parser.add_argument('-s',action='store',dest='sigma')
parser.add_argument('-a',action='store',dest='alpha')
```

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

### Example usage of normPRED

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/image-processing-helpers.md

Shows how to import and use normPRED from u2net_test. It creates a random tensor, applies sigmoid, then normalizes it, asserting that the result has min 0 and max 1.

```python
import torch
from u2net_test import normPRED

pred = torch.sigmoid(torch.randn(1, 1, 320, 320))
dn = normPRED(pred)          # values stretched to [0, 1]
assert dn.max().item() == 1.0 and dn.min().item() == 0.0
```

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

### Run training or inference

Source: https://github.com/xuebinqin/u-2-net/blob/master/README.md

Commands to initiate the training or testing process. Ensure you are in the U-2-Net directory before running these.

```bash
python u2net_train.py
```

```bash
python u2net_test.py
```

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

### ToTensorLab usage in test pipeline

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/data-loader.md

Example of integrating ToTensorLab into a torchvision Compose pipeline, as used in u2net_test.py.

```python
from torchvision import transforms
from data_loader import RescaleT, ToTensorLab

# test pipeline used by u2net_test.py
test_transform = transforms.Compose([RescaleT(320), ToTensorLab(flag=0)])
```

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

### normPRED(d)

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/image-processing-helpers.md

Documentation for the normPRED helper function, including its signature, parameters, return value, and usage example.

```APIDOC
## normPRED(d)

### Description
Min-max normalization of a prediction map. Linearly scales the input tensor so that its minimum becomes 0 and its maximum becomes 1. Used to stretch sigmoid probability maps to the full [0, 1] range before saving.

### Signature
```python
def normPRED(d):
    ma = torch.max(d)
    mi = torch.min(d)
    dn = (d-mi)/(ma-mi)
    return dn
```

### Parameters
- **d** (torch.Tensor) - Required - Raw prediction tensor (any shape)

### Returns
- **dn** (torch.Tensor) - `d` linearly scaled so that its min is 0 and its max is 1 (`(d - mi) / (ma - mi)`)

### Example
```python
import torch
from u2net_test import normPRED

pred = torch.sigmoid(torch.randn(1, 1, 320, 320))
dn = normPRED(pred)          # values stretched to [0, 1]
assert dn.max().item() == 1.0 and dn.min().item() == 0.0
```

### Source Locations
- `u2net_test.py:25`
- `u2net_portrait_test.py:25`
- `u2net_human_seg_test.py:24`
- `u2net_portrait_demo.py:91`
- `u2net_portrait_composite.py:28`
```

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

### Launch Gradio Web Demo with PaddleHub

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/cli-scripts.md

Executes the web interface using the PaddleHub backend. Requires paddlepaddle and paddlehub packages.

```python
model = hub.Module(name='U2Net')

def infer(img):
  result = model.Segmentation(
      images=[cv2.imread(img.name)],
      paths=None,
      batch_size=1,
      input_size=320,
      output_dir='output',
      visualization=True)
  return result[0]['front'][:,:,::-1], result[0]['mask']

inputs = gr.inputs.Image(type='file', label="Original Image")
outputs = [gr.outputs.Image(type="numpy", label="Front"), gr.outputs.Image(type="numpy", label="Mask")]
gr.Interface(infer, inputs, outputs, title=title, description=description, article=article, examples=examples).launch()
```

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

### Example Usage: Variant A

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/image-processing-helpers.md

Demonstrates how to invoke the standard segmentation save_output function using a normalized prediction tensor.

```python
import torch
from u2net_test import save_output, normPRED

pred = normPRED(torch.rand(1, 1, 320, 320))
save_output('test_data/test_images/img01.jpg', pred, 'test_data/u2net_results/')
```

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

### SalObjDataset.__init__

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/data-loader.md

Constructor for SalObjDataset. Initializes the dataset with image and label file path lists and an optional transform.

```APIDOC
## SalObjDataset.__init__

### Description
Initializes the dataset with image and label file path lists and an optional transform.

### Parameters
- **img_name_list** (list[str]) - Required - Image file paths (read with skimage.io.imread)
- **lbl_name_list** (list[str]) - Required (may be []) - Label file paths; when empty, zero-value labels are generated. Pairing is positional with img_name_list
- **transform** (callable) - Optional - A transform (or transforms.Compose chain) applied to the sample dict

### Request Example
```python
SalObjDataset(
    img_name_list=['a.jpg', 'b.jpg'],
    lbl_name_list=[],
    transform=transforms.Compose([RescaleT(320), ToTensorLab(flag=0)])
)
```
```

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

### U2NETP.__init__

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/model-u2net.md

Constructor for the U2NETP class. Initializes the network with specified input and output channels.

```APIDOC
## U2NETP.__init__

### Description
Initializes the U2NETP model with the given input and output channel counts.

### Method
Constructor

### Signature
`__init__(self, in_ch=3, out_ch=1)`

### Parameters
- **in_ch** (int) - Optional, default `3` - Number of input channels.
- **out_ch** (int) - Optional, default `1` - Number of output channels per side map.

### Returns
None
```

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

### U2NET.__init__

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/model-u2net-refactor.md

Initializes the U2NET network with a configuration dictionary and output channel count. Builds the network layers based on the provided configuration.

```APIDOC
## U2NET.__init__

### Description
Initializes the U2NET network with a configuration dictionary and output channel count. Builds the network layers based on the provided configuration.

### Method
Constructor

### Signature
```python
U2NET(cfgs, out_ch)
```

### Parameters
- `cfgs` (dict[str, list]) - Required - Stage configuration. Each key defines a stage in order. Each value is a list: `[name, (height, in_ch, mid_ch, out_ch, dilated?), side]`. `side` ≤ 0 means no side head; `side` > 0 is the side head input channel count and the head is registered as `side<last-char-of-name>`. The `(..., dilated)` element is optional.
- `out_ch` (int) - Required - Output channels of every side head and of the fused map. The fusion conv is built as `nn.Conv2d(height * out_ch, out_ch, 1)` where `height = (len(cfgs) + 1) // 2`.

### Returns
None
```

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

### U2NET_lite()

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/model-u2net-refactor.md

Builds the lightweight U²-Net model with out_ch=1 fixed. Returns a U2NET instance.

```APIDOC
## Function U2NET_lite()

### Description
Builds the lightweight U²-Net (equivalent architecture to model/u2net.py:U2NETP) with out_ch=1 fixed. Returns a U2NET instance.

### Signature
```python
def U2NET_lite():
    lite = {
        # cfgs for building RSUs and sides
        # {stage : [name, (height(L), in_ch, mid_ch, out_ch, dilated), side]}
        'stage1': ['En_1', (7, 3, 16, 64), -1],
        'stage2': ['En_2', (6, 64, 16, 64), -1],
        'stage3': ['En_3', (5, 64, 16, 64), -1],
        'stage4': ['En_4', (4, 64, 16, 64), -1],
        'stage5': ['En_5', (4, 64, 16, 64, True), -1],
        'stage6': ['En_6', (4, 64, 16, 64, True), 64],
        'stage5d': ['De_5', (4, 128, 16, 64, True), 64],
        'stage4d': ['De_4', (4, 128, 16, 64), 64],
        'stage3d': ['De_3', (5, 128, 16, 64), 64],
        'stage2d': ['De_2', (6, 128, 16, 64), 64],
        'stage1d': ['De_1', (7, 128, 16, 64), 64],
    }
    return U2NET(cfgs=lite, out_ch=1)
```

### Parameters
None.

### Returns
- **U2NET** - A U2NET model instance with out_ch=1.

### Source
`model/u2net_refactor.py:152`.
```

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

### Download Pretrained Model Weights

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/cli-scripts.md

Downloads u2net and u2net_portrait weights from Google Drive to the local saved_models directory. Requires the gdown package.

```python
os.makedirs('./saved_models/u2net', exist_ok=True)
os.makedirs('./saved_models/u2net_portrait', exist_ok=True)

gdown.download('https://drive.google.com/uc?id=1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ',
    './saved_models/u2net/u2net.pth', quiet=False)
gdown.download('https://drive.google.com/uc?id=1IG3HdpcRiDoWNookbncQjeaPN28t90yW',
    './saved_models/u2net_portrait/u2net_portrait.pth', quiet=False)
```

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

### Clone U-2-Net Repository

Source: https://github.com/xuebinqin/u-2-net/blob/master/README.md

Use this command to download the U-2-Net repository from GitHub.

```bash
git clone https://github.com/NathanUA/U-2-Net.git
```

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

### Example Usage: Variant B with Error Handling

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/image-processing-helpers.md

Demonstrates how to invoke the portrait composite save_output function, including a try-except block to handle potential errors like image read failures.

```python
import torch
from u2net_portrait_composite import save_output

pred = torch.rand(1, 1, 512, 512)
try:
    save_output('p.jpg', pred, 'out/', sigma=20, alpha=0.5)
except ValueError as e:
    print('blur/fuse failed:', e)   # e.g. image read errors surface here
```

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

### Usage example for U2NET_full and U2NET_lite

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/model-u2net-refactor.md

Basic usage of U2NET_full and U2NET_lite from model.u2net_refactor. The forward pass returns a list of 7 tensors; unpacking order is reversed compared to model/u2net.py: fused, d6, d5, d4, d3, d2, d1. The fused tensor is the final prediction and d1 is the shallowest side map. The output shape is [1, 1, 320, 320] for a 320x320 input.

```python
import torch
from model.u2net_refactor import U2NET_full, U2NET_lite

net = U2NET_full()                       # out_ch fixed to 1
net.eval()

maps = net(torch.rand(1, 3, 320, 320))   # list of 7 tensors
fused, d6, d5, d4, d3, d2, d1 = maps     # NOTE: reversed vs. model/u2net.py

# fused is the final prediction; d1 is the shallowest side map
print(fused.shape)                       # torch.Size([1, 1, 320, 320])

lite = U2NET_lite()                      # lightweight variant
```

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

### U2NET_lite() - Build lightweight U²-Net

Source: https://github.com/xuebinqin/u-2-net/blob/master/_autodocs/api-reference/model-u2net-refactor.md

Builds the lightweight U²-Net with out_ch=1 fixed, equivalent to model/u2net.py:U2NETP. Returns a U2NET instance.

```python
def U2NET_lite():
    lite = {
        # cfgs for building RSUs and sides
        # {stage : [name, (height(L), in_ch, mid_ch, out_ch, dilated), side]}
        'stage1': ['En_1', (7, 3, 16, 64), -1],
        'stage2': ['En_2', (6, 64, 16, 64), -1],
        'stage3': ['En_3', (5, 64, 16, 64), -1],
        'stage4': ['En_4', (4, 64, 16, 64), -1],
        'stage5': ['En_5', (4, 64, 16, 64, True), -1],
        'stage6': ['En_6', (4, 64, 16, 64, True), 64],
        'stage5d': ['De_5', (4, 128, 16, 64, True), 64],
        'stage4d': ['De_4', (4, 128, 16, 64), 64],
        'stage3d': ['De_3', (5, 128, 16, 64), 64],
        'stage2d': ['De_2', (6, 128, 16, 64), 64],
        'stage1d': ['De_1', (7, 128, 16, 64), 64],
    }
    return U2NET(cfgs=lite, out_ch=1)
```

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

### Run Portrait Composition

Source: https://github.com/xuebinqin/u-2-net/blob/master/README.md

Execute the portrait composition script after placing the model file in ./saved_models/u2net_portrait/. The -s flag sets the Gaussian sigma for blurring, and -a sets the alpha weight for the original image.

```bash
python u2net_portrait_composite.py -s 20 -a 0.5
```
Tokens:
Raw