ultimate-image-classification

A black box module to use on any image classification problem using Tensorflow 2 and Keras. Plug in your data, choose which pre-trained model to finetune, and start training!

Features

Model selection (pre-trained or not). You can choose from the following models:

• VGG16, VGG19
• Densenet121, Densenet169, Densenet201
• Xception
• ResNet50, ResNet101, ResNet152
• InceptionV3
• InceptionResNetV2
• NASNetMobile, NASNetLarge
• MobileNet, MobileNetV2
• EfficientNetB0 to EfficientNetB7

Data Augmentations

Automatic data augmentations to the images like:

• Rotations
• Scaling
• Flipping
• Color augmentations
• Contrast augmentations
• Shear
• Translation

We use cpu workers and queues to prepare the data fetching, normalization, and augmentaions to not bottleneck the gpu during training

Downscaling CNN

An optional CNN to downscale the image taken from this paper.
Use it in case you feel that the normal image resizing takes too much information from the image and want the downscaling to be learnable.

Learning

There are many parameters you can control during the training very easily from the configs file:

• Control train/test split
• Control input image size
• Single label or multi-label classification. The loss and final activation will change automatically.
• Specify the architecture of the fully connected layers at the end
• Specify which layers to train and which to freeze
• Number of layers to be popped from the pre-trained model
• Optimizer type (Adam, SGD, RMSprop, Adagrad, Adadelta, Adamax, Nadam)
• Batch size
• learning rate and learning rate decay policy
• Automatic class weighting in case of imbalanced classes
• Automatic best checkpoint saving
• Specify model path and continue training from a checkpoint

Evaluation

We want to make your life easier when it comes to evaluation so we added the following options:

• Choose output path that will contain the model, configs used, and training logs
• Automatic tensorboard support and logging
• Draw activations on the input images using GradCam
• Get accuracy, precision, recall, and confusion matric for single label classification
• Get AUC-ROC, precision, recall, F1, and automatic best threshold for exact match for multilabel classification

Usage

The project was last tested on a virtual environment of python 3.7, pip 23.2.1, and MacOS

• Download this repo
• Have python 3.6+
• pip install -r requirements.txt
• Have a csv describing the data using path and class columns like:
  
  
  • Note that if it's a multilabel problem the classes should be in the following format: class1$class2$class3...
  • If your data is in a structure where each folder represent a class and the images are inside that folder like:
      images\
        class1\
          img1
          img2
        class2\
          img1
    You can run make_csv_from_folders.py to make the csv automatically
• After having a csv for the data you can split it into training and testing by running split_train_test.py which will generate a csv for training and testing. Of course you can also specify them manually.
• Open configs.py and link to the training and testing csvs, and change any training parameters
• Run train.py
• Link to the saved model in configs.py in 'load_model_path' and run test.py
• Run draw_activations.py to get activation highlighting on the original images. The images will be saved next to the model (optional)
• Run compress weights to get smaller saved model (optional)
• Draw training log by running draw_training_log.py (optional)

Example

Note: The following is a brief overview of an example usecase, for a complete walkthrough please visit this blog.

This is an example usecase on a small dataset of about 450 images. We have 3 classes: gym, pool, and operating room.

• Download images from here
• Extract inside 'data' folder to have 'data/images'
• Run train.py

Results

• These are the results when running on the default configs.py using a pre-trained Densenet121.

• Test set results:

Acknowledgements

• Used Grad-Cam implementation here
• Used & extended on many ideas from this repo.