DEV Community

CodeTrade India Pvt. Ltd. for CodeTrade India

Posted on

1

Effortless Object Detection In TensorFlow With Pre-Trained Models

Effortless Object Detection In TensorFlow With Pre-Trained Models-CodeTrade

Object detection is a crucial task in computer vision that involves identifying and locating objects within an image or a video stream. The implementation of object detection has become more accessible than ever before with advancements in deep learning libraries like TensorFlow.

In this blog post, we will walk through the process of performing object detection using a pre-trained model in TensorFlow, complete with code examples. Let’s start.

Explore More: How To Train TensorFlow Object Detection In Google Colab: A Step-by-Step Guide

Steps to Build Object Detection Using Pre-Trained Models in TensorFlow

Before diving into the code, you must set up your environment and prepare your dataset. In this example, we’ll use a pre-trained model called ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8 from TensorFlow’s model zoo, which is trained on the COCO dataset.

1. Data Preparation

First, let’s organize our data into the required directory structure:

import os
import shutil
import glob
import xml.etree.ElementTree as ET
import pandas as pd

# Create necessary directories
os.mkdir('data')

# Unzip your dataset into the 'data' directory
# My dataset is 'Fruit_dataset.zip'
# Replace the path with your dataset's actual path
!unzip /content/drive/MyDrive/Fruit_dataset.zip -d /content/data

# Move image and annotation files to their respective folders
# Adjust paths according to your dataset structure
# This code assumes that your dataset contains both 'jpg' and 'xml' files
# and organizes them into 'annotations_train', 'images_train', 'annotations_test', and 'images_test' folders.
# You may need to adapt this structure to your dataset.
# images & annotations for test data
for dir_name, _, filenames in os.walk('/content/data/test_zip/test'):
    for filename in filenames:
        if filename.endswith('xml'):
            destination_path = '/content/data/test_zip/test/annotations_test'
        elif filename.endswith('jpg'):
            destination_path = '/content/data/test_zip/test/images_test'
        source_path = os.path.join(dir_name, filename)
        try:
            shutil.move(source_path, destination_path)
        except:
            pass

# images & annotations for training data 
for dir_name, _, filenames in os.walk('/content/data/train_zip/train'):
    for filename in filenames:
        if filename.endswith('xml'):
            destination_path = '/content/data/train_zip/train/annotations_train'
        elif filename.endswith('jpg'):
            destination_path = '/content/data/train_zip/train/images_train'
        source_path = os.path.join(dir_name, filename)
        try:
            shutil.move(source_path, destination_path)
        except:
            pass
Enter fullscreen mode Exit fullscreen mode

2. Convert XML Annotations to CSV

To train a model, we need to convert the XML annotation files into a CSV format that TensorFlow can use. We’ll create a function for this purpose:

import glob
import xml.etree.ElementTree as ET
import pandas as pd

def xml_to_csv(path):
    classes_names = []
    xml_list = []

    for xml_file in glob.glob(path + '/*.xml'):
        tree = ET.parse(xml_file)
        root = tree.getroot()
        for member in root.findall('object'):
            classes_names.append(member[0].text)
            value = (root.find('filename').text,
                     int(root.find('size')[0].text),
                     int(root.find('size')[1].text),
                     member[0].text,
                     int(member[4][0].text),
                     int(member[4][1].text),
                     int(member[4][2].text),
                     int(member[4][3].text))
            xml_list.append(value)
    column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax']
    xml_df = pd.DataFrame(xml_list, columns=column_name)
    classes_names = list(set(classes_names))
    classes_names.sort()
    return xml_df, classes_names

# Convert XML annotations to CSV for both training and testing data
for label_path in ['/content/data/train_zip/train/annotations_train', '/content/data/test_zip/test/annotations_test']:
    xml_df, classes = xml_to_csv(label_path)
    xml_df.to_csv(f'{label_path}.csv', index=None)
    print(Successfully converted {label_path} xml to csv.')
Enter fullscreen mode Exit fullscreen mode

This code outputs CSV files summarizing image annotations. Each file details the bounding boxes and corresponding class labels for all objects within an image.

3. Create TFRecord Files

The next step is to convert our data into TFRecords. This format is essential for training TensorFlow object detection models efficiently. We’ll utilize the generate_tfrecord.py script included in the TensorFlow Object Detection API.

#Usage:  
#!python generate_tfrecord.py output.csv output.pbtxt /path/to/images output.tfrecords

# For train.record
!python generate_tfrecord.py /content/data/train_zip/train/annotations_train.csv /content/label_map.pbtxt /content/data/train_zip/train/images_train/ train.record

# For test.record
!python generate_tfrecord.py /content/data/test_zip/test/annotations_test.csv /content/label_map.pbtxt /content/data/test_zip/test/images_test/ test.record
Enter fullscreen mode Exit fullscreen mode

Ensure that you have a label_map.pbtxt file containing class labels and IDs in your working directory or adjust the path accordingly.

Read a complete Article Here:

Effortless Object Detection In TensorFlow With Pre-Trained Models

Image of Timescale

Timescale – the developer's data platform for modern apps, built on PostgreSQL

Timescale Cloud is PostgreSQL optimized for speed, scale, and performance. Over 3 million IoT, AI, crypto, and dev tool apps are powered by Timescale. Try it free today! No credit card required.

Try free

Top comments (0)

Billboard image

The Next Generation Developer Platform

Coherence is the first Platform-as-a-Service you can control. Unlike "black-box" platforms that are opinionated about the infra you can deploy, Coherence is powered by CNC, the open-source IaC framework, which offers limitless customization.

Learn more

👋 Kindness is contagious

Immerse yourself in a wealth of knowledge with this piece, supported by the inclusive DEV Community—every developer, no matter where they are in their journey, is invited to contribute to our collective wisdom.

A simple “thank you” goes a long way—express your gratitude below in the comments!

Gathering insights enriches our journey on DEV and fortifies our community ties. Did you find this article valuable? Taking a moment to thank the author can have a significant impact.

Okay