A repository for everything Vision Transformers.
-
Image Classification
- ViT Base Patch 16 | 224x224: Torchvision pretrained weights
- ViT Base Patch 32 | 224x224: Torchvision pretrained weights
- ViT Tiny Patch 16 | 224x224: Timm pretrained weights
- Vit Tiny Patch 16 | 384x384: Timm pretrained weights
- Swin Transformer Tiny Patch 4 Window 7 | 224x224: Official Microsoft weights
- Swin Transformer Small Patch 4 Window 7 | 224x224: Official Microsoft weights
- Swin Transformer Base Patch 4 Window 7 | 224x224: Official Microsoft weights
- Swin Transformer Large Patch 4 Window 7 | 224x224: No pretrained weights
- MobileViT S
- MobileViT XS
- MobileVit XXS
-
Object Detection
- DETR ResNet50 (COCO pretrained)
- DETR ResNet50 DC5 (COCO pretrained)
- DETR ResNet101 (COCO pretrained)
- DETR ResNet101 DC5 (COCO pretrained)
- Quick Setup
- Importing Models and Usage
- DETR Video Inference Commands (COCO pretrained models)
- Examples
pip install vision-transformers
git clone https://github.com/sovit-123/vision_transformers.git
cd vision_transformersInstallation in the environment of your choice:
- Install requirements on RTX 30/40 (Ampere and Ada Lovelace) series and T4/P100 GPUs.
pip install -r requirements.txt
- OR install requirements for RTX 50 series and Blackwell GPUs. First install PyTorch >= 2.8 with CUDA >= 12.9
pip install -r requirements_blackwell.txt
Finally, install the package.
pip install .Replace num_classes=1000 with you own number of classes.
from vision_transformers.models import vit model = vit.vit_b_p16_224(num_classes=1000, pretrained=True) # model = vit.vit_b_p32_224(num_classes=1000, pretrained=True) # model = vit.vit_ti_p16_224(num_classes=1000, pretrained=True)
from vision_transformers.models import swin_transformer model = swin_transformer.swin_t_p4_w7_224(num_classes=1000, pretrained=True) # model = swin_transformer.swin_s_p4_w7_224(num_classes=1000, pretrained=True) # model = swin_transformer.swin_b_p4_w7_224(num_classes=1000, pretrained=True) # model = swin_transformer.swin_l_p4_w7_224(num_classes=1000)
- Clone the repository, install the requirements, and the library from source.
From the vision_transformers directory:
- If you have no validation split
python tools/train_classifier.py --data data/diabetic_retinopathy/colored_images/ 0.15 --epochs 5 --model vit_ti_p16_224
-
In the above command:
-
data/diabetic_retinopathy/colored_images/represents the data folder where the images will be inside the respective class folders -
0.15represents the validation split as the dataset does not contain a validation folder
-
-
If you have validation split
python tools/train_classifier.py --train-dir data/plant_disease_recognition/train/ --valid-dir data/plant_disease_recognition/valid/ --epochs 5 --model vit_ti_p16_224
- In the above command:
--train-dirshould be path to the training directory where the images will be inside their respective class folders.--valid-dirshould be path to the validation directory where the images will be inside their respective class folders.
vit_b_p32_224
vit_ti_p16_224
vit_ti_p16_384
vit_b_p16_224
swin_b_p4_w7_224
swin_t_p4_w7_224
swin_s_p4_w7_224
swin_l_p4_w7_224
mobilevit_s
mobilevit_xs
mobilevit_xxs
- The datasets annotations should be in XML format. The dataset (according to
--dataflag) given in following can be found here => https://www.kaggle.com/datasets/sovitrath/aquarium-data
python tools/train_detector.py --model detr_resnet50 --epochs 2 --data data/aquarium.yaml
Replace weights and input file path as per your requirement.
python tools/inference_image_detect.py --weights runs/training/res_1/best_model.pth --input image.jpg
You can also provide the path to a directory to run inference on all images in that directory.
python tools/inference_image_detect.py --weights runs/training/res_1/best_model.pth --input image_directory
Replace weights and input file path as per your requirement. You can add --show to the command to visualize the detection on screen.
python tools/inference_video_detect.py --weights runs/training/res_1/best_model.pth --input video.mp4
All commands to be executed from the root project directory (vision_transformers)
python tools/inference_video_detect.py --model detr_resnet50 --show --input example_test_data/video_1.mp4 detr_resnet50_dc5 <path/to/your/file> detr_resnet101 detr_resnet101_dc5
# Track all COCO classes. python tools/inference_video_detect.py --track --model detr_resnet50 --show --input example_test_data/video_1.mp4 detr_resnet50_dc5 <path/to/your/file> detr_resnet101 detr_resnet101_dc5 # Track only person class (for DETR, object indices start from 2 for COCO pretrained models). Check `data/test_video_config.yaml` for more information. python tools/inference_video_detect.py --track --model detr_resnet50 --show --input ../inference_data/video_4.mp4 --classes 2 # Track person and motocycle classes (for DETR, object indices start from 2 for COCO pretrained models). Check `data/test_video_config.yaml` for more information. python tools/inference_video_detect.py --track --model detr_resnet50 --show --input ../inference_data/video_4.mp4 --classes 2 5
Just provide the path to the trained weights instead of a model.
python tools/inference_video_detect.py --track --weights runs/training/res_1/best_model.pth --show --input ../inference_data/video_4.mp4