Image classification in machine learning: A full guide (2024)

Have you ever wondered how reverse image search works or how your smartphone can recognize you and your contacts' faces in a photo?

This is all due to image classification, a critical component of machine learning (ML) that’s shaping the way our technology views the world.

An image showing how reverse image search on Google works. On the left, you'll find a woman dressed in a pink button up set. The right has results from Google that mimic this look. — *Reverse image search on Google.*

In this guide, I’ll walk you through the fundamentals of image classification, how it's applied in various industries, the technology that powers it, and what you need to understand before diving in yourself.

By the end of this blog, you’ll gain:

The knowledge to train an image classification model
The skills to deploy that model to your devices
Expert tips and tricks from seasoned computer vision (CV) professionals

Let’s dive in!

Machine learning drives image classification

To understand the basics of image classification, it's essential to recognize that it's rooted in ML.

You may already be using ML to streamline tasks and eliminate manual effort. ML empowers computers to learn patterns from data and make decisions autonomously.

Instead of explicitly coding each step, you feed data into an ML model, which then identifies patterns and predicts outcomes.

*Showing the evolution of search results in Netflix, as ML models were implemented. (source*)

ML is being used all around us—in personalized recommendations on Netflix, spam filtering in your email, and even in predicting traffic on your GPS. But when it comes to training models to understand and work with images, that's where computer vision (CV) comes into play.

At the core of CV is image classification, which serves as the starting point for many more complex tasks, like object detection and image segmentation, within the field.

What is image classification within machine learning?

Image classification uses ML algorithms to analyze the presence of items in an image and categorize them accordingly.

For instance, an image classification model might look at the below picture, recognize patterns (like the shape of pepperoni and pizza), and find that it belongs to the class “pepperoni_pizza,”—without indicating where the pizza is within the image.

A pepperoni pizza being classified as multiple objects using the Viam app. — *A pepperoni pizza being classified as multiple objects using the* *Viam app*.

Unsupervised learning

Unsupervised learning is used to train a dataset with minimal human interference. That means the model analyzes the images from the dataset and looks for patterns and similarities—colors, textures, shapes—without guidance.

The model then clusters the images into groups, without assigning what the common thread is.

In unsupervised learning, if images of a fish tank were presented, the algorithm might cluster similar images into groups based on visual similarities, such as grouping images of different fish types. However, it wouldn’t be able to label these clusters with specific names, like 'goldfish' or 'angelfish,' since it doesn't have predefined categories or labels.

Supervised learning

Supervised learning relies heavily on human input early on, requiring each image in the dataset to be manually labeled with the correct classification or category.

For instance, if you're working with images of a fish tank, a person would determine and assign the specific categories for labeling, such as “fish,” “plants,” or “tank accessories.”

Supervised learning models can categorize images in different ways, such as through single-label classification or multi-label classification.

Single-label classification

If you’re working with simple tasks that require only one label per image, this model type is ideal for you. In single-label classification, you’ll only choose one label for each image, even if the image could appropriately be grouped under multiple categories.

This could be especially useful for specific cases, like a smart pet feeder for cats, where the camera only needs to determine whether food is present in the image to dispense more food.

An image of the training data on Viam’s Data tab—showing images with bowls empty and with food. — *An image of the training data on Viam’s* *Data* *tab—showing images with bowls empty and with food.*

Multi-label classification

If you’re working with more complex tasks that require multiple labels, multi-label classification is the way to go.

Imagine your smart pet feeder stores food for both dogs and cats, requiring it to distinguish between the two. A multi-label classification model would make this possible, allowing the feeder to dispense treats when either or both animals are detected in the frame.

A smart pet dispenser classifying the image as a dog with 98% accuracy on the Viam app. — *A smart pet dispenser classifying the image as a dog with 98% accuracy.*

Did you know that with Viam, you can train single-label classification models and multi-label classification models, in under an hour?

Image classification vs. object detection

While image classification determines the class of a specific object, object detection takes this a step further by locating the object within the frame. This uses bounding boxes—rectangular boxes defined by x and y coordinates—to determine the position and size of an object within an image in CV.

An image showing the difference between object detection and image classification. — *Showing the difference between image classification and object detection within* *Viam’s app*.

An image classification model might detect a cat or dog within an image, but object detection models would be able to find where they are too.

This is especially useful for use cases where it’s necessary to locate the object, such as:

Autonomous navigation
Quality assurance
Face and person recognition
Traffic flow management

