Если вам кажется, что начать работу с нейросетями - это сложно, то этот материал для вас!
Итак, YOLO (You Only Look Once) — нейронная сеть, предназначенная работы с объектами на изображениях и может решать следующие задачи:
Детекция - обнаружение объектов
Сегментация - разделение изображения на области, которые относятся к каждому объекту
Классификация - определение что же находится на изображении
Поиск ключевых точек тела - для определения позы человека
Трекинг объектов - потоковая обработка, при которой для каждого объекта возможно сохранять и использовать историю местоположения
Также в этой статье также будет рассмотрено:
Предсказание движения на основе трекинга
Создание собственного датасета для дообучения модели детекции новых объектов
Отличительной особенностью 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
Процесс обработки изображения, при котором все изображение целиком будет относиться к определенному классу.
Детекция и классификация решают разные задачи и имеют свои особенности и области применения. Вот основные причины, почему не всегда следует использовать детекцию вместо классификации:
Сложность задачи:
Классификация: Определяет, к какому классу относится весь объект или изображение целиком. Например, классификация фотографии как "собака" или "кошка".
Детекция: Находит объекты внутри изображения и определяет их классы и местоположение (например, где на фотографии находится собака и где кошка).
Ресурсы и вычислительная мощность:
Классификация: Обычно требует меньше вычислительных ресурсов, так как анализируется весь объект или изображение целиком без необходимости определения его частей.
Детекция: Более вычислительно затратна, так как требует анализа изображения для поиска объектов и определения их границ.
Скорость выполнения:
Классификация: Быстрее, так как выполняется одна операция определения класса для всего изображения.
Детекция: Медленнее, так как требует многократного анализа изображения для поиска всех объектов и их классификации.
Сложность реализации:
Классификация: Проще в реализации и настройке, так как обучается на меньших и более структурированных данных.
Детекция: Более сложна в реализации, требует больше данных и времени на обучение, особенно если нужно обнаруживать объекты разного размера и формы.
Целостное определение класса объекта: Если нужно определить класс всего объекта или изображения, а не его частей. Например, определить, что изображено на фотографии (собака или кошка).
Ограниченные вычислительные ресурсы: В условиях, когда ограничены ресурсы для вычислений и требуется быстрая обработка данных.
Более точное определение "подклассов". Например, можно находить буквы на изображении с помощью детекции, а затем более точно определять символ при помощи классификации!
Множественные объекты на изображении: Если на изображении может быть несколько объектов разных классов, и нужно определить их местоположение и классы. Например, обнаружение автомобилей и пешеходов на улице.
Анализ сложных сцен: Когда нужно анализировать сложные сцены, где важно не только определить, какие объекты присутствуют, но и где они находятся.
Применение в реальном времени: В задачах, где необходимо отслеживать объекты в реальном времени, например, в системах видеонаблюдения.
Пример кода для одного изображения:
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
и УРА, можно пользоваться!