utils/load_input.py

import cv2
from ultralytics import YOLO
import torch


# 读取摄像头（返回摄像头）
def load_camera(num="http://192.168.1.3:8080",
                width=1920,
                height=1080):
    cap = cv2.VideoCapture(num)

    # 设置摄像头参数
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, width)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)

    if cap.isOpened():
        return cap
    else:
        return None


# 读取视频
def load_video(video_path):
    video = cv2.VideoCapture(video_path)
    return video


# 加载模型
def load_model(model_path,
               yaml_path=None,
               task='segment'):
    # 加载（改动过结构）的模型
    if yaml_path:
        model = YOLO(yaml_path, task=task).load(model_path)
    else:
        model = YOLO(model_path, task=task)
    return model


# 加载模型到图形计算设备
def load_device(model):
    # 获取图形计算设备信息
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

    # 将模型切换到图形计算设备上
    model.to(device)

    return model


# 将摄像头设置成生成器
def capture_frames(num=0):
    cap = load_camera(num)
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        yield frame
    cap.release()


# 使用生成器形成视频流
def generate_capture_frames():
    for frame in capture_frames():
        ret, buffer = cv2.imencode('.jpg', frame)
        frame = buffer.tobytes()
        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')