*The object detection model showing how it could be used with traffic flow management.*

To learn more about object detection, head to our guide.

What are the steps involved in image classification?

While Viam handles much of the heavy lifting for you with its no-code user experience—removing many of the steps—you might still want to understand the process.

Here’s a quick overview of how it works:

1. Dataset curation

Before you can train your model, you’ll need to have a good amount of data that’s fit for training. This means the data will need to be carefully curated, diverse, and aligned to the categories or classes you want the model to learn.

This data sometimes has to be preprocessed manually, which means you’ll need to:

Resize the images: Make sure all the images are the same size (for example, making every image 256x256 pixels).
Normalize the data: Adjust the representations of the images so they're standardized so the computer can compare them fairly.
Apply augmentation: Create slightly different versions of the images (like flipping, rotating, or zooming in) to give the computer more examples to learn from.

If you train your model on Viam, the platform automatically takes care of this preprocessing for you.

2. Model selection and training the machine learning model

Selecting the right model architecture

After preparing your dataset, choose an appropriate model architecture for image classification, such as a convolutional neural network (CNN).

CNNs are effective for analyzing images because they process them layer by layer, similar to how the human brain recognizes visual patterns. These networks start by identifying simple features like edges and textures, and as they progress through layers, they learn to recognize more complex patterns, ultimately distinguishing between different categories or classes in the images.

*A diagram showing how convolutional neural networks (CNNs) operate, similarly to the primary visual cortex of the brain. (source*)

Training the machine learning model

Once the model architecture is defined, it’s time to train the model.

When you train an image classification model, you’re teaching a computer to recognize and categorize different images. You begin by feeding the model a large collection of images, like showing it pictures of different fish species, and labeling each one.

*A gif showing images being tagged within Viam.*

As the model processes these images, it learns to recognize patterns that distinguish one category from another. It continually adjusts to improve accuracy and is tested with new images to ensure it's truly learning, not just memorizing. Over time, the model becomes better at correctly classifying images, allowing it to confidently identify new ones as "an angelfish" or "a goldfish."

*Showing how the training process works for image classification.*

3. Machine learning model deployment

Deploying an image classification model involves taking the trained model and integrating it into a real-world application, where it can start making predictions on new, unseen images.

Once deployed, the model can instantly classify images as they are uploaded or captured, providing real-time results.

Build an image classification model with Viam

With Viam, building an image classification model is straightforward. You don’t need to do any coding—just follow the intuitive steps below, and you can have your model up and running in less than an hour.

And if you need any more guidance, check out our documentation for a full step-by-step walkthrough.

Step 1 - Create a well-rounded dataset to train on

To begin training your model, you’ll need a solid dataset with at least 10 images—though we recommend using more for better results. With Viam’s Data Management Service, you’re able to:

Gather data from any camera, whether it's your phone, computer, or other models, directly to the Viam app.
Upload the data you’re looking to train.

Tips from an ML expert on creating a curated dataset

I sat down with Tahiya Salam, our in-house ML expert who has a PhD in just that. Some of her top tips when compiling a dataset are:

Ensure diversity: "Get images from a variety of lighting conditions, angles, and distances to ensure diversity within the data,” Tahiya advises. This helps the model learn to recognize objects in different scenarios, making it more robust and accurate in real-world applications.
Include negative examples: “Include images with and without the object you’re looking to classify. This helps the model distinguish the target object from the background and reduces the chances of false positives by teaching it what the object is not.”
Balance your classes: “Make sure that each category or class in your dataset has a roughly equal number of images. An imbalanced dataset can lead the model to favor one class over others, reducing its overall accuracy.”
Match your training images to your intended use case: “Use images that reflect the quality and conditions of your production environment. For example, if you plan to use a low-quality camera in production, train with low-quality images. Similarly, if your model will run all day, capture images in both daylight and nighttime conditions.”
Regularly update your dataset: “As you gather more data or as your project evolves, continue to add new images to your dataset. This keeps your model up-to-date and improves its performance over time.”

Step 2 - Train the image classification model

How to label images for image classification

Once your images are all populated into Viam’s data tab, it’s time to label the images. All you have to do is click on the image you’re looking to tag, type in or select the image tag, and select the dataset name in the dropdown.

Tips from an ML expert on labeling images

