Начало работы с YOLO: просто и понятно

Начало работы с YOLO: просто и понятно

Если вам кажется, что освоить работу с нейросетями сложно, эта статья для вас! Сегодня мы познакомим с одной из самых популярных моделей компьютерного зрения — YOLO (You Only Look Once), которая делает работу с объектами на изображениях доступной даже для тех, у кого минимальные навыки программирования на Python. Подробно про эффектив­ный алго­ритм и особенности его работы нам расскажет основатель школы робототехники RobotX и автор курсов по технологии компьютерного зрения на Stepik — Степан Бурмистров.

Что такое YOLO и какие задачи она решает?

YOLO — это нейронная сеть, которая используется для работы с изображениями и решает следующие задачи:

  • Детекция объектов — обнаружение объектов и их местоположения на изображении.
  • Сегментация — разделение изображения на области, относящиеся к каждому объекту.
  • Классификация — определение того, что изображено на картинке.
  • Поиск ключевых точек тела — определение позы человека.
  • Трекинг объектов — потоковая обработка, позволяющая отслеживать и сохранять историю перемещений объектов.

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

Легкий старт с YOLO: подготовка и установка

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

pip install ultralytics

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

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

Начало работы с YOLO: просто и понятно

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

from ultralytics import YOLO import cv2 import numpy as np import os # Загрузка модели YOLOv8 model = YOLO('yolov8n.pt') # Список цветов для различных классов colors = [ (255, 0, 0), (255, 0, 0), (255, 255, 0), (255, 255, 0), (255, 255, 255), (255, 0, 255), (192, 192, 192), (128, 128, 128), (128, 0, 128), (128, 128, 0), (0, 128, 0), (128, 0, 0), (128, 128, 128), (0, 128, 128), (72, 61, 139), (47, 79, 79), (47, 79, 47), (0, 206, 209), (148, 0, 211), (255, 20, 147) ]

В первой части кода выбирается модель. YOLO автоматически скачает её при первом запуске скрипта. Есть несколько моделей на выбор:

  • yolov8n.pt
  • yolov8s.pt
  • yolov8m.pt
  • yolov8l.pt
  • yolov8x.pt

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

# функция для обработки изображения def process_image(image_path): # загрузка изображения image = cv2.imread(image_path) results = model(image) # Получение оригинального изображения и результатов 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.8, color, 2) new_image_path = os.path.split(image_path)[0] + '_yolo/' + os.path.split(image_path)[1] cv2.imwrite(new_image_path, image) # Cохранение объектов в текстовый файл 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(): for detail in details: f.write(f"Coordinates: {detail[0]}, {detail[1]}, {detail[2]}, {detail[3]}\n") print(f"Processed image: {image_path}") print(f"Saved bounding-box image to {new_image_path}") print(f"Saved data to {text_file_path}") process_image('test.png')

После запуска скрипта нейросеть определит объекты на изображении и вернёт их координаты и классы (например, "person", "car", "bus", "traffic light"). Итак, после выполнения данного скрипта видим результат:

Начало работы с YOLO: просто и понятно

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

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)

Такой формат позволяет легко работать с координатами объектов и использовать их для дальнейшего анализа или обработки.

Полный код можно найти по этой ссылке.

Дополнительные возможности YOLO

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

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

Бонус: При покупке курса до 10 ноября используйте промокод YOLODISCOUNT и получите скидку 15%!

Не упустите шанс начать изучение компьютерного зрения и нейросетей вместе с нами! Следите за обновлениями на нашем блоге и платформе Stepik — впереди ещё много интересных статей из этой серии. Учитесь, развивайтесь и становитесь профессионалами!

Начать дискуссию