Как создать свой собственный Object Detector

Создание кастомного распознания объектов с использованием YOLOv3 с ImageAI

Цель данной статьи помочь начинающему девелоперу в разработке своей системы обнаружения объектов, направляя на все ключевые моменты для обучения успешной модели.

Для примера будем тренировать модель с помощью ImageAI, а также открытой Python библиотеки которая разрешит Вам натренировать модель в пару строк кода.

Вступление

В данной статье мы рассмотрим все необходимые шаги в создании детектора изображений, от получения датасета до тестирования нашей модели.

Вот наши шаги:

  • Сбор данных
  • Разметка датасета
  • Увеличение данных
  • Тренировка модели
  • Оценка модели

Сбор данных

Перед тем как мы начнем тренировать нашу модель мы должны собрать как можно больше данных, в нашем случае это изображения. Минимум должно быть 200 изображений в зависимости от сложности нашего детектора.

Если мы хотим что б наша модель была максимально надежна, наши изображения должны отличатся друг от друга освещением и фоном, что бы наша модель могла делать лучше выводы.

Эти изображения могут быть сделаны на ваш телефон или скачены с интернета.

Тут вы можете увидеть мой под-датасет распознания пилюль.

После получения дадасета мы должны обрезать наши изображения, потому что некоторые из них имеют большой размер. Маленький размер фото означает что наша модель будет работать быстрее.

Для этого я написал следующий скрипт:

import os
from PIL import Image
# Define here your paths
original_images_path = '...'
resized_images_path = '...'
new_width = 416
new_height = 416
for image_file in os.listdir(original_images_path):
    ima = Image.open(original_images_path + image_file)
    new_ima = ima.resize((new_width, new_height), Image.ANTIALIAS)
    new_ima.save(resized_images_path + image_file)

Разметка данных

Это самы утомительный этап, поскольку вы должны разметить каждый простой объект который вы хотите чтоб модель распознала.

Вы можете разметить ваши изображения с помощью LabelImg. Запустите следующую команду в терминале для установки.

pip3 install labelImg

После установки вы можете приступать к разметки фотографий выполнив команду.

labelImg path_to_images/

Затем вам нужно создать прямоугольную рамку с соответствующей меткой, для каждого объекта на изображении.

После разметки всех объектов на изображении вы должны сохранить в  PascalVOC формат генерируемый .xml файл

Увеличение данных

Наличие большого набора данных имеет решающее значение для производительности нашей модели глубокого обучения, поэтому мы можем дополнить данные, которые у нас уже есть, с помощью библиотеки imgaug.

Этот github репозиторий объясняет и предоставляет код как дополнить изображения данными из размеченных прямоугольников. Этот процесс можно описать тремя этапами:

  • Конвертация .xml файл в .csv формат
  • Применить дополненные данные
  • Конфертировать результирующий .csv файл в множество .xml файлов

Сперва мы определим ф-ю конвертации наших .xml файлов в .csv

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


def xml_to_csv(path):
    xml_list = []
    for xml_file in glob.glob(path + '/*.xml'):
        tree = ET.parse(xml_file)
        root = tree.getroot()
        for member in root.findall('object'):
            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)
    return xml_df

Затем мы определим функцию обратного процесса.

import csv
import xml.etree.cElementTree as ET

def csv_to_xml(csv_path, resized_images_path, labels_path, folder):
    f = open(csv_path, 'r')
    reader = csv.reader(f)
    header = next(reader)
    old_filename = None
    for row in reader:
        filename = row[0]
        if filename == old_filename:
            object = ET.SubElement(annotation, 'object')
            ET.SubElement(object, 'name').text = row[3]
            ET.SubElement(object, 'pose').text = 'Unspecified'
            ET.SubElement(object, 'truncated').text = '0'
            ET.SubElement(object, 'difficult').text = '0'
            bndbox = ET.SubElement(object, 'bndbox')
            ET.SubElement(bndbox, 'xmin').text = row[4]
            ET.SubElement(bndbox, 'ymin').text = row[5]
            ET.SubElement(bndbox, 'xmax').text = row[6]
            ET.SubElement(bndbox, 'ymax').text = row[7]
        else:
            if old_filename is not None:
                labels_file = old_filename.replace('.jpg', '.xml')
                tree = ET.ElementTree(annotation)
                tree.write(labels_path + labels_file)

            annotation = ET.Element('annotation')
            ET.SubElement(annotation, 'folder').text = folder
            ET.SubElement(annotation, 'filename').text = filename
            ET.SubElement(annotation, 'path').text =     resized_images_path + filename
            source = ET.SubElement(annotation, 'source')
            ET.SubElement(source, 'database').text = 'Unknown'
            size = ET.SubElement(annotation, 'size')
            ET.SubElement(size, 'width').text = row[1]
            ET.SubElement(size, 'height').text = row[2]
            ET.SubElement(size, 'depth').text = '3'
            ET.SubElement(annotation, 'segmented').text = '0'

            object = ET.SubElement(annotation, 'object')
            ET.SubElement(object, 'name').text = row[3]
            ET.SubElement(object, 'pose').text = 'Unspecified'
            ET.SubElement(object, 'truncated').text = '0'
            ET.SubElement(object, 'difficult').text = '0'
            bndbox = ET.SubElement(object, 'bndbox')
            ET.SubElement(bndbox, 'xmin').text = row[4]
            ET.SubElement(bndbox, 'ymin').text = row[5]
            ET.SubElement(bndbox, 'xmax').text = row[6]
            ET.SubElement(bndbox, 'ymax').text = row[7]
        old_filename = filename
    f.close()

