Skip to content

Getting Started Image Classification

geekjr edited this page Aug 17, 2021 · 4 revisions

Here, you will see how to get started wit quickai for image classification. We will be training a neural network to classify images of an Acura Integra Type R 2001 vs. an Acura RL Sedan 2012. These images were taken from the Stanford car dataset found here on Kaggle. Start by downloading the train folder from examples in the repository. Inside of that train folder, there will be two folders, one for each class.

Next. copy that downloaded folder to a place of your choosing, and next to the train folder create a new Python script called image_classification_cars.py. Open up that file and start by importing the ImageClassification class from quickai:

from quickai import ImageClassification

Once you have imported that class, lets call it to train our model:

ImageClassification("vgg16", "./train", "cars", save_ios=True)

In the above line of code we call the ImageClassification with vgg16 as the model_type, ./train as path, and cars as save. model_type is the parameter that specifies what model architecture to use in our model, path is the path that contains our training data, and save is the name of the .h5 file to save the model to. The save_ios parameter is a boolean for if an Apple version(.mlmodel) of the model should be saved.

Now, go ahead and run image_classification_cars.py. After a few seconds, the model training process will start, and once that is done, you will see a new file in your directory called cars.h5. That file contains your model. Congrats, you have just trained your first model using quickai!

Check out the next tutorial to see how to use this model for inference. The .mlmodel can be used in Xcode just like any other .mlmodel.