You only look once, или YOLO, — эффектив­ный алго­ритм, который поз­воля­ет выделять объ­екты на изоб­ражении. В этой статье мы исполь­зуем его, что­бы написать на Python прог­рамму для опре­деле­ния чис­ла людей в помеще­нии, а по дороге опро­буем его в раз­гадыва­нии кап­чи.

Ес­ли ты смот­рел «Тер­минатор», то пом­нишь кад­ры из глаз T-800: он смот­рел по сто­ронам и опре­делял раз­ные объ­екты. Тог­да о такой машине мож­но было толь­ко меч­тать, а сегод­ня ее мож­но смас­терить самому из готовых час­тей.

Кадр из фильма «Терминатор»
Кадр из филь­ма «Тер­минатор»

Рас­позна­вание объ­ектов сегод­ня при­гож­дает­ся для решения самых раз­ных задач: клас­сифика­ции видов рас­тений и живот­ных, рас­позна­вания лиц, опре­деле­ния габари­тов объ­ектов — и это далеко не пол­ный спи­сок.

 

Какие бывают алгоритмы

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

R-CNN, Region-Based Convolutional Neural Network

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

Об­работ­ка в режиме реаль­ного вре­мени: не под­держи­вает­ся.

Fast R-CNN, Fast Region-Based Convolutional Neural Network

Под­ход ана­логи­чен алго­рит­му R-CNN. Но вмес­то того, что­бы пред­варитель­но выделять реги­оны, мы переда­ем вход­ное изоб­ражение в CNN для соз­дания свер­точной кар­ты приз­наков, где затем будет про­исхо­дить выбороч­ный поиск, а пред­ска­зание клас­са объ­ектов выпол­няет спе­циаль­ный слой Softmax.

Об­работ­ка в режиме реаль­ного вре­мени: не под­держи­вает­ся.

Faster R-CNN, Faster Region-Based Convolutional Neural Network

По­доб­но Fast R-CNN, изоб­ражение переда­ется в CNN соз­дания свер­точной кар­ты приз­наков, но вмес­то алго­рит­ма выбороч­ного поис­ка для прог­нозиро­вания пред­ложений по реги­онам исполь­зует­ся отдель­ная сеть.

Об­работ­ка в режиме реаль­ного вре­мени: под­держи­вает­ся при высоких вычис­литель­ных мощ­ностях.

YOLO, You Only Look Once

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

Об­работ­ка в режиме реаль­ного вре­мени: под­держи­вает­ся!

Как видишь, YOLO пока что луч­ший вари­ант для обна­руже­ния и рас­позна­вания обра­зов. Он отли­чает­ся высокой ско­ростью и точ­ностью обна­руже­ния объ­ектов, а еще этот алго­ритм мож­но исполь­зовать в про­ектах на Android и Raspberry Pi с помощью нет­ребова­тель­ного tiny-вари­анта сети, с которым мы с тобой сегод­ня будем работать.

Tiny-вари­ант нес­коль­ко про­игры­вает в точ­ности пол­ноцен­ному вари­анту сети, но и тре­бует мень­шей вычис­литель­ной мощ­ности, что поз­волит запус­тить про­ект, который мы сегод­ня будем делать, как на сла­бом компь­юте­ре, так и, при желании, на смар­тфо­не.

 

Пишем код

Что­бы написать лег­ковес­ное при­ложе­ние для обна­руже­ния объ­ектов на изоб­ражении, нам с тобой понадо­бят­ся:

До­пол­нитель­но уста­новим биб­лиоте­ки OpenCV и NumPy:

pip install opencv-python
pip install numpy

Те­перь напишем при­ложе­ние, которое будет находить объ­екты на изоб­ражении при помощи YOLO и отме­чать их.

Мы поп­робу­ем обой­ти CAPTCHA с изоб­ражени­ями гру­зови­ков — класс truck в датасе­те COCO. Допол­нитель­но мы пос­чита­ем количес­тво обна­ружен­ных объ­ектов нуж­ного нам клас­са и выведем всю информа­цию на экран.

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

import cv2
import numpy as np
def apply_yolo_object_detection(image_to_process):
"""
Распознавание и определение координат объектов на изображении
:param image_to_process: исходное изображение
:return: изображение с размеченными объектами и подписями к ним
"""
height, width, depth = image_to_process.shape
blob = cv2.dnn.blobFromImage(image_to_process, 1 / 255, (608, 608), (0, 0, 0), swapRB=True, crop=False)
net.setInput(blob)
outs = net.forward(out_layers)
class_indexes, class_scores, boxes = ([] for i in range(3))
objects_count = 0
# Запуск поиска объектов на изображении
for out in outs:
for obj in out:
scores = obj[5:]
class_index = np.argmax(scores)
class_score = scores[class_index]
if class_score > 0:
center_x = int(obj[0] * width)
center_y = int(obj[1] * height)
obj_width = int(obj[2] * width)
obj_height = int(obj[3] * height)
box = [center_x - obj_width // 2, center_y - obj_height // 2, obj_width, obj_height]
boxes.append(box)
class_indexes.append(class_index)
class_scores.append(float(class_score))
# Выборка
chosen_boxes = cv2.dnn.NMSBoxes(boxes, class_scores, 0.0, 0.4)
for box_index in chosen_boxes:
box_index = box_index[0]
box = boxes[box_index]
class_index = class_indexes[box_index]
# Для отладки рисуем объекты, входящие в искомые классы
if classes[class_index] in classes_to_look_for:
objects_count += 1
image_to_process = draw_object_bounding_box(image_to_process, class_index, box)
final_image = draw_object_count(image_to_process, objects_count)
return final_image

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

def draw_object_bounding_box(image_to_process, index, box):
"""
Рисование границ объекта с подписями
:param image_to_process: исходное изображение
:param index: индекс определенного с помощью YOLO класса объекта
:param box: координаты области вокруг объекта
:return: изображение с отмеченными объектами
"""
x, y, w, h = box
start = (x, y)
end = (x + w, y + h)
color = (0, 255, 0)
width = 2
final_image = cv2.rectangle(image_to_process, start, end, color, width)
start = (x, y - 10)
font_size = 1
font = cv2.FONT_HERSHEY_SIMPLEX
width = 2
text = classes[index]
final_image = cv2.putText(final_image, text, start, font, font_size, color, width, cv2.LINE_AA)
return final_image

До­бавим фун­кцию, которая выведет количес­тво рас­познан­ных объ­ектов на изоб­ражении.

def draw_object_count(image_to_process, objects_count):
"""
Подпись количества найденных объектов на изображении
:param image_to_process: исходное изображение
:param objects_count: количество объектов искомого класса
:return: изображение с подписанным количеством найденных объектов
"""
start = (45, 150)
font_size = 1.5
font = cv2.FONT_HERSHEY_SIMPLEX
width = 3
text = "Objects found: " + str(objects_count)
# Вывод текста с обводкой (чтобы было видно при разном освещении картинки)
white_color = (255, 255, 255)
black_outline_color = (0, 0, 0)
final_image = cv2.putText(image_to_process, text, start, font, font_size, black_outline_color, width * 3, cv2.LINE_AA)
final_image = cv2.putText(final_image, text, start, font, font_size, white_color, width, cv2.LINE_AA)
return final_image

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

def start_image_object_detection():
"""
Анализ изображения
"""
try:
# Применение методов распознавания объектов на изображении от YOLO
image = cv2.imread("assets/truck_captcha.png")
image = apply_yolo_object_detection(image)
# Вывод обработанного изображения на экран
cv2.imshow("Image", image)
if cv2.waitKey(0):
cv2.destroyAllWindows()
except KeyboardInterrupt:
pass

А теперь мы соз­дадим фун­кцию main, в которой нас­тро­им нашу сеть и поп­робу­ем запус­тить ее.

if __name__ == '__main__':
# Загрузка весов YOLO из файлов и настройка сети
net = cv2.dnn.readNetFromDarknet("yolov4-tiny.cfg", "yolov4-tiny.weights")
layer_names = net.getLayerNames()
out_layers_indexes = net.getUnconnectedOutLayers()
out_layers = [layer_names[index[0] - 1] for index in out_layers_indexes]
# Загрузка из файла классов объектов, которые умеет обнаруживать YOLO
with open("coco.names.txt") as file:
classes = file.read().split("\n")
# Определение классов, которые будут приоритетными для поиска на изображении
# Названия находятся в файле coco.names.txt
# В данном случае определяется грузовик для прохождения CAPTCHA
classes_to_look_for = ["truck"]
start_image_object_detection()

Да­вай пос­мотрим, как алго­ритм YOLO спра­вил­ся с тес­том прос­той CAPTCHA.

Исходная CAPTCHA
Ис­ходная CAPTCHA
Результат применения YOLO
Ре­зуль­тат при­мене­ния YOLO

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

Ка­жет­ся, пока что вос­ста­ние машин нам не гро­зит! 🙂

 

Модифицируем приложение

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

Что­бы задача была на «живом при­мере», мы вос­поль­зуем­ся пуб­личной камерой, уста­нов­ленной в одном из бар­бершо­пов Лон­дона. Из‑за его скром­ной пло­щади находить­ся внут­ри может не боль­ше десяти человек.

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

def start_video_object_detection():
"""
Захват и анализ видео в режиме реального времени
"""
while True:
try:
# Захват картинки с видео
video_camera_capture = cv2.VideoCapture("http://81.130.136.82:82/mjpg/video.mjpg")
while video_camera_capture.isOpened():
ret, frame = video_camera_capture.read()
if not ret:
break
# Применение методов распознавания объектов на кадре видео от YOLO
frame = apply_yolo_object_detection(frame)
# Вывод обработанного изображения на экран с уменьшением размера окна
frame = cv2.resize(frame, (1920 // 2, 1080 // 2))
cv2.imshow("Video Capture", frame)
if cv2.waitKey(0):
break
video_camera_capture.release()
cv2.destroyAllWindows()
except KeyboardInterrupt:
pass

Так­же нам пот­ребу­ется нем­ного модифи­циро­вать фун­кцию main, что­бы теперь запус­кать обра­бот­ку видео вмес­то обра­бот­ки изоб­ражения.

if __name__ == '__main__':
# Загрузка весов YOLO из файлов и настройка сети
net = cv2.dnn.readNetFromDarknet("yolov4-tiny.cfg", "yolov4-tiny.weights")
layer_names = net.getLayerNames()
out_layers_indexes = net.getUnconnectedOutLayers()
out_layers = [layer_names[index[0] - 1] for index in out_layers_indexes]
# Загрузка из файла классов объектов, которые умеет обнаруживать YOLO
with open("coco.names.txt") as file:
classes = file.read().split("\n")
# Определение классов, которые будут приоритетными для поиска на изображении
# Названия находятся в файле coco.names.txt
# В данном случае определяется грузовик для прохождения CAPTCHA и человек для анализа видео
classes_to_look_for = ["truck", "person"]
start_video_object_detection()

По­луча­ем резуль­тат: шесть из семи человек были рас­позна­ны.

Результат обработки видео
Ре­зуль­тат обра­бот­ки видео

Мож­но добавить и дру­гие полез­ные фун­кции: нап­ример, отправ­лять на поч­ту или в Telegram сооб­щение о том, что в бар­бершоп набилось мно­гова­то людей.

 

Итоги

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

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

Здесь пот­ребу­ются алго­рит­мы, опре­деля­ющие в пер­спек­тиве габари­ты объ­ектов и работа­ющие с трех­мерным прос­транс­твом. Такие исполь­зуют­ся для опре­деле­ния тран­спортных средств в само­управля­емых авто­моби­лях — Aggregate View Object Detection или YOLO 3D Oriented Object Bounding Box Detection.

Ис­ходни­ки про­екта, с которы­ми ты смо­жешь порабо­тать над решени­ем подоб­ных задач, смот­ри в мо­ем репози­тории на GitHub.

Оцени статью:

Что тебе понравилось больше всего?
Что тебе не понравилось больше всего?