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.
-
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
You can check the detailed usage @ https://medium.com/hardware-interfacing/custom-image-classifier-using-transfer-learning-in-pytorch-framework-c2f7f5155239
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.
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
In order to predict from trained model, place your image to be predicted in Predict_Image and run:
python predict.py
In order to segregate a set of images into different folders, place the images in Predict_Image folder and run:
python predict.py --seg
PyTorch - 0.4.1
NumPy - 1.15.4
OpenCV - 3.4.4
Cuda - Optional
PyTorch Tutorials - ( https://pytorch.org/tutorials/ )
Medium Blogs - ( https://medium.com/hardware-interfacing/ )