Effortless Object Detection In TensorFlow With Pre-Trained Models

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

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.')

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

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