Trasnformer NN ued for segmentation with additional UNet in the end of Transformer for human segmentation.
Code provided in this directory able to perform training of UNet based on VIT(Vision Transformer) output as well as further inference of whole NN structure.
- NN
- Transformer - Contains all files related to transformer model(structure, utils(loading, inferencing etc.))
- UNet - Contains all files related to the UNet model(structure, utils)
- Dataset
- Utils
- data - directory for your data used for UNet trianing possible to use any (e.g This small dataset)
To train your own UNet model based on VIT please configure /NN/train.py file for image, annotation and transformer model paths. Trasnformer model used in this work found in here . Transformer models checkpoints and variants could be found here