Code for the paper "PortraitNet: Real-time portrait segmentation network for mobile device".
From left to right: input, segmentation mask, background replacement, background bluring.
- python 3.8
- PyTorch 1.10.1+cu111
- pip install easydict matplotlib tqdm opencv-python scipy pyyaml numpy
-
EG1800 Since several image URL links are invalid in the original EG1800 dataset, we finally use 1447 images for training and 289 images for validation.
-
Supervise-Portrait Supervise-Portrait is a portrait segmentation dataset collected from the public human segmentation dataset Supervise.ly using the same data process as EG1800.
Please modify the input image file path and the background image file path in the demo.py file. Then run the following command:
python demo.py- Download the datasets (EG1800 or Supervise-Portriat). If you want to training at your own dataset, you need to modify data/datasets.py and data/datasets_portraitseg.py.
- Prepare training/testing files, like data/select_data/eg1800_train.txt and data/select_data/eg1800_test.txt.
- Select and modify the parameters in the folder of config.
- Start the training with single gpu:
cd myTrain
python train.py