python Обнаружение YOLOv4 в реальном времени с экрана рабочего стола

Есть простой код использования Yolo для вывода видеоролика либо камеры с ПК:

import cv2
import time

Conf_threshold = 0.4
NMS_threshold = 0.4
COLORS = [(0, 255, 0), (0, 0, 255), (255, 0, 0),
          (255, 255, 0), (255, 0, 255), (0, 255, 255)]

class_name = []
with open('Resources\coco.names.txt', 'r') as f:
    class_name = [cname.strip() for cname in f.readlines()]
# print(class_name)
net = cv2.dnn.readNet('Resources\yolov4-tiny.weights', 'Resources\yolov4-tiny.cfg')
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA_FP16)

model = cv2.dnn_DetectionModel(net)
model.setInputParams(size=(416, 416), scale=1/255, swapRB=True)

cap = cv2.VideoCapture('test.mp4')

starting_time = time.time()
frame_counter = 0
while True:
    ret, frame = cap.read()
    frame_counter += 1
    if ret == False:
        break
    classes, scores, boxes = model.detect(frame, Conf_threshold, NMS_threshold)
    for (classid, score, box) in zip(classes, scores, boxes):
        color = COLORS[int(classid) % len(COLORS)]
        label = "%s : %f" % (class_name[classid], score)
        cv2.rectangle(frame, box, color, 1)
        cv2.putText(frame, label, (box[0], box[1]-10),
                   cv2.FONT_HERSHEY_COMPLEX, 0.5, color, 1)
    endingTime = time.time() - starting_time
    fps = frame_counter/endingTime
    # print(fps)
    cv2.putText(frame, f'FPS: {fps}', (5, 35),
               cv2.FONT_HERSHEY_COMPLEX, 0.7, (0, 255, 0), 2)
    cv2.imshow('frame', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        cap.release()
        cv2.destroyAllWindows()

Но мне нужно чтобы вместо видеоролика test.mp4 был захват экрана. Для создания скриншотов с экрана использовался этот код:

import numpy as np
import cv2
from mss import mss

def screen_record_efficiency():
    bbox = {'top': 0, 'left': 0, 'width': 800, 'height': 600}

    sct = mss()

    font = cv2.FONT_HERSHEY_SIMPLEX

    while 1:
        # grab image
        sct_img = np.array(sct.grab(bbox))

        # display image
        cv2.imshow('screen', sct_img)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            cv2.destroyAllWindows()
            break

Не могу понять каким образом можно вместо видеофайла загрузить массив с изображениями со второго кода и как это вообще объединить


Ответы (0 шт):