Поиск сущностей на изображении с помощью YOLO object detection в Google Colab

Благодаря Python библиотеке YOLO object detection можно производить поиск различных сущностей на изображении (например, машины, велосипеды и т.д.), при этом используя минимум кода и получая результат в удобном виде.

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

Ссылка на блокнот Gooogle Colab: https://colab.research.google.com/drive/1DSSavMhfjskGrR-RU1hV5mcy1JtzyxsP?usp=sharing

Сначала устанавливаем модуль «ultralytics»:

# Устанавливаем ultralytics
!pip install ultralytics

Далее подключаем необходимые модули, а также даем доступ к Google Диску:

# Подключаем необходимые модули
from ultralytics import YOLO
import cv2
import numpy as np
import os
import sys
# Подключаем Google Диск
from google.colab import drive
drive.mount('/content/drive')

Затем загружаем модель, в примере загружается базовая «yolov8n» (она загружается автоматически), но можно и загрузить свою обученную (тогда необходимо предварительно загрузить ее в папку на диске и указать путь к ней). И также в данном блоке устанавливается функция, отвечающая за распознавание объектов на изображении:

# Загрузка модели YOLOv8
model = YOLO('yolov8n.pt')
# model = YOLO('/content/drive/MyDrive/ml/detection/yolo/train2/my_14-10-2025_250.pt')

# Список цветов для различных классов
colors = [
    (255, 0, 0)
]

# Функция для обработки изображения
def process_image(image_path):
    # Загрузка изображения
    image = cv2.imread(image_path)
    results = model(image)[0]
    # results = model.predict(image, save=True, imgsz=640, conf=0.5)[0]
    # print(results)


    # Получение оригинального изображения и результатов
    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"Обработка {image_path}:")
    print(f"Результат сохранен в {new_image_path}")
    # print(f"Лог сохранен в {text_file_path}")

В финале запускаем функцию распознавания, указав путь к файлу распознаваемого изображения на Google Диске (где «/content/drive/MyDrive/» — это корень вашего Google Диска и далее от него путь до файла):

# запускаем распознавание
process_image('/content/drive/MyDrive/ml/detection/yolo/train2/test1.jpg')

В результате данный пример сформирует новый файл изображения с именем «_yolo» в этой же папке, где лежит и исходное, на котором будут выделены b подписаны распознанные объекты:

Также можно включить сохранение данных распознавания в текстовый лог (для этого в коде есть закомментированный блок). Обратите внимание, что если YOLO не найдет распознаваемое изображение — то возьмет произвольное из своей библиотеки.