DEV Community

Cover image for Open Source AI Image Classifier with Automatic Dataset Creator
Emirhan Akdeniz for SerpApi

Posted on • Originally published at github.com

Open Source AI Image Classifier with Automatic Dataset Creator

Serapis AI Image Classifier is a program that allows you to automatically create image datasets using SerpApi's Google Images Scraper API, finetune a ResNet50 model, and classify images using the trained model.


Installation

You can install these dependencies using the following command:

pip install requirements.txt
Enter fullscreen mode Exit fullscreen mode

Usage

You can use Serapis AI Image Classifier in one of the following three modes:


Create a dataset and train a new model from scratch

To create a dataset and train a new model from scratch, you will need to provide a list of labels and an image to use as a reference for the scraping process.
SerpApi API Key is necessary for this mode for the program to automatically scrape images you will use in your database using SerpApi's Google Images Scraper API.

You can register to SerpApi to claim free credits.

python serapis.py --train --labels eagle, bull, lion, man --image-path lionimage.jpg --api-key <SerpApi-API-KEY>
Enter fullscreen mode Exit fullscreen mode

Use old scraped images and train a new model

To use old scraped images and train a new model, you will need to provide a list of labels and specify that you want to use old images with the --use-old-images flag.
You can also put your images in the images/ folder and also add enter them in images/catalogue.csv to manually train models using your own dataset.

python serapis.py --train --labels eagle, bull, lion, man --use-old-images --image-path lionimage.jpg
Enter fullscreen mode Exit fullscreen mode

Use a previously trained model

To use a previously trained model, you will need to provide the path to the trained model and an image to classify.

python serapis.py --model-path models/1.pth --image-path lionimage.jpg
Enter fullscreen mode Exit fullscreen mode

Dialogue Mode

You can also navigate the program by not providing any arguments and using the dialogue mode:

python serapis.py
Enter fullscreen mode Exit fullscreen mode

The Output

Classified Image of a Lion

The output will give you the answer:

The image contains Lion
Enter fullscreen mode Exit fullscreen mode

Optional Arguments:

-h, --help                    Help to Nagigate
--train                       Whether to train a new model
--model-path MODEL_PATH       Pretrained Model path you want to use
--dialogue                    Whether to use dialogue to navigate through the program
--use-old-images              Whether to use old images you have downloaded to train a new model
--api-key API_KEY             SerpApi API Key
--limit LIMIT                 Number of images you want to scrape at most for each label
--labels LABELS [LABELS ...]  Labels you want to use to train a new model
--image-path IMAGE_PATH       Path to the image you want to classify
Enter fullscreen mode Exit fullscreen mode

I am grateful to readers for their time and attention. I hope this repository helps in your tasks, or helps you grasp the possibilities you have when you have SERP data at your disposal.

Top comments (0)