Skip to content

General image classification framework implemented by Pytorch for all image classification task.

License

Notifications You must be signed in to change notification settings

Daisy-Zhang/Image-Classification-Pytorch

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

12 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Pytorch Image Classification

General image classification framework implemented by Pytorch for almost all image classification task.

structure

/checkpoints

This directory will store all models you trained.

/data

Please put all your training or test data in this directory and follow the original directory structure. I provide two kinds of dataset format: Custom and ImageFolder. Users can choose one of them for your convinence. For example, if you want to use ImageFolder, you can put all your training images in /data/ImageFolder/train and your test images in /data/ImageFolder/test. The further child directory in /data/ImageFolder/train is the classes folder of your images, such as /data/ImageFolder/train/ClassA...

/log

All log record file will be stored in this directory.

/models

You can put all the network model you design in this directory. I already provide three classic networks: VGG16, ResNet18,ResNet50, GoogLeNet.

env

You could use following command to install all dependencies:

pip -r requirements.txt

PS: for the pytorch version, early version may still be available.

data

I provide two kinds of dataloader in dataset.py: ImageFolder dataloader and Custom dataloader. Users can choose any one of them. All you need is change following part in train.py and test.py:

train_loader = get_imagefolder_train_loader()
#train_loader = get_custom_train_loader()
print('get train loader done')
val_loader = get_imagefoler_val_loader()
#val_loader = get_custom_val_loader()
print('get val loader done')

And do not forget put your images in /data directory following original structure I provide.

config

Users can change the config setting in conf.py as they need, such as IMAGE_SIZE, EPOCH et. al.

train

Users can use following command to start training model:

python train.py -model=vgg16 -gpu
  • -model: choose one model in /models.
  • -gpu: use gpu to train model.

PS: Users can also use following command to get usage:

python train.py -h

test

Users can use following command to evaluate model:

python -model=vgg16 -weights=YOUR_WEIGHT_PATH -gpu -data_path=YOUR_TEST_DATA_PATH
  • -model: choose the model structure of your trained model.
  • -weights: the model weights path.
  • -gpu: use gpu to test model.
  • -data_path: the test data path. Users can modify the my_eval() function in utils.py to better test your own test data.

PS: Users can also use following command to get usage:

python test.py -h

others

utils.py: some utils function used in train.py and test.py. Users can modify this file for their convinence.

If this repo do you a favor, a star is my pleasure :)

And if you find any problem, please contact me or open an issue.

About

General image classification framework implemented by Pytorch for all image classification task.

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages