Если вам кажется, что начать работу с нейросетями - это сложно, то этот материал для вас!

Итак, YOLO (You Only Look Once) — нейронная сеть, предназначенная работы с объектами на изображениях и может решать следующие задачи:

• Детекция - обнаружение объектов

• Сегментация - разделение изображения на области, которые относятся к каждому объекту

• Классификация - определение что же находится на изображении

• Поиск ключевых точек тела - для определения позы человека

• Трекинг объектов - потоковая обработка, при которой для каждого объекта возможно сохранять и использовать историю местоположения

Также в этой статье также будет рассмотрено:

• Предсказание движения на основе трекинга

• Создание собственного датасета для дообучения модели детекции новых объектов

Подготовка к работе с YOLO

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

pip install ultralytics

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

Детекция объектов на фото

Детекция объектов - определение местоположения объектов и их классов на изображении.

Для примера возьмем картинку и определим объекты на ней:

Для использования нейросети YOLO напишем скрипт:

from ultralytics import YOLO
import cv2
import numpy as np
import os

# Загрузка модели YOLOv8
model = YOLO('yolov8n.pt')

# Список цветов для различных классов
colors = [
    (255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (0, 255, 255),
    (255, 0, 255), (192, 192, 192), (128, 128, 128), (128, 0, 0), (128, 128, 0),
    (0, 128, 0), (128, 0, 128), (0, 128, 128), (0, 0, 128), (72, 61, 139),
    (47, 79, 79), (47, 79, 47), (0, 206, 209), (148, 0, 211), (255, 20, 147)
]

# Функция для обработки изображения
def process_image(image_path):
    # Загрузка изображения
    image = cv2.imread(image_path)
    results = model(image)[0]
    
    # Получение оригинального изображения и результатов
    image = results.orig_img
    classes_names = results.names
    classes = results.boxes.cls.cpu().numpy()
    boxes = results.boxes.xyxy.cpu().numpy().astype(np.int32)

    # Подготовка словаря для группировки результатов по классам
    grouped_objects = {}

    # Рисование рамок и группировка результатов
    for class_id, box in zip(classes, boxes):
        class_name = classes_names[int(class_id)]
        color = colors[int(class_id) % len(colors)]  # Выбор цвета для класса
        if class_name not in grouped_objects:
            grouped_objects[class_name] = []
        grouped_objects[class_name].append(box)

        # Рисование рамок на изображении
        x1, y1, x2, y2 = box
        cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
        cv2.putText(image, class_name, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

    # Сохранение измененного изображения
    new_image_path = os.path.splitext(image_path)[0] + '_yolo' + os.path.splitext(image_path)[1]
    cv2.imwrite(new_image_path, image)

    # Сохранение данных в текстовый файл
    text_file_path = os.path.splitext(image_path)[0] + '_data.txt'
    with open(text_file_path, 'w') as f:
        for class_name, details in grouped_objects.items():
            f.write(f"{class_name}:\n")
            for detail in details:
                f.write(f"Coordinates: ({detail[0]}, {detail[1]}, {detail[2]}, {detail[3]})\n")

    print(f"Processed {image_path}:")
    print(f"Saved bounding-box image to {new_image_path}")
    print(f"Saved data to {text_file_path}")


process_image('test.png')

В начале кода — выбор модели (она скачается автоматически при первом запуске скрипта). Кроме минимальной yolov8n.pt, доступны еще несколько:
yolov8n.pt yolov8s.pt yolov8m.pt yolov8l.pt yolov8x.pt

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

Сравнительные характеристики моделей, а также возможность их запуска на различных устройствах описаны в статье Степана Жданова:
https://habr.com/ru/articles/822917/

В результате работы нейросети получаем объект boxes, который содержит информацию о координатах, найденных на изображении объектов, а также принадлежности их к классу (person, car, bus, traffic light и т.д.)

Итак, после выполнения данного скрипта видим результат:

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

car:
Coordinates: (842, 681, 1180, 894)
Coordinates: (254, 849, 524, 971)
Coordinates: (49, 620, 425, 857)
stop sign:
Coordinates: (407, 560, 470, 626)
Coordinates: (267, 494, 341, 557)
traffic light:
Coordinates: (334, 157, 451, 426)
Coordinates: (938, 97, 1031, 312)
Coordinates: (86, 481, 130, 602)
person:
Coordinates: (578, 711, 710, 990)
Coordinates: (715, 723, 750, 798)
Coordinates: (715, 852, 864, 976)
Coordinates: (241, 897, 385, 1012)
truck:
Coordinates: (52, 620, 425, 859)

Детекция объектов в видео-потоке

После практики с одним изображением перейдем к обработке видео. В целом, оно не сильно сложнее, т.к. видео это всего лишь последовательность изображений! Отличительной особенностью следующего кода является только то, что для записи видео используются специальные кодеки для формата MP4.

from ultralytics import YOLO
import cv2
import numpy as np

# Загрузка модели YOLOv8
model = YOLO('yolov8n.pt')

# Список цветов для различных классов
colors = [
    (255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (0, 255, 255),
    (255, 0, 255), (192, 192, 192), (128, 128, 128), (128, 0, 0), (128, 128, 0),
    (0, 128, 0), (128, 0, 128), (0, 128, 128), (0, 0, 128), (72, 61, 139),
    (47, 79, 79), (47, 79, 47), (0, 206, 209), (148, 0, 211), (255, 20, 147)
]

# Открытие исходного видеофайла
input_video_path = 'input.mp4'
capture = cv2.VideoCapture(input_video_path)

# Чтение параметров видео
fps = int(capture.get(cv2.CAP_PROP_FPS))
width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))

# Настройка выходного файла
output_video_path = 'detect.mp4'
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
writer = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))

