travistangvh / TorchScene

Scene recognition tool based on pytorch. Provide training, test and deployment functions, as well as many pretrained models.

TorchScene

A scene recognition tool based on pytorch

Installation

conda env create -f environment.yml python=3.9


conda activate torch-scene
pip install tensorboard==2.7.0
conda install pytorch==1.11.0 torchvision==0.12.0 -c pytorch
conda install -y -c conda-forge tensorboard
export PYTHONPATH=$PYTHONPATH:$(pwd)

Model Zoo (Pretrained Models)

Please refer [Model Zoo]

Step by step

Please download the data from [Place2 Data]

1. Download the data

These images are 256x256 images, in a more friendly directory structure that in train and val split the images are organized such as train/reception/00003724.jpg and val/raft/000050000.jpg

sh download_data_pytorch.sh

2. Train the model

For train configs, please refer to conf/train.yaml

python tools/train.py batch_size=128

3. carry out the inferring

Donwload pretrained weights from [Google Drive], or use the weights you trained in local. For infer configs, please refer to conf/infer.yaml

python tools/infer.py img_path=root/TorchScene/imgs/12.jpg weight_path=root/TorchScene/checkpoints/vision_transformer1654506720.0398195.ckpt

4. Convert a model to TorchScript

python scripts/convert_torchscript.py

Acknowledge

The dataset and basic code comes from [MIT Place365]

Thanks for the great work!