DOCSEARCHbyAlgolia
Submit Docs

U²-Net (U Square Net)

repository·master·Indexed 27 days ago

https://github.com/xuebinqin/u-2-net

A deep learning model for Salient Object Detection (SOD) using a nested U-structure. It supports tasks such as portrait generation, human segmentation, and extracting salient objects from images. The repository provides scripts for inference and training using both a full model (u2net.pth) and a lightweight version (u2netp.pth), as well as specialized models for human segmentation (u2net_human_seg.pth) and portraits (u2net_portrait.pth).

Tokens
1.1K
Snippets
5
Records
6
Agent score
45%

What's inside U²-Net

  1. Run human segmentation with u2net_human_seg.pth

    master

    To perform human segmentation (e.g., for portrait or body segmentation) using the specialized human segmentation model:

    1. Setup Model: Download u2net_human_seg.pth and place it in ./saved_models/u2net_human_seg/.
    2. Prepare Data: Place your target images in a directory, for example: ./test_data/test_human_images/.
    3. Run Inference: Execute the following command:
      python u2net_human_seg_test.py
    4. Results: The output will be saved in a directory such as ./test_data/u2net_test_human_images_results/.

    Note: Due to the labeling accuracy of the Supervisely Person Dataset, this model may not provide hair-level accuracy, but it is more robust for general human segmentation tasks than the standard DUTS-TR trained model.

    python u2net_human_seg_test.py

  3. Set up U2-Net for Salient Object Detection

    master

    Follow these steps to set up the environment and run inference or training:

    1. Clone the repository:

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

    2. Download pre-trained models:

      • For the full model (u2net.pth, 176.3 MB), place it in ./saved_models/u2net/.
      • For the lightweight model (u2netp.pth, 4.7 MB), place it in ./saved_models/u2netp/.
      • Download links are available via Google Drive or Baidu Pan (as specified in the repository README).

    3. Run processes: Navigate to the U-2-Net directory and execute the following:

      • To train: python u2net_train.py
      • To test/inference: python u2net_test.py

    Note: You can switch between models by changing the model_name variable within u2net_train.py or u2net_test.py to either 'u2net' or 'u2netp'.

    git clone https://github.com/NathanUA/U-2-Net.git
# After downloading models to correct directories:
python u2net_test.py

  4. Generate portraits from your own dataset

    master

    To generate portraits from custom images using the u2net_portrait_demo.py script, follow these steps:

    1. Setup Model: Clone the repo and place u2net_portrait.pth in ./saved_models/u2net_portrait/.
    2. Prepare Images: Place your images in ./test_data/test_portrait_images/your_portrait_im/.
      • Requirement: To ensure high quality, the human head region in the input image should be close to or larger than 512x512, and the background should be relatively clear.
    3. Run Prediction: Execute the following command:
      python u2net_portrait_demo.py
    4. Results: Outputs will be saved in ./test_data/test_portrait_images/your_portrait_results/.

    Note on Implementation: Unlike the standard test script, u2net_portrait_demo.py includes a face detection step. It detects the largest face, then crops, pads, and resizes the Region of Interest (ROI) to 512x512 before feeding it to the network.

    python u2net_portrait_demo.py

  5. Generate portraits using the APDrawingGAN testing set

    master

    To run inference on the standard APDrawingGAN testing set:

    1. Setup Model: Place u2net_portrait.pth in ./saved_models/u2net_portrait/.
    2. Prepare Data: Download the train and test set from APDrawingGAN. Since these images are stitched side-by-side (512x1024), you must split them into two 512x512 images and place them in ./test_data/test_portrait_images/portrait_im/.
    3. Run Inference: Execute the following command:
      python u2net_portrait_test.py
    4. Results: Outputs will be saved in ./test_data/test_portrait_images/portrait_results.
    python u2net_portrait_test.py

  6. Composite original images with generated portraits

    master

    You can fuse an original image with a generated portrait using the u2net_portrait_composite.py script. This allows for compositing different styles by blending the original and the portrait.

    Prerequisites:

    1. Clone the repository:
      git clone https://github.com/NathanUA/U-2-Net.git
    2. Download u2net_portrait.pth and place it in ./saved_models/u2net_portrait/.

    Execution: Run the following command:

    python u2net_portrait_composite.py -s 20 -a 0.5

    Arguments:

    • -s: The sigma of the Gaussian function used for blurring the original image.
    • -a: The alpha weights of the original image used during fusion.