while True:
    # Захват кадра
    ret, frame = capture.read()
    if not ret:
        break

    # Обработка кадра с помощью модели YOLO
    results = model(frame)[0]

    # Получение данных об объектах
    classes_names = results.names
    classes = results.boxes.cls.cpu().numpy()
    boxes = results.boxes.xyxy.cpu().numpy().astype(np.int32)
    # Рисование рамок и подписей на кадре
    for class_id, box, conf in zip(classes, boxes, results.boxes.conf):
        if conf>0.5:
            class_name = classes_names[int(class_id)]
            color = colors[int(class_id) % len(colors)]
            x1, y1, x2, y2 = box
            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
            cv2.putText(frame, class_name, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

    # Запись обработанного кадра в выходной файл
    writer.write(frame)

# Освобождение ресурсов и закрытие окон
capture.release()
writer.release()

Для детекции объектов при помощи веб камеры - достаточно в строке
capture = cv2.VideoCapture(input_video_path)
указать индекс камеры для подключения
capture = cv2.VideoCapture(0)
Пример полного кода:
https://github.com/stepanburmistrov/YoloV8/blob/main/Detection_Camera.py

Сегментация изображения

Сегментация - разделение изображения на классы. Один из самых эффектных способов разобраться с этим процессом будет удаление фона вокруг человека с фото/видео.

Для начала, возьмем фотографию (кадр из предыдущего видео) и применим модель YoloV8 предназначенную для сегментации. Как и в случае с детекцией выбор есть:
yolov8n-seg.pt yolov8s-seg.pt yolov8m-seg.pt yolov8l-seg.pt yolov8x-seg.pt

import cv2
import numpy as np
from ultralytics import YOLO
import os

# Загрузка модели YOLOv8
model = YOLO('yolov8x-seg.pt')

colors = [
    (255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (0, 255, 255),
    (255, 0, 255), (192, 192, 192), (128, 128, 128), (128, 0, 0), (128, 128, 0),
    (0, 128, 0), (128, 0, 128), (0, 128, 128), (0, 0, 128), (72, 61, 139),
    (47, 79, 79), (47, 79, 47), (0, 206, 209), (148, 0, 211), (255, 20, 147)
]

def process_image(image_path):
    # Проверка наличия папки для сохранения результатов
    if not os.path.exists('results'):
        os.makedirs('results')
    
    # Загрузка изображения
    image = cv2.imread(image_path)
    image_orig = image.copy()
    h_or, w_or = image.shape[:2]
    image = cv2.resize(image, (640, 640))
    results = model(image)[0]
    
    classes_names = results.names
    classes = results.boxes.cls.cpu().numpy()
    masks = results.masks.data.cpu().numpy()

    # Наложение масок на изображение
    for i, mask in enumerate(masks):
        color = colors[int(classes[i]) % len(colors)]
        
        # Изменение размера маски перед созданием цветной маски
        mask_resized = cv2.resize(mask, (w_or, h_or))
        
        # Создание цветной маски
        color_mask = np.zeros((h_or, w_or, 3), dtype=np.uint8)
        color_mask[mask_resized > 0] = color

        # Сохранение маски каждого класса в отдельный файл
        mask_filename = os.path.join('results', f"{classes_names[classes[i]]}_{i}.png")
        cv2.imwrite(mask_filename, color_mask)

        # Наложение маски на исходное изображение
        image_orig = cv2.addWeighted(image_orig, 1.0, color_mask, 0.5, 0)


    # Сохранение измененного изображения
    new_image_path = os.path.join('results', os.path.splitext(os.path.basename(image_path))[0] + '_segmented' + os.path.splitext(image_path)[1])
    cv2.imwrite(new_image_path, image_orig)
    print(f"Segmented image saved to {new_image_path}")

process_image('segmentation_test.png')

В результате получим вот такое изображение:

А также в папку results сохранятся маски каждого найденного класса, для удобного исследования

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

Пример кода, решающего эту задачу:

import cv2
import numpy as np
from ultralytics import YOLO
import os

model = YOLO('yolov8n-seg.pt')

# Цвет для выделения объектов класса "person"
person_color = (0, 255, 0)  # Зеленый цвет

def process_image(image_path):
    frame = cv2.imread(image_path)
    if frame is None:
        print("Ошибка: не удалось загрузить изображение")
        return

    image_orig = frame.copy()
    h_or, w_or = frame.shape[:2]
    image = cv2.resize(frame, (640, 640))
    results = model(image)[0]

    classes = results.boxes.cls.cpu().numpy()
    masks = results.masks.data.cpu().numpy()

    # Создаем зеленый фон
    green_background = np.zeros_like(image_orig)
    green_background[:] = person_color

    # Наложение масок на изображение
    for i, mask in enumerate(masks):
        class_name = results.names[int(classes[i])]
        if class_name == 'person':
            color_mask = np.zeros((640, 640, 3), dtype=np.uint8)
            resized_mask = cv2.resize(mask, (640, 640), interpolation=cv2.INTER_NEAREST)
            color_mask[resized_mask > 0] = person_color

            color_mask = cv2.resize(color_mask, (w_or, h_or), interpolation=cv2.INTER_NEAREST)

            mask_resized = cv2.resize(mask, (w_or, h_or), interpolation=cv2.INTER_NEAREST)
            green_background[mask_resized > 0] = image_orig[mask_resized > 0]

    # Сохраняем обработанное изображение с добавлением суффикса '_segmented'
    base_name, ext = os.path.splitext(image_path)
    output_path = f"{base_name}_removed_BG{ext}"
    cv2.imwrite(output_path, green_background)
    print(f"Processed image saved to {output_path}")

    cv2.imshow('Processed Image', green_background)  # Показываем обработанное изображение
    cv2.waitKey(0)
    cv2.destroyAllWindows()

# Путь к изображению, которое необходимо обработать
image_path = 'test.jpg'
process_image(image_path)

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

В качестве еще одного примера для использования — автоматическое создание стикеров:
https://github.com/stepanburmistrov/YoloV8/blob/main/Segmentation_Image2Sticker.py

Классификация

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

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

1. Сложность задачи:

   • Классификация: Определяет, к какому классу относится весь объект или изображение целиком. Например, классификация фотографии как "собака" или "кошка".

   • Детекция: Находит объекты внутри изображения и определяет их классы и местоположение (например, где на фотографии находится собака и где кошка).

2. Ресурсы и вычислительная мощность:

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

   • Детекция: Более вычислительно затратна, так как требует анализа изображения для поиска объектов и определения их границ.

3. Скорость выполнения:

   • Классификация: Быстрее, так как выполняется одна операция определения класса для всего изображения.

   • Детекция: Медленнее, так как требует многократного анализа изображения для поиска всех объектов и их классификации.

4. Сложность реализации:

   • Классификация: Проще в реализации и настройке, так как обучается на меньших и более структурированных данных.

   • Детекция: Более сложна в реализации, требует больше данных и времени на обучение, особенно если нужно обнаруживать объекты разного размера и формы.

Когда использовать классификацию

1. Целостное определение класса объекта: Если нужно определить класс всего объекта или изображения, а не его частей. Например, определить, что изображено на фотографии (собака или кошка).

2. Ограниченные вычислительные ресурсы: В условиях, когда ограничены ресурсы для вычислений и требуется быстрая обработка данных.

3. Более точное определение "подклассов". Например, можно находить буквы на изображении с помощью детекции, а затем более точно определять символ при помощи классификации!

Когда использовать детекцию

1. Множественные объекты на изображении: Если на изображении может быть несколько объектов разных классов, и нужно определить их местоположение и классы. Например, обнаружение автомобилей и пешеходов на улице.

2. Анализ сложных сцен: Когда нужно анализировать сложные сцены, где важно не только определить, какие объекты присутствуют, но и где они находятся.

3. Применение в реальном времени: В задачах, где необходимо отслеживать объекты в реальном времени, например, в системах видеонаблюдения.

Пример кода для одного изображения:

import cv2
import numpy as np
from ultralytics import YOLO
import os

model = YOLO('yolov8n-cls.pt')


def process_image(img):
    # Обработка кадра с помощью модели
    results = model(img)[0]

    # Отображение результатов классификации на изображении
    if results.probs is not None:
        # Доступ к вершинам классификации
        top1_idx = results.probs.top1  # Индекс класса с наивысшей вероятностью
        top1_conf = results.probs.top1conf.item()  # Вероятность для класса с наивысшей вероятностью
        class_name = results.names[top1_idx]  # Получаем имя класса по индексу

        # Отображаем класс и вероятность на кадре
        label = f"{class_name}: {top1_conf:.2f}"
        cv2.putText(img, label, (50, 50),
                    cv2.FONT_HERSHEY_SIMPLEX, 2,
                    (255, 0, 0), 3)

    return image

image = cv2.imread("dog.jpg")
image = process_image(image)
cv2.imwrite('result.jpg', image)

Поиск ключевых точек тела - определения позы человека

Способой применения данной модели можно найти множество, вот некоторые из них:

• Тренировки спортсменов: Помощь в анализе движений спортсменов для улучшения их техники

• Реабилитация: Мониторинг и корректировка движений пациентов во время реабилитационных упражнений.

• Обнаружение падений: Автоматическое обнаружение падений и других опасных ситуаций для пожилых людей или работников на производстве.

• Анимация: Создание реалистичных движений для анимационных персонажей в фильмах и видеоиграх.

• Виртуальная и дополненная реальность: Реалистичное отслеживание движений пользователей для создания интерактивных VR и AR приложений.

• Анализ поведения клиентов: Изучение движения и поведения клиентов в магазинах для оптимизации выкладки товаров и улучшения обслуживания.

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

Пример кода для обработки изображения (а с помощью предыдущих примеров легко сделать обработку видео-файла или камеры):

from ultralytics import YOLO
import cv2
import numpy as np
import os

# Загрузка модели YOLOv8-Pose
model = YOLO('yolov8n-pose.pt')

# Словарь цветов для различных классов
colors = {
    'white': (255, 255, 255),
    'red': (0, 0, 255),
    'blue': (255, 0, 0)
}

def draw_skeleton(image, keypoints, confs, pairs, color):
    for (start, end) in pairs:
        if confs[start] > 0.5 and confs[end] > 0.5:
            x1, y1 = int(keypoints[start][0]), int(keypoints[start][1])
            x2, y2 = int(keypoints[end][0]), int(keypoints[end][1])
            if (x1, y1) != (0, 0) and (x2, y2) != (0, 0):  # Игнорирование точек в (0, 0)
                cv2.line(image, (x1, y1), (x2, y2), color, 2)

def process_image(image_path):
    # Загрузка изображения
    image = cv2.imread(image_path)
    if image is None:
        print("Ошибка: не удалось загрузить изображение")
        return

    # Обработка изображения с помощью модели
    results = model(image)[0]

    # Проверка на наличие обнаруженных объектов
    if hasattr(results, 'boxes') and hasattr(results.boxes, 'cls') and len(results.boxes.cls) > 0:
        classes_names = results.names
        classes = results.boxes.cls.cpu().numpy()
        boxes = results.boxes.xyxy.cpu().numpy().astype(np.int32)

        # Обработка ключевых точек
        if results.keypoints:
            keypoints = results.keypoints.data.cpu().numpy()
            confs = results.keypoints.conf.cpu().numpy()
            
            for i, (class_id, box, kp, conf) in enumerate(zip(classes, boxes, keypoints, confs)):
                draw_box=False
                if draw_box:
                    class_name = classes_names[int(class_id)]
                    color = colors['white']
                    x1, y1, x2, y2 = box
                    cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
                    cv2.putText(image, class_name, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

                # Визуализация ключевых точек с номерами
                for j, (point, point_conf) in enumerate(zip(kp, conf)):
                    if point_conf > 0.5:  # Фильтрация по уверенности
                        x, y = int(point[0]), int(point[1])
                        if (x, y) != (0, 0):  # Игнорирование точек в (0, 0)
                            cv2.circle(image, (x, y), 5, colors['blue'], -1)
                            cv2.putText(image, str(j), (x + 5, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, colors['blue'], 2)

                # Рисование скелета
                draw_skeleton(image, kp, conf, [(5, 7), (7, 9), (6, 8), (8, 10)], colors['white']) # Руки
                draw_skeleton(image, kp, conf, [(11, 13), (13, 15), (12, 14), (14, 16)], colors['red']) # Ноги
                draw_skeleton(image, kp, conf, [(5, 11), (6, 12)], colors['blue']) # Тело

    # Сохранение и отображение результатов
    output_path = os.path.splitext(image_path)[0] + "_pose_detected.jpg"
    cv2.imwrite(output_path, image)
    print(f"Сохранено изображение с результатами: {output_path}")

    cv2.imshow('YOLOv8-Pose Detection', image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

# Путь к изображению для обработки
image_path = 'd.jpg'
process_image(image_path)

Трекинг объектов

Трекинг объектов — это процесс записи и анализа перемещений объектов в видео или последовательности изображений. Этот процесс включает присвоение уникальных идентификаторов каждому обнаруженному объекту и отслеживание его местоположения с течением времени. Благодаря этому подходу можно решать множество задач в различных областях.

Применения трекинга объектов:

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

• Транспорт и логистика: Управление дорожным движением, мониторинг транспортных средств для оптимизации потока движения, автономные транспортные средства и предотвращение столкновений.

• Розничная торговля: Анализ поведения покупателей для оптимизации расположения товаров, улучшения пользовательского опыта и противокражные системы для мониторинга подозрительных действий.

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

• Медицина и здравоохранение: Реабилитация, отслеживание движений пациентов для мониторинга их прогресса, анализ походки и других движений для выявления нарушений и заболеваний.

• Робототехника и взаимодействие человек-компьютер: Навигация роботов, обеспечение безопасного передвижения роботов в динамической среде, жестовое управление устройствами и приложениями.

Пример кода для обработки видео:

from collections import defaultdict
import cv2
import numpy as np
from ultralytics import YOLO

# Загрузка модели YOLOv8
model = YOLO("yolov8x.pt")

# Открытие видео файла
video_path = "input.mp4"
cap = cv2.VideoCapture(video_path)

# Проверка успешного открытия видео
if not cap.isOpened():
    print(f"Ошибка открытия {video_path}")
    exit()


# Получение FPS видео
fps = cap.get(cv2.CAP_PROP_FPS)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

# Настройка VideoWriter для сохранения выходного видео
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter('out.mp4', fourcc, fps, (width, height))

# Создание словаря для хранения истории треков объектов
track_history = defaultdict(lambda: [])

# Цикл для обработки каждого кадра видео
while cap.isOpened():
    # Считывание кадра из видео
    success, frame = cap.read()

    if not success:
        print("Конец видео")
        break

    # Применение YOLOv8 для отслеживания объектов на кадре, с сохранением треков между кадрами
    results = model.track(frame, persist=True)

    # Проверка на наличие объектов
    if results[0].boxes is not None and results[0].boxes.id is not None:
        # Получение координат боксов и идентификаторов треков
        boxes = results[0].boxes.xywh.cpu()  # xywh координаты боксов
        track_ids = results[0].boxes.id.int().cpu().tolist()  # идентификаторы треков

        # Визуализация результатов на кадре
        annotated_frame = results[0].plot()

        # Отрисовка треков
        for box, track_id in zip(boxes, track_ids):
            x, y, w, h = box  # координаты центра и размеры бокса
            track = track_history[track_id]
            track.append((float(x), float(y)))  # добавление координат центра объекта в историю
            if len(track) > 30:  # ограничение длины истории до 30 кадров
                track.pop(0)

            # Рисование линий трека
            points = np.hstack(track).astype(np.int32).reshape((-1, 1, 2))
            cv2.polylines(annotated_frame, [points], isClosed=False, color=(230, 230, 230), thickness=10)

        # Отображение аннотированного кадра
        cv2.imshow("YOLOv8 Tracking", annotated_frame)
        out.write(annotated_frame)  # запись кадра в выходное видео
    else:
        # Если объекты не обнаружены, просто отображаем кадр
        cv2.imshow("YOLOv8 Tracking", frame)
        out.write(frame)  # запись кадра в выходное видео

    # Прерывание цикла при нажатии клавиши 'Esc'
    if cv2.waitKey(1) == 27:
        break

# Освобождение видеозахвата и закрытие всех окон OpenCV
cap.release()
out.release()  # закрытие выходного видеофайла
cv2.destroyAllWindows()

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

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

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

def predict_position(track, future_time, fps):
    if len(track) < 2:
        return track[-1]

    N = min(len(track), 25)
    track = np.array(track[-N:])

    times = np.arange(-N + 1, 1)

    A = np.vstack([times, np.ones(len(times))]).T
    k_x, b_x = np.linalg.lstsq(A, track[:, 0], rcond=None)[0]
    k_y, b_y = np.linalg.lstsq(A, track[:, 1], rcond=None)[0]

    future_frames = future_time * fps
    future_x = k_x * future_frames + b_x
    future_y = k_y * future_frames + b_y

Все тоже видео, но уже с обработкой предсказания:

Код для обработки видео с предсказанием

Создание собственного датасета

И теперь самое интересное — как же обучить нейросеть YOLOV8 обрабатывать не только автомобили и собачек, но и объекты с которыми необходимо работать. Для этого разработчиками предусмотрена возможность дообучить модель.
Почему дообучить? Потому что обучается не вся модель, а только последние слои нейронной сети. Это позволяет эффективно и быстро тренировать модель не ограниченном датасете, который реально разметить руками! Начнем!

Задача пройти полный путь и получить качественный результат, который в дальнейшем можно будет улучшать, увеличивая сложность объектов и условия их появления!
Представляю нашего героя - "Кубик" (который на самом деле прямоугольный параллелепипед). Именно его мы научим распознавать нашу модель!

— Сразу отвечу на вопрос: "А почему не OpenCV? Ведь ее тут достаточно".
— Да, более того, мы с помощью нее автоматически разметим датасет. А не ее, т.к. задача - обучить модель определять нужные объекты!

Снимаем видео с объектом, захватывая также кадры, где его нет, т.е. просто фон!

Следующим шагом разделяем видео на отдельные изображения

import cv2
import os
import time

# Путь к видеофайлу
video_path = '000.mp4'
# Папка для сохранения изображений
output_folder = 'output_images'
# Интервал между кадрами (каждый n-й кадр будет сохранен)
frame_interval = 1  # Можно изменить на 2, 5 и т.д.

os.makedirs(output_folder, exist_ok=True)

# Открытие видеофайла
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
    print(f"Ошибка открытия видеофайла: {video_path}")
    exit()

frame_count = 0
saved_count = 0

while True:
    ret, frame = cap.read()
    if not ret:
        break

    if frame_count % frame_interval == 0:
        # Получение текущего времени в виде временной метки
        timestamp = int(time.time() * 1000)  # Используем миллисекунды для большей точности
        output_path = os.path.join(output_folder, f'{timestamp}_frame_{saved_count:05d}.jpg')
        cv2.imwrite(output_path, frame)
        print(f"Сохранено: {output_path}")
        saved_count += 1

    frame_count += 1

cap.release()
print("Разделение видео на фотографии завершено.")

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

Теперь нужно разметить данные, т.е. указать места, где же находится искомый объект.

Автоматизированная разметка

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

import cv2
import numpy as np

def nothing(*arg):
    pass

cv2.namedWindow( "result" ) 
cv2.namedWindow( "settings" )


cv2.createTrackbar('h1', 'settings', 0, 180, nothing)
cv2.createTrackbar('s1', 'settings', 0, 255, nothing)
cv2.createTrackbar('v1', 'settings', 0, 255, nothing)
cv2.createTrackbar('h2', 'settings', 180, 180, nothing)
cv2.createTrackbar('s2', 'settings', 255, 255, nothing)
cv2.createTrackbar('v2', 'settings', 255, 255, nothing)

while True:
    img = cv2.imread('000.jpg')
    h,w,_=img.shape
    img=cv2.resize(img,(w//5,h//5))
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV )
 
    # считываем значения бегунков
    h1 = cv2.getTrackbarPos('h1', 'settings')
    s1 = cv2.getTrackbarPos('s1', 'settings')
    v1 = cv2.getTrackbarPos('v1', 'settings')
    h2 = cv2.getTrackbarPos('h2', 'settings')
    s2 = cv2.getTrackbarPos('s2', 'settings')
    v2 = cv2.getTrackbarPos('v2', 'settings')
    h_min = np.array((h1, s1, v1), np.uint8)
    h_max = np.array((h2, s2, v2), np.uint8)
    img_bin = cv2.inRange(hsv, h_min, h_max)
    cv2.imshow('result', img_bin)
    cv2.imshow('original', img)
    ch = cv2.waitKey(5)
    if ch == 27:
        break
cv2.destroyAllWindows()

Теперь необходимо записать полученные значения в формате:
lower_hsv = np.array([64, 54, 167])
upper_hsv = np.array([180, 255, 255])

Следующий скрипт применит данный фильтр ко всем изображениям и сохранит координаты найденного объекта.
ВАЖНО! Этот скрипт для решения простой задачи - где в кадре один объект одного класса, в хороших условиях. Для ручной разметки скрипт дальше!

import cv2
import os
import numpy as np

input_folder = 'output_images'
output_folder = 'dataset/train'
output_images_folder = os.path.join(output_folder, 'images')
output_labels_folder = os.path.join(output_folder, 'labels')

os.makedirs(output_images_folder, exist_ok=True)
os.makedirs(output_labels_folder, exist_ok=True)

lower_hsv = np.array([89, 71, 120])
upper_hsv = np.array([180, 255, 255])

def find_mask(image, lower_hsv, upper_hsv):
    hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    mask = cv2.inRange(hsv_image, lower_hsv, upper_hsv)
    return mask

def find_bounding_rect(mask):
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if contours:
        contour = max(contours, key=cv2.contourArea)
        x, y, w, h = cv2.boundingRect(contour)
        return x, y, x + w, y + h
    else:
        return None

def normalize_coordinates(x1, y1, x2, y2, img_width, img_height):
    x_center = (x1 + x2) / 2 / img_width
    y_center = (y1 + y2) / 2 / img_height
    width = abs(x2 - x1) / img_width
    height = abs(y2 - y1) / img_height
    return x_center, y_center, width, height

for filename in os.listdir(input_folder):
    if filename.endswith(('.jpg', '.jpeg', '.png')):
        image_path = os.path.join(input_folder, filename)
        image = cv2.imread(image_path)
        resized_image = cv2.resize(image, (640,640))
        mask = find_mask(resized_image, lower_hsv, upper_hsv)
        bounding_rect = find_bounding_rect(mask)

        if bounding_rect is not None: 
            x1, y1, x2, y2 = bounding_rect
            x_center, y_center, width, height = normalize_coordinates(x1, y1, x2, y2, 640, 640)

            # Сохранение изображения
            output_image_path = os.path.join(output_images_folder, filename)
            cv2.imwrite(output_image_path, resized_image)

            # Сохранение 
            label_filename = os.path.splitext(filename)[0] + '.txt'
            label_file_path = os.path.join(output_labels_folder, label_filename)
            with open(label_file_path, 'w') as f:
                f.write(f"0 {x_center} {y_center} {width} {height}\n")

            print(f"Processed and saved {filename}")

print("Подготовка датасета завершена.")

В процессе обработки файлы изменяются до размера 640*640 пикселей для передачи в нейросеть. Также для каждого файла записывается TXT файл, в котором содержится информация о расположении объекта.
Важно, что координаты указываются не в пикселях, а в виде значений от 0 до 1, указывая отношение местоположение точки относительно размера кадра.

Структура папок на данном этапе выглядит следующим образом:

dataset/
├── train/
│   ├── images/
│   │   ├── img1.jpg
│   │   ├── img2.jpg
│   │   ├── ...
│   ├── labels/
│   │   ├── img1.txt
│   │   ├── img2.txt
│   │   ├── ...
├── output_images/
│   ├── img1.jpg
│   ├── img2.jpg
│   ├── ...
├── Dataset_video2images.py
└── Dataset_HSV_Markup.py

После разметки хорошо бы проверить, а как же качественно были размечены данные!

import cv2
import os

# Папки с изображениями и метками
images_path = 'dataset/train/images'
labels_path = 'dataset/train/labels'

# Папка для сохранения изображений с нарисованными прямоугольниками
output_folder = 'checked_images'
os.makedirs(output_folder, exist_ok=True)

# Цвета для классов (можно добавить больше цветов, если классов больше)
colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0)]

# Чтение всех файлов изображений и меток
images = [f for f in os.listdir(images_path) if f.endswith(('.jpg', '.jpeg', '.png'))]

# Функция для преобразования координат из нормализованных значений в пиксели
def denormalize_coordinates(x_center, y_center, width, height, img_width, img_height):
    x_center *= img_width
    y_center *= img_height
    width *= img_width
    height *= img_height
    x1 = int(x_center - width / 2)
    y1 = int(y_center - height / 2)
    x2 = int(x_center + width / 2)
    y2 = int(y_center + height / 2)
    return x1, y1, x2, y2

# Обработка изображений
for image_file in images:
    image_path = os.path.join(images_path, image_file)
    label_file = os.path.splitext(image_file)[0] + '.txt'
    label_path = os.path.join(labels_path, label_file)

    # Проверка наличия файла меток
    if not os.path.exists(label_path):
        print(f"Label file not found for image: {image_file}")
        continue

    # Загрузка изображения
    image = cv2.imread(image_path)
    img_height, img_width = image.shape[:2]

    # Чтение файла меток и рисование прямоугольников
    with open(label_path, 'r') as f:
        for line in f:
            cls, x_center, y_center, width, height = map(float, line.strip().split())
            x1, y1, x2, y2 = denormalize_coordinates(x_center, y_center, width, height, img_width, img_height)
            color = colors[int(cls) % len(colors)]
            cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
            cv2.putText(image, f'class {int(cls)}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)

    # Сохранение изображения с нарисованными прямоугольниками
    output_path = os.path.join(output_folder, image_file)
    cv2.imwrite(output_path, image)
    print(f"Saved checked image: {output_path}")

print("Проверка разметки завершена.")

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

Ручная разметка

Для работы с более сложными данными, множеством классов требуется ручная разметка. Существуют, конечно, онлайн-сервисы, но это не наш метод. Пишем сами:

import cv2
import os
import yaml
import shutil

# Путь к папке с исходными изображениями
full_images_path = 'output_images'
# Путь к папке для сохранения обработанных данных
dataset_path = 'dataset'
train_images_path = os.path.join(dataset_path, 'train', 'images')
train_labels_path = os.path.join(dataset_path, 'train', 'labels')
valid_images_path = os.path.join(dataset_path, 'valid', 'images')
valid_labels_path = os.path.join(dataset_path, 'valid', 'labels')
test_images_path = os.path.join(dataset_path, 'test', 'images')
test_labels_path = os.path.join(dataset_path, 'test', 'labels')
ready_images_path = os.path.join(full_images_path, 'ready')

os.makedirs(train_images_path, exist_ok=True)
os.makedirs(train_labels_path, exist_ok=True)
os.makedirs(valid_images_path, exist_ok=True)
os.makedirs(valid_labels_path, exist_ok=True)
os.makedirs(test_images_path, exist_ok=True)
os.makedirs(test_labels_path, exist_ok=True)
os.makedirs(ready_images_path, exist_ok=True)


window_name = 'Annotation Tool'
current_class = 0
colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0)]
annotations = []

# Функция для масштабирования изображения
def resize_image(image, size=(640, 640)):
    return cv2.resize(image, size)

# Обработка событий мыши
drawing = False
ix, iy = -1, -1

def draw_rectangle(event, x, y, flags, param):
    global ix, iy, drawing, annotations, current_class
    
    if event == cv2.EVENT_LBUTTONDOWN:
        drawing = True
        ix, iy = x, y
    elif event == cv2.EVENT_MOUSEMOVE:
        if drawing:
            image = param['original_image'].copy()
            for annotation in annotations:
                cls, x1, y1, x2, y2 = annotation
                cv2.rectangle(image, (x1, y1), (x2, y2), colors[cls], 2)
            cv2.rectangle(image, (ix, iy), (x, y), colors[current_class], 2)
            cv2.imshow(window_name, image)
    elif event == cv2.EVENT_LBUTTONUP:
        drawing = False
        annotations.append((current_class, ix, iy, x, y))
        image = param['original_image'].copy()
        for annotation in annotations:
            cls, x1, y1, x2, y2 = annotation
            cv2.rectangle(image, (x1, y1), (x2, y2), colors[cls], 2)
        cv2.imshow(window_name, image)
    elif event == cv2.EVENT_RBUTTONDOWN:  # Удаление последней рамки
        if annotations:
            removed_annotation = annotations.pop()
            image = param['original_image'].copy()  # Вернемся к оригинальному изображению
            for annotation in annotations:
                cls, x1, y1, x2, y2 = annotation
                cv2.rectangle(image, (x1, y1), (x2, y2), colors[cls], 2)
            cv2.imshow(window_name, image)
        else:
            image = param['original_image'].copy()
            cv2.imshow(window_name, image)

# Обновление файла data.yaml
def update_data_yaml():
    data_yaml_path = 'data.yaml'
    data = {
        'train': 'dataset/train/images',
        'val': 'dataset/valid/images',
        'test': 'dataset/test/images',
        'nc': 4,
        'names': ['class0', 'class1', 'class2', 'class3']
    }
    with open(data_yaml_path, 'w') as f:
        yaml.dump(data, f, default_flow_style=None, sort_keys=False)
    print(f"Updated {data_yaml_path}")

# Загрузка и обработка изображений
for filename in os.listdir(full_images_path):
    if filename.endswith(('.jpg', '.jpeg', '.png')):
        image_path = os.path.join(full_images_path, filename)
        image = cv2.imread(image_path)
        image = resize_image(image)
        original_image = image.copy()
        annotations = []

        cv2.namedWindow(window_name)
        cv2.setMouseCallback(window_name, draw_rectangle, param={'original_image': original_image})

        while True:
            image_with_annotations = original_image.copy()
            for annotation in annotations:
                cls, x1, y1, x2, y2 = annotation
                cv2.rectangle(image_with_annotations, (x1, y1), (x2, y2), colors[cls], 2)
            cv2.imshow(window_name, image_with_annotations)
            key = cv2.waitKey(1) & 0xFF

            if key == ord(' '):  # Нажатие пробела для сохранения
                # Сохранение изображения
                output_image_path = os.path.join(train_images_path, filename)
                cv2.imwrite(output_image_path, original_image)
                print(f"Saved image to {output_image_path}")

                # Сохранение текстовых данных
                label_filename = os.path.splitext(filename)[0] + '.txt'
                label_file_path = os.path.join(train_labels_path, label_filename)
                with open(label_file_path, 'w') as f:
                    for annotation in annotations:
                        cls, x1, y1, x2, y2 = annotation
                        x_center = (x1 + x2) / 2 / 640
                        y_center = (y1 + y2) / 2 / 640
                        width = abs(x2 - x1) / 640
                        height = abs(y2 - y1) / 640
                        f.write(f"{cls} {x_center} {y_center} {width} {height}\n")
                print(f"Saved labels to {label_file_path}")

                # Обновление data.yaml
                update_data_yaml()

                # Перемещение обработанного изображения
                ready_image_path = os.path.join(ready_images_path, filename)
                shutil.move(image_path, ready_image_path)
                print(f"Moved image to {ready_image_path}")
                break
            elif key == 27:  # Нажатие Esc для пропуска изображения
                print("Skipped image")
                break
            elif key in [ord(str(i)) for i in range(10)]:  # Выбор класса
                current_class = int(chr(key))
                print(f"Selected class: {current_class}")

        cv2.destroyAllWindows()

Краткая инструкция:
— Выбираете один из 4 классов с помощью цифр 0,1,2,3 на клавиатуре
— Рисуем рамки вокруг объектов
— Удалить последнюю рамку - правая кнопка мыши
— Перейти к следующему кадру - пробел

Теперь, когда все данные размечены осталось подготовить несколько файлов и можно запускать обучение!

Финальные приготовления

Для качественного обучения и проверки работы необходимо разделить все данные на 3 части:
— train (обучающие данные) - 70%
— test (тестовые данные, для проверки во время обучения) - 20 %
— valid (проверочные данные, для тестирования модели после обучения) - 10 %

Вручную делать это неудобно, поэтому автоматизируем процесс с помощью скрипта:

import os
import shutil
import random

# Параметры для разделения данных
test_percent = 0.2  # Процент данных для тестирования
valid_percent = 0.1  # Процент данных для проверки

# Путь к папке с данными
dataset_path = 'dataset'
train_images_path = os.path.join(dataset_path, 'train', 'images')
train_labels_path = os.path.join(dataset_path, 'train', 'labels')
valid_images_path = os.path.join(dataset_path, 'valid', 'images')
valid_labels_path = os.path.join(dataset_path, 'valid', 'labels')
test_images_path = os.path.join(dataset_path, 'test', 'images')
test_labels_path = os.path.join(dataset_path, 'test', 'labels')

os.makedirs(valid_images_path, exist_ok=True)
os.makedirs(valid_labels_path, exist_ok=True)
os.makedirs(test_images_path, exist_ok=True)
os.makedirs(test_labels_path, exist_ok=True)

# Получение всех файлов изображений и соответствующих меток
images = [f for f in os.listdir(train_images_path) if f.endswith(('.jpg', '.jpeg', '.png'))]
labels = [f for f in os.listdir(train_labels_path) if f.endswith('.txt')]

# Убедимся, что количество изображений и меток совпадает
images.sort()
labels.sort()

# Проверка на соответствие количества изображений и меток
if len(images) != len(labels):
    print("Количество изображений и меток не совпадает.")
    exit()

# Перемешивание данных
data = list(zip(images, labels))
random.shuffle(data)
images, labels = zip(*data)

# Разделение данных
num_images = len(images)
num_test = int(num_images * test_percent)
num_valid = int(num_images * valid_percent)
num_train = num_images - num_test - num_valid

# Перемещение данных в соответствующие папки
def move_files(file_list, source_image_dir, source_label_dir, dest_image_dir, dest_label_dir):
    for file in file_list:
        image_path = os.path.join(source_image_dir, file)
        label_path = os.path.join(source_label_dir, os.path.splitext(file)[0] + '.txt')
        shutil.move(image_path, os.path.join(dest_image_dir, file))
        shutil.move(label_path, os.path.join(dest_label_dir, os.path.splitext(file)[0] + '.txt'))

# Перемещение тестовых данных
move_files(images[:num_test], train_images_path, train_labels_path, test_images_path, test_labels_path)

# Перемещение валидационных данных
move_files(images[num_test:num_test + num_valid], train_images_path, train_labels_path, valid_images_path, valid_labels_path)

# Оставшиеся данные остаются в папке train

print(f"Перемещено {num_test} изображений в папку test.")
print(f"Перемещено {num_valid} изображений в папку valid.")
print(f"Осталось {num_train} изображений в папке train.")

Еще один важный этап - создание файла data.yaml, с информацией и папках, файлах и названиях классов в будущей модели. Структура файла выглядит так:

train: dataset/train/images 
val: dataset/valid/images 
test: dataset/test/images

nc: 4 
names: [class0, class1, class2, class3]

Теперь все готово, и структура проекта выглядит следующим образом:

dataset/
├── train/
│   ├── images/
│   │   ├── img1.jpg
│   │   ├── img2.jpg
│   │   ├── ...
│   ├── labels/
│   │   ├── img1.txt
│   │   ├── img2.txt
│   │   ├── ...
├── test/
│   ├── images/
│   │   ├── img3.jpg
│   │   ├── img4.jpg
│   │   ├── ...
│   ├── labels/
│   │   ├── img3.txt
│   │   ├── img4.txt
│   │   ├── ...
├── valid/
│   ├── images/
│   │   ├── img5.jpg
│   │   ├── img6.jpg
│   │   ├── ...
│   ├── labels/
│   │   ├── img5.txt
│   │   ├── img6.txt
│   │   ├── ...
├── output_images/
│   ├── img1.jpg
│   ├── img2.jpg
│   ├── ...
├── data.yaml
├── Dataset_video2images.py
└── Dataset_HSV_Markup.py

Обучение

Всего несколько параметров, с которыми можно и нужно работать в самом начале:
epochs = 500 — количество эпох обучения. Выбирается индивидуально для задачи. При обучении YOLO останавливается автоматически, если улучшения результатов не происходит несколько эпох.

batch = 64 - размер "пакета изображений" передаваемый за один раз в нейросеть. Изменяется в зависимости от количества доступной памяти.

imgsz = 640 - размер изображений

import os
from ultralytics import YOLO

current_dir = os.path.dirname(os.path.abspath(__file__))

data_path = os.path.join(current_dir, 'data.yaml')
model = YOLO(os.path.join(current_dir, 'yolov8n.pt'))


epochs = 500
batch = 64
imgsz = 640
if __name__ == '__main__':
    results = model.train(data=data_path,
                      epochs=epochs, 
                      batch=batch, 
                      imgsz=imgsz, 
                      name='red',
                      device='cuda')

Запускаем обучение и ждем!

После завершения процесса (время сильно может отличаться в зависимости от множества параметров) в папке runs будут лежать графики и прочие данные процесса обучения, а также файлы модели best.pt и last.pt

Осталось только в скрипте для детекции заменить файл модели на best.pt и УРА, можно пользоваться!