Затем мы определим наш данные с imgaug, для примера напишем следующее

from imgaug.augmentables.bbs import BoundingBoxesOnImage
from imgaug import augmenters as iaa

aug = iaa.SomeOf(2, [
    iaa.Affine(scale=(0.5, 1.5)),
    iaa.Affine(rotate=(-60, 60)),
    iaa.Affine(translate_percent={"x": (-0.3, 0.3), "y": (-0.3,   0.3)}),
    iaa.Fliplr(1),
    iaa.Multiply((0.5, 1.5)),
    iaa.GaussianBlur(sigma=(1.0, 3.0)),
    iaa.AdditiveGaussianNoise(scale=(0.03 * 255, 0.05 * 255)),
    iaa.Add((-25, 25)),
    iaa.MotionBlur(k=15),
    iaa.MultiplySaturation((0.5, 1.5)),
    iaa.LogContrast(gain=(0.6, 1.4)),
    iaa.Flipud(1)
])

В завершение применим этот пайплайн к нашему датасету и сохраним оба дополнительных изображения и их дополнительные метки. Вы также должны определить функции image_aug() и bbs_obj_to_df(), которые лежат в данном репозитории github.

if __name__ == '__main__':
    xml_df = xml_to_csv('labels_path/')
    for i in range(10):
        augmented_images_df = image_aug(xml_df, 'resized_images/', 'resized_images/', 'aug{}_'.format(i), aug)
        augmented_images_df.to_csv('aug{}_images.csv'.format(i), index=False)
        csv_to_xml(csv_path='aug{}_images.csv'.format(i),
                   resized_images_path='resized_images/', 
                   labels_path='labels_path/',
                   folder='resized_images')
        os.remove('aug{}_images.csv'.format(i))

Тренинг модели

После всего этого препроцессинга мы наконец можем обучить нашу модель с помощью ImageAI.

Прежде всего нужно установить библиотеку.

pip3 install imageai

Для использования трансферного обучения вы можете скачать модель отсюда.

Затем нам нужно разделить наш набор данных на папки train и validation.

dataset
├─ train
│ └── images
│ └── annotations
└─ validation
├── images
├── annotations

Это можно сделать с помощью следующего кода.

import os
import shutil
import random

images_path = 'resized_images/'
labels_path = 'labels/'
train_path = 'dataset/train/'
validation_path = 'dataset/validation/'


for image_file in os.listdir(images_path):
    labels_file = image_file.replace('.jpg', '.xml')
    if random.uniform(0, 1) > 0.2:
        shutil.copy(images_path + image_file, train_path + 'images/'
                    + image_file)
        shutil.copy(labels_path + labels_file, train_path +
                    'annotations/' + labels_file)
    else:
        shutil.copy(images_path + image_file, validation_path +

И мы можем начать тренировку с помощью следующего сценария.

from imageai.Detection.Custom import DetectionModelTrainer

trainer = DetectionModelTrainer()
trainer.setModelTypeAsYOLOv3()
trainer.setDataDirectory(data_directory="your_dataset_path")
trainer.setTrainConfig(object_names_array=["obj1", "obj2"], batch_size=16, num_experiments=200, train_from_pretrained_model="pretrained-yolov3.h5")
trainer.trainModel()

Оценка модели

После того как мы обучили нашу модель, мы можем оценить каждый чекпоинт обучения через mAP.

from imageai.Detection.Custom import DetectionModelTrainer

trainer = DetectionModelTrainer()
trainer.setModelTypeAsYOLOv3()
trainer.setDataDirectory(data_directory="dataset")
metrics = trainer.evaluateModel(model_path="dataset/models", json_path="dataset/json/detection_config.json", iou_threshold=0.5, object_threshold=0.9, nms_threshold=0.5)

Вывод должен выглядеть примерно так.

Используем нашу модель для обнаружения

Наконец, мы можем использовать нашу модель для обнаружения объектов на одном изображении.

from imageai.Detection.Custom import CustomObjectDetection

detector = CustomObjectDetection()
detector.setModelTypeAsYOLOv3()
detector.setModelPath("detection_model-ex-006--loss-0024.905.h5")
detector.setJsonPath("detection_config.json")
detector.loadModel()
detections = detector.detectObjectsFromImage(input_image="ima.jpg", output_image_path="ima-detected.jpg", minimum_percentage_probability=90)
view raw

Вы можете настроить параметр Minim_percentage_probability, чтобы показать обнаруженные объекты с большей или меньшей достоверностью, получив результат, подобный следующему.

Теперь я хотел бы услышать от вас, как это работает!

Спасибо за прочтение!

Перевод статьи, источник.

Post Created 8

Related Posts

Begin typing your search above and press enter to search. Press ESC to cancel.

Back To Top