asimniazi63 / pytorch-hair-segmentation

pytorch-hair-segmentation

pytorch-hair-segmentation

Implementation of pytorch semantic segmentation with figaro-1k.

• tutorial document : https://pytorchhair.gitbook.io/project/ (kor)

Prerequisites

opencv-contrib-python 3.4.4
pytorch 0.4.1
torchvision 0.2.1
numpy 1.14.5

Downloading dataset

# specify a directory for dataset to be downloaded into, else default is ./data/
sh data/figaro.sh #<directory>

Running trainer

# sample execution

python3 main.py \
  --networks mobilenet \
  --dataset figaro \
  --data_dir ./data/Figaro1k \
  --scheduler ReduceLROnPlateau \
  --batch_size 4 \
  --epochs 5 \
  --lr 1e-3 \
  --num_workers 2 \
  --optimizer adam \
  --img_size 256 \
  --momentum 0.5 \
  --ignite True

• You should add your own model script in networks and make it avaliable in get_network in ./networks/__init__.py

Running docker & train

with ignite

docker run davinnovation/pytorch-hairsegment:cpu python main.py

with no-ignite

docker run -p davinnovation/pytorch-hairsegment:cpu python main.py --ignite False

Evaluating model

# sample execution

python3 evaluate.py \
  --networks pspnet_resnet101 \
  --ckpt_dir [path to checkpoint] \
  --dataset figaro \
  --data_dir ./data/Figaro1k \
  --save_dir ./overlay/ \
  --use_gpu True

Evaluation result on figaro testset

Model	IoU	F1-score	Checkpoint
pspnet_resnet101	0.92	0.96	link
pspnet_squeezenet	0.88	0.91	link
deeplabv3plus	0.80	0.89	-

Sample visualization

• Red: GT / Blue: Segmentation Map