ФАЙЛ "YOLO.PY"

Данный код использует модель YOLOv8n для обнаружения объектов в видеопотоке и передает данные о них (bounding box, уверенность, классы) на сервер в формате JSON. В случае ошибки соединение с сервером закрывается корректно.

ИМПОРТ БИБЛИОТЕК

import sys
import socket
import json
from ultralytics import YOLO
  • sys: Импортируется для работы с аргументами командной строки.

  • socket: Импортируется для создания и управления сетевыми соединениями.

  • json: Импортируется для работы с данными в формате JSON.

  • YOLO из ultralytics: Импортируется для использования модели YOLOv8n.

ИНИЦИАЛИЗАЦИЯ МОДЕЛИ YOLOv8n

model = YOLO('yolov8n.pt')

Загружается предварительно обученная модель YOLOv8n из файла yolov8n.pt.

ПОДКЛЮЧЕНИЕ К СЕРВЕРУ

if len(sys.argv) > 1:
    port_data = sys.argv[-1]
    port_data = port_data.split(',')
    ports = port_data[0].split(':')
    rtsp_url = port_data[1]
    socks = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    socks.connect((ports[0], int(ports[1])))
else:
    print("данных нет")
    exit()

Проверяется, есть ли аргументы командной строки. Если аргументы есть, извлекаются данные порта и RTSP URL из последнего аргумента. Создается сокет и подключается к серверу по указанному IP-адресу и порту. Если аргументов нет, выводится сообщение "данных нет" и программа завершается.

ОСНОВНАЯ ЧАСТЬ КОДА

try:
    results = model(rtsp_url, stream=True, verbose=False)
    for result in results:
        boxes = result.boxes.xyxy.tolist()
        conf = result.boxes.conf.tolist()
        YOLO_cls = result.boxes.cls.tolist()

        data = {
            'boxes': boxes,
            'conf': conf,
            'class': YOLO_cls
        }
        if boxes == [] and conf == [] and YOLO_cls == []:
            data = None

        json_data = json.dumps(data)
        socks.sendall(json_data.encode('utf-8'))

except Exception as e:
    print(f"Произошла ошибка: {e}")

finally:
    socks.close()

В блоке try:

  • Запускается обработка видеопотока с использованием модели YOLO.

  • Для каждого кадра извлекаются данные о bounding box (boxes), уверенности (conf) и классах (cls).

  • Данные упаковываются в словарь data.

  • Если данные отсутствуют (boxes, conf, cls пусты), словарь data устанавливается в None.

  • Данные конвертируются в JSON-строку и отправляются на сервер.

В блоке except:

  • Если возникает ошибка, она выводится в консоль.

В блоке finally:

  • Сокет закрывается после завершения работы или при возникновении ошибки.

Процесс подготовки кода:

  1. Создайте файл и назовите yolo.py

  1. Поэтапно скопируйте код в файл

  1. Установка библиотек: в данном случае не предустановлена библиотека ultralytics. Если она у вас не скачана, следуйте этим шагам:

    1. Откройте терминал VS Code (Terminal -> New Termial или комбинацией Ctrl + Shift + `)

    2. В терминале поочередно введите и запустите следующие команды: pip install ultralytics

Last updated