Image Classification Model Selector

Abstract

Image classification is a part of computer vision, in which the digital system categorizes the entire image. Deep Learning (DL) models are widely used for image classification. However, creating DL models is resource-intensive and time-consuming, and requires extensive knowledge in the DL domain. Google Teachable Machines (GTM) is a website that outputs a trained model given the dataset, however, GTM uses only the MobileNet model and does not balance the image dataset which affects the model’s accuracy. This paper proposes a tool that automates the steps in building and training an image classification model. Using this tool does not require any extensive knowledge in DL. The tool automates the image data pre-processing steps, model building, model training, and model testing to output the best model for the given image classification dataset based on the test accuracy. The tool is tested on two datasets (each balanced and unbalanced dataset): a custom construction dataset and a Minet dataset. Both datasets are also used to train models using the GTM website. Due to the automated pre-processing steps, the average increase in the accuracy is 14.55% in the construction dataset and 3.91% in the Minet dataset. Comparing to the GTM models, the tool produced model with 8.33% more accuracy on the construction dataset and model with 14.07% more accuracy on the Minet dataset. The models trained by the proposed tool have better accuracy compared to the models obtained using GTM. Thus, using the image classification model selector facilitates the creation of an effective image classification model for the target dataset.