Mask R-CNN от новичка до профессионала

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

Ну так распишем подробнее наши задачи:

• Фиксировать людей и автомобили — выделять их на изображении и генерировать соответствующие экземпляры классов с необходимыми полями.
• Определять номер авто, если он попал в кадр определенной камеры
• Сравнить текущий кадр с предыдущим на равенство объектов, чтобы мы могли узнать

Ок, подумал я, и взял в руки толстую змею, python, значится. Было решено использовать нейронную сетку Mask R-Cnn в связи с ее простотой и современными характеристиками. Также, разумеется, для манипуляция с изображениями будем использовать OpenCV.

Установка среды

Будем использовать Windows 10, из-за того, что ты вероятнее всего, используешь именно ее.
Подразумевается, что у тебя уже присутствует 64 битный Python. Если же нет, то можно скачать пакет, например, отсюда

Установка пакетов

git clone https://github.com/matterport/Mask_RCNN
cd Mask_RCNN
pip3 install -r requirements.txt
python3 setup.py install

Если по какой-то причине не удается собрать из исходников, существует версия из pip:

pip3 install mrcnn --user

К пакету, разумеется, поставятся все зависимости.

Этап 1. Создание простейшей программы-распознавателя.

Сделаем необходимые импорты

import os
import cv2

import mrcnn.config
import mrcnn
from mrcnn.model import MaskRCNN

Нейросеть требует создания конфига с переопределенными полями

class MaskRCNNConfig(mrcnn.config.Config):
    NAME = "coco_pretrained_model_config"
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1
    DETECTION_MIN_CONFIDENCE = 0.8 # минимальный процент отображения прямоугольника
    NUM_CLASSES = 81

Укажем расположение файла с весами. Пусть в данном примере он будет лежать в папке с этим файлом. Если его нет, то он скачается.

import mrcnn.utils
DATASET_FILE = "mask_rcnn_coco.h5"
if not os.path.exists(DATASET_FILE):
    mrcnn.utils.download_trained_weights(DATASET_FILE)

Создадим нашу модель с настройками выше

model = MaskRCNN(mode="inference", model_dir="logs", config=MaskRCNNConfig())
model.load_weights(DATASET_FILE, by_name=True)

И пожалуй, начнем обработку всех изображений в каталоге images в текущей директории.

IMAGE_DIR = os.path.join(os.getcwd(), "images")
for filename in os.listdir(IMAGE_DIR):
    image = cv2.imread(os.path.join(IMAGE_DIR, filename))
    rgb_image = image[:, :, ::-1]
    detections = model.detect([rgb_image], verbose=1)[0]

Что же мы увидим в detections?

    print(detections)

Например, что-то похожее:

{'rois': array([[ 303, 649, 542, 1176],[ 405, 2, 701, 319]]), 'class_ids': array([3, 3]), 
'scores': array([0.99896, 0.99770015], dtype=float32), 
'masks': array()}

В данном случае нашли 2 объекта.
rois — массивы координат левого нижнего и правого верхнего угла
class_ids — числовые идентификаторы найденных объектов, пока нам нужно знать, что 1 — человек, 3 — машина, 8 — грузовик.
scores — насколько модель уверена в решении, этот параметр можно отсеивать через DETECTION_MIN_CONFIDENCE в конфиге, отсекая все неподходящие варианты.
masks — контур объекта. Данные используются для рисования маски объекта. Т.к. они достаточно объемны, и не предназначены для понимания человеком, приводить в статье их не буду.

Ок, мы могли бы на этом остановиться, но мы же хотим посмотреть на изображение, которое обычно выдают гайды по использованию нейросеток с красиво выделенными объектами?

Проще было бы вызвать функцию mrcnn.visualize.display_instances, но мы не будем так делать, напишем свою.

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

def visualize_detections(image, masks, boxes, class_ids, scores):
    import numpy as np
    bgr_image = image[:, :, ::-1]

    CLASS_NAMES = ['BG',"person", "bicycle", "car", "motorcycle", "bus", "truck"]
    COLORS = mrcnn.visualize.random_colors(len(CLASS_NAMES))

    for i in range(boxes.shape[0]):        
        y1, x1, y2, x2 = boxes[i]

        classID = class_ids[i]            
        label = CLASS_NAMES[classID]
        font = cv2.FONT_HERSHEY_DUPLEX
        color = [int(c) for c in np.array(COLORS[classID]) * 255]
        text = "{}: {:.3f}".format(label, scores[i])
        size = 0.8
        width = 2

        cv2.rectangle(bgr_image, (x1, y1), (x2, y2), color, width)
        cv2.putText(bgr_image, text, (x1, y1-20), font, size, color, width)        

Хотя одним из основных преимуществ этой нейронной сети является решение задач Instance segmentation — получение контуров объектов, мы пока этим не воспользовались, разберем это.

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

mask = masks[:, :, i]  # берем срез
image = mrcnn.visualize.apply_mask(image, mask, color, alpha=0.6) # рисование маски

Результат:

Этап II. Первые успехи. Распознавание номеров машин.

Для распознавания нам нужен четкий кадр машины вблизи, поэтому было решено брать только кадры с КПП, а потом сравнивать на похожесть(об этом в следующей главе). Этот способ, правда, дает слишком большую неточность, т.к. машины могут быть очень похожими визуально и мой алгоритм пока не может избегать такие ситуации.

