Vindicar

"""
Принимаем видео с вебкамеры, и пытаемся сегментировать его с помощью YOLOv8.
Выводим в окне маски объектов.
Потребуется установить пакеты следующей командой:
pip install ultralytics
"""
from sys import argv
from pathlib import Path

import numpy
import cv2
import torch

COLORS = [
    (64, 128, 128),
    (128, 64, 64),
    (64, 128, 64),
    (64, 64, 128),
    (128, 64, 128),
]

VIDEO_SOURCE = 0  # 0 - вебкамера. Строка - имя файла или URL видео потока.
script_dir = Path(argv[0]).parent.resolve()
# определяем, на каком устройстве будет выполняться модель
device = (
    'cuda' if torch.cuda.is_available() else
    'mps' if torch.backends.mps.is_available() else
    'cpu'
)

import ultralytics


def text_at(  # рисует текст по центру
        img: numpy.ndarray,  # на чём
        text: str,  # что
        pos: tuple[int, int],  # где центр
        font=cv2.FONT_HERSHEY_COMPLEX,  # каким шрифтом
        font_scale: float = 1.0,  # в каком масштабе (не кегль!)
        font_thickness: int = 2,  # насколько толстые линии
        text_color=(255, 0, 0),  # каким цветом
        bg_color=None  # на каком фоне
        ):
    x, y = pos
    (text_w, text_h), _ = cv2.getTextSize(text, font, font_scale, font_thickness)
    if bg_color is not None:
        cv2.rectangle(img,
                      (x - text_w // 2, y - text_h // 2),
                      (x + text_w // 2, y + text_h // 2),
                      bg_color, -1)
    cv2.putText(img,
                text,
                (x - text_w // 2, y + text_h // 2 - 1),
                font, font_scale, text_color, font_thickness)


model = ultralytics.YOLO(
    script_dir / 'yolov8m-seg.pt'  # имя файла модели указывает на её тип
)

# цикл работы с видео
video = cv2.VideoCapture(VIDEO_SOURCE)
if not video.isOpened():
    raise SystemExit('Не удалось открыть источник видео')
overlay_image: numpy.ndarray = numpy.zeros(  # изображение-оверлей для отображения разметки кадра
    (
        int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)),
        int(video.get(cv2.CAP_PROP_FRAME_WIDTH)),
        3
     ),
    numpy.uint8
)

try:
    while True:
        success, frame = video.read()  # читаем кадр
        if not success:
            raise SystemExit('Видео закончилось')
        # для YOLO предобработка не требуется
        outputs = model.predict(  # получаем отклик сети
            source=frame,
            conf=0.15,  # минимальная степень уверенности в объекте
            device=device,  # устройство
            verbose=False  # чтобы не спамило логом в консоль
        )
        output = outputs[0]  # модель всегда возвращает список результатов, даже для одного изображения
        if output.masks is not None:  # если были получены маски
            masks = output.masks.cpu().xy  # список масок как координат вершин многоугольников (контуров)
            boxes = output.boxes.cpu()  # ограничивающие прямоугольники
            image_output = []  # список аннотаций, которые нужно будет вывести
            for box, mask in zip(boxes, masks):  # формируем список аннотаций
                class_id = int(box.cls.item())  # номер класса, соответствующего очередной области
                class_name = model.names[class_id]  # имя класса, соответствующее области
                pts = numpy.array(mask).astype(numpy.int32)  # массив координат точек контура Nx2
                image_output.append((class_id, class_name, pts))
            overlay_image.fill(0)  # очищаем изображение-оверлей
            for (class_id, class_name, pts) in image_output:  # закрашиваем каждую маску цветом
                cv2.fillPoly(
                    img=overlay_image,
                    pts=pts[numpy.newaxis, ...],  # для работы fillPoly() массив должен быть вида 1xNx2
                    color=COLORS[class_id % len(COLORS)],  # цвет массива определяем по номеру класса для стабильности
                )
            for (class_id, class_name, pts) in image_output:  # выводим надписи отдельным циклом, чтобы их не закрасило
                M = cv2.moments(pts)  # вычисляем моменты контура
                # используем их для расчёта координат центроида контура
                cX = int(M["m10"] / M["m00"])
                cY = int(M["m01"] / M["m00"])
                # выводим имя класса
                text_at(overlay_image, class_name,
                        (cX, cY),
                        text_color=(255, 255, 255),
                        bg_color=(1, 1, 1))
            # накладываем оверлей с разметкой на кадр
            alpha = 0.8  # вес оверлея, чем меньше - тем он прозрачнее.
            apply = overlay_image > 0  # изменяем только те части кадра, где есть оверлей
            frame[apply] = alpha * overlay_image[apply] + (1 - alpha) * frame[apply]  # обожаю numpy
        cv2.imshow('Press Esc to exit', frame)  # показываем результат
        if cv2.waitKey(10) == 27:  # если пользователь нажал Esc, выходим
            break
finally:
    video.release()  # не забываем отключиться от источника видео в итоге