This repo contains code and a pre-trained model for clothes segmentation.
This repo contains training code, inference code and pre-trained model for Cloths Parsing from human portrait. Here clothes are parsed into 3 category: Upper body(red), Lower body(green) and Full body(yellow)
This model works well with any background and almost all poses. For more samples visit samples.md
U2NET : This project uses an amazing U2NET as a deep learning model. Instead of having 1 channel output from u2net for typical salient object detection task it outputs 4 channels each respresting upper body cloth, lower body cloth, fully body cloth and background. Only categorical cross-entropy loss is used for a given version of the checkpoint.
Dataset : U2net is trained on 45k images iMaterialist (Fashion) 2019 at FGVC6 dataset. To reduce complexity, I have clubbed the original 42 categories from dataset labels into 3 categories (upper body, lower body and full body). All images are resized into square ¯\_(ツ)_/¯
768 x 768 px for training. (This experiment was conducted with 768 px but around 384 px will work fine too if one is retraining on another dataset).
train
folder which contains training images and train.csv
which is label csv file in options/base_options.py
python setup_model_weights.py
and it will generate weights after model surgey in prev_checkpoints
folder.options/base_options.py
like checkpoint saving folder, logs folder etc.distributed = False
in options/base_options.py
, for multi gpu set it to True
.python train.py
python -m torch.distributed.launch --nnodes=1 --node_rank=0 --nproc_per_node=4 --use_env train.py
trained_checkpoint
folder.input_images
folderpython infer.py
for inference.output_images