Было решено использовать готовую либу от украинского производителя nomeroff-net (не реклама). Т.к. почти весь код можно найти в примерах к модели, то приводить полное описание не буду.

Скажу лишь, что эту функцию можно запускать с исходным изображением или распознанную машину можно вырезать из кадра и подать в эту функцию.

import sys
import matplotlib.image as mpimg
import os
sys.path.append(cfg.NOMEROFF_NET_DIR)
from NomeroffNet import  filters, RectDetector, TextDetector, OptionsDetector, Detector, textPostprocessing

nnet = Detector(cfg.MASK_RCNN_DIR, cfg.MASK_RCNN_LOG_DIR)
nnet.loadModel("latest")

rectDetector = RectDetector()
optionsDetector = OptionsDetector()
optionsDetector.load("latest")
textDetector = TextDetector.get_static_module("ru")()
textDetector.load("latest")

def detectCarNumber(imgPath: str) -> str:
    img = mpimg.imread(imgPath)
    NP = nnet.detect([img])

    cvImgMasks = filters.cv_img_mask(NP)

    arrPoints = rectDetector.detect(cvImgMasks)
    zones = rectDetector.get_cv_zonesBGR(img, arrPoints)

    regionIds, stateIds, _c = optionsDetector.predict(zones)
    regionNames = optionsDetector.getRegionLabels(regionIds)

    # find text with postprocessing by standart  
    textArr = textDetector.predict(zones)
    textArr = textPostprocessing(textArr, regionNames)
    return textArr

textArr на выходе будет представлять массив строк с номерами машин, найденных на кадре, например:
["К293РР163"], или [""], [] — если подходящие номера не были найдены.

Этап III. Опознаем объекты на схожесть.

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

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

Предложу для этих целей sift алгоритм. Оговоримся, что он не входит в основную часть OpenCV, поэтому нам необходимо доставить contrib модули дополнительно. К сожалению, алгоритм запатентован и его использование в коммерческих программах ограничено. Но мы нацелены на научно-исследовательскую деятельность, не так ли?

pip3 install opencv-contrib-python --user

~~ Перегружаем оператор == ~~ Пишем функцию, принимающую 2 сравниваемых объекта в виде матриц. Например, мы их получаем после вызова функции cv2.open(path)

Напишем реализацию нашего алгоритма.

def compareImages(img1, img2) -> bool:
    sift = cv2.xfeatures2d.SIFT_create()

Находим ключевые точки и дескрипторы с помощью SIFT. Пожалуй, хелп для этих функций приводить не буду, ведь его всегда его можно вызвать в интерактивной оболочке как help(somefunc)

    kp1, des1 = sift.detectAndCompute(img1, None)
    kp2, des2 = sift.detectAndCompute(img2, None)

Настроим наш алгоритм.

    FLANN_INDEX_KDTREE = 0
    indexParams = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
    searchParams = dict(checks=50)
    flann = cv2.FlannBasedMatcher(indexParams, searchParams)

Теперь запустим его.

    matches = flann.knnMatch(des1, des2, k=2)

Посчитаем сходства между изображениями

    matchesCount = 0 
    for m, n in matches:
        if m.distance < cfg.cencitivity*n.distance:
            matchesCount += 1

    return matchesCount > cfg.MIN_MATCH_COUNT

А теперь, попробуем его использовать
Для этого, после обнаружения объектов, нам нужно их вырезать с исходного изображения

Я не смог написать ничего лучше, чем сохранить на медленную память, а потом уже оттуда считывать.

    def extractObjects(objects, binaryImage, outputImageDirectory, filename=None):
        for item in objects:
            y1, x1, y2, x2 = item.coordinates
            # вырежет все объекты в отдельные изображения
            cropped = binaryImage[y1:y2, x1:x2]
            beforePoint, afterPoint = filename.split(".")
            outputDirPath = os.path.join(os.path.split(outputImageDirectory)[0], "objectsOn" + beforePoint)
                if not os.path.exists(outputDirPath):
                    os.mkdir(outputDirPath)
                coordinates = str(item).replace(" ", ",")
                pathToObjectImage = "{}{}.jpg".format(item.type, coordinates)
                cv2.imwrite(os.path.join(outputDirPath, str(pathToObjectImage)), cropped)

Теперь у нас есть объекты в каталоге <outputImageDirectory>/objectsOn<imageFilename>

Теперь, если у нас есть как минимум 2 таких каталога, то мы можем сравнивать объекты в них. Запустим функцию, написанную ранее

if compareImages(previousObjects, currentObjects):
    print(“Эти объекты одинаковы!”)

Или мы можем сделать другое действие, вроде маркировки этих объектов одинаковым айдишником.

Конечно, как и все нейронные сети, эта склонна давать иногда ошибочные результаты.

В целом, мы выполнили 3 задачи, поставленные в начале, так что будем закругляться. Я сомневаюсь, что эта статья открыла глаза людям, которые написали хотя бы одну программу, решающую задачи image recognition/image segmentation, но я надеюсь, что я помог хотя бы одному начинающему разработчику).

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