Skip to content

amrit-das/custom_image_classifier_pytorch

Repository files navigation

PyTorch Based Custom Image Classifier

Creating machine learning models, for Image Clasification, built with the help of PyTorch Framework. The API can be used for training models based on custom datasets. Its a ready to deploy platform to provide custom image training sets as input to a ResNet18 based transfer learning approach. Contributions are most Welcome as this repository is still under building. We are trying to add more and more features and make the work versatile.

Dataset : Contains the training and testing datasets
models : Contains the trained models/checkpoints
Predict Image : This folder is used to store the images/videoes to be predicted/segmented
train.py : To train the data
predict.py : To predict from the trained model or segregate the images into separate folders
hyper_params.json : The json file contains all the hyper patameters and related notations.

Table of Contents

       

Installation

  • PyTorch can be installed from the website pytorch.org

    *If installing for ARM devices, follow the Link

  • Example dataset can be downloaded from this Link

  • Clone the repository and Install requirements

sudo apt install git
git clone https://github.com/amrit-das/custom_image_classifier_pytorch.git
cd custom_image_classifier_pytorch
sudo pip install -r requirements.txt

Usage

You can check the detailed usage @ https://medium.com/hardware-interfacing/custom-image-classifier-using-transfer-learning-in-pytorch-framework-c2f7f5155239

Hyperparameters

Before training the network, you can tune the network to your hyper parameters using the 'hyper_params.json' script. But it is advisable not to change any parameters that you may not be aware of, doing so might mess up the network.

Training

Once modified the hyperparameters, put the training data in Dataset/train and testing data in Dataset/val and run the following code

python train.py

Predicting

In order to predict from trained model, place your image to be predicted in Predict_Image and run:

python predict.py

Segregating

In order to segregate a set of images into different folders, place the images in Predict_Image folder and run:

python predict.py --seg

Version

PyTorch - 0.4.1

NumPy - 1.15.4

OpenCV - 3.4.4

Cuda - Optional

Reference

PyTorch Tutorials - ( https://pytorch.org/tutorials/ )

Medium Blogs - ( https://medium.com/hardware-interfacing/ )