Leverage pre-trained models: "If you want to streamline the labeling process, consider starting with a pre-trained model. Viam allows you to bypass some of the manual work, making it easier to get started quickly."
Be consistent with labels: "Consistency is crucial," she emphasizes. "Make sure to use the same label for similar images throughout your dataset. Inconsistent labeling can confuse the model and lead to inaccurate results."
Use diverse examples: "Include a variety of examples for each label," she says. "This means labeling images from different angles, lighting conditions, and contexts to ensure the model learns to recognize the object in any scenario."
Incorporate negative labels: "Don’t forget to label images where the object isn’t present. This helps the model learn what the object is not, reducing false positives."

How to train your model

After creating a dataset, all you have to do is:

Navigate to the “Data" tab.
Click on the dataset you want to train a model from.
Fill in the relevant information: The model’s name, task type (in this case, select either single or multi-label image classification), and associated labels.
Click the “Train model” button on your dataset’s page.

Step 3 - Deploy your image classification model

Now comes the exciting part: watching your model in action on your machine. With Viam, the deployment process is quick and seamless, requiring as little as 5 minutes to complete.

Tips from an ML expert on deploying a model to your machine

Thoroughly test for accuracy: "Before you deploy, it's essential to test your model for accuracy," says Tahiya. "Use a diverse set of images that mimic real-world scenarios to ensure your model performs well across all categories."
Plan for edge cases: "Don't forget to think about edge cases," she suggests. "Consider how your model will handle unexpected or rare inputs. Testing these cases before deployment can help prevent misclassifications and make your model more robust."
Make sure all necessary components and services are configured correctly: “When deploying your model, you’ll be configuring the ML Model Service and an ML Model Vision Service to visualize the predictions your model makes. Triple check they’re both properly set up.”

What is an example of image classification?

Now that you know how to create an image classification model, take a look at the following use cases to spark ideas on how you can apply these models in your everyday projects.

Streaming platforms

Streaming platforms use image classification to enhance user experience and streamline content management. By automatically classifying and tagging visual elements in video thumbnails, movie posters, and other media assets, these platforms can efficiently organize and recommend content to users based on their preferences.

For example, image classification can identify and categorize genres, detect actors, or recognize specific visual themes, enabling personalized content suggestions.

E-commerce

If you’ve ever used an image search feature—uploading a photo to find similar products—or browsed a “similar items” section, you’ve directly benefited from image classification technology.

A shopping sites recommendations for "Styles You May Like." — *Styles you might also like are recommended at the bottom of an e-commerce image for a jacket.*

By automatically categorizing and tagging product images, these platforms see to it that items are organized within relevant categories, making it easier for customers to find exactly what they’re looking for.

Security and surveillance

By automatically identifying and categorizing objects, activities, or individuals within video feeds, image classification enables real-time threat detection and response.

For instance, it can be used to recognize authorized users (allowing for an alarm to be disarmed), identify unauthorized users, or identify weapons.

*A look at an image classification model trained on Viam, used to detect authorized individuals and disarm an alarm. Head to our “DIY home automation projects featuring Raspberry Pi (2024)“ to learn more about it and build one yourself.*

Manufacturing

Image classification is a powerful tool for quality control, making sure that products meet the required standards before reaching the market.

For example, in food manufacturing, image classification models can automatically inspect products such as donuts and cookies on assembly lines, classifying them as either quality-approved or defective.

Social media

Image classification is essential for managing and curating large amounts of user-generated content. Platforms use image classification to automatically identify and tag visual content, which helps fuel features like photo organization, personalized content recommendations, and targeted advertising.

Additionally, image classification is the reason you don't constantly encounter inappropriate content online. It effectively detects images containing violence, nudity, or hate symbols, triggering systems to filter them out, ensuring a safer and more enjoyable user experience.

Wildlife monitoring

Image classification contributes to reducing vessel strikes on wildlife. For example, Viam’s partnership with the Whale and Vessel Safety Taskforce (WAVs), supports the development and implementation of technology and monitoring tools within the marine and boating communities to mitigate these risks.

An example of how image classification can help monitor the ocean for specific species. A camera's image of a whale being shown in the Viam app, with an image classification model classifying it as a whale with a 98% accuracy. — *An example of how image classification can help monitor the ocean for specific species.*

Get started with image classification today

Ready to train and deploy your own image classification model? Give Viam a try—it’s free to use and there’s no-code necessary.

For more insights into other types of computer vision models you can deploy on your devices, be sure to check out “Your object detection guide from a computer vision expert (2024).”

‍

Technical Reviewers: Tahiya Salam, Nick Hehr