Try-On Module

Building a Character Preserving Virtual Try-On System using deep learning. The project is initiated as college project work.

A web app is developed where users can give image input to the model using a web browser and FastAPI connects this model to frontend.

Resources

Enviroment Setup

git clone https://github.com/rosxnb/viton.git
cd viton
python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

Model Generated Data

The data is uploaded on google drive. You may download and view the data with tensorboardX library using command tensorboard --logdir tensorboard/train_tom_200K

Note:

The result directory contains warp-cloth and warp-mask generated by GMM model in training and testing phase which are required to be in path ./data/train/ for training of TOM model and for testing TOM model, they need to be in ./data/test/

To run model on CPU, skip --use_cuda or pass --no-use_cuda flag when running files train.py and test.py

GMM Training

Trained on GPU NVIDIA GeForce RTX 3050 with 4GB RAM and CUDA Version: 12.2. Training took around 4hrs for 200K steps.

Train command:

python train.py --stage GMM --name train_gmm_200K --workers 4 --save_count 5000 --shuffle --use_cuda

Test command:

python test.py --stage GMM --name test_gmm_200K --workers 4 --datamode test --data_list test_pairs.txt --checkpoint checkpoints/train_gmm_200K/gmm_final.pth --use_cuda

TOM Training

Trained on GPU NVIDIA GeForce RTX 3050 with 4GB RAM and CUDA Version: 12.2. Training took around 13hrs for 200K steps.

Generate wrap data for TOM train module:

python test.py --stage GMM --name generate_gmm_200K --workers 6 --datamode train --data_list train_pairs.txt --checkpoint checkpoints/train_gmm_200K/gmm_final.pth --use_cuda

Train command:

python train.py --stage TOM --name train_tom_200K --workers 6 --save_count 50000 --shuffle --use_cuda

Make sure to copy test data generated by GMM module todata/test/ folder before running test command on TOM.

Test command:

python test.py --stage TOM --name test_tom_200K --workers 6 --datamode test --data_list test_pairs.txt --checkpoint checkpoints/train_tom_200K/tom_final.pth --use_cuda

Generated Single Output

python infer.py

Note: The generated images are shown as matplotlib plots and is not saved locally.

References

The code in this repo is done with reference to following papers:

@inproceedings{wang2018toward,
	title={Toward Characteristic-Preserving Image-based Virtual Try-On Network},
	author={Wang, Bochao and Zheng, Huabin and Liang, Xiaodan and Chen, Yimin and Lin, Liang},
	booktitle={Proceedings of the European Conference on Computer Vision (ECCV)},
	pages={589--604},
	year={2018}
}