Yolov5 implements target detection system (using thread pool)

1. Introduction

The system contains the following parts:

1. Read data from camera
2. The main thread outputs data to the window
3. The background thread completes the computer vision target detection process and shares the results with the main thread
4. The main thread draws the results on the output image

Download yolvo5:

The source code of YOLOv5 is placed on Github at: https://github.com/ultralytics/yolov5

I use yolov5s.pt, yolvo5 learning reference target detection YOLOv5 introduction

Install the required modules for Yolov5

In the command line,
enter "pip install -r requirements.txt" and wait for the installation to complete. (The pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple command will be faster!!)

2. Steps

1. Import the library

import cv2 as cv
import torch
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from threading import Thread

2. Turn on the camera

cap = cv.VideoCapture(0)

3. Window output

def show_frame(self):
        print("thread1 open!")
        while True:
            self.ret, self.frame = cap.read()
            if self.state:
                cv.imshow("res", self.a)
                if cv.waitKey(10) & 0xFF == ord('q'):# 按q退出
                    break
        self.ret = False
        self.thread_state = False
        cap.release()
        cv.destroyAllWindows()
        print("t1 finished")

4. Computer vision target detection reasoning process

def process_frame(self):
        print("thread2 open!")
        model = torch.hub.load(self.path, "custom",
                               self.pt_path, source="local")
        while True:
            if self.ret:
                res = model(self.frame)
                res.render()  # 打标签
                self.a = res.ims[0]
                self.state = True
            elif self.thread_state is False:
                break
        print("t2 finished")

5. Use thread pool management

The thread pool helps you manage threads. You no longer need to create a thread for each task to process the task. When a task needs to be executed, it will apply for a thread from the thread pool. Some threads in the thread pool will be used to execute the task. If not, it will wait. Others will release the thread after the execution of the task is completed, and the waiting task can use the released thread to execute. Operated. For thread pool usage, please refer to How to use python thread pool

The concurrent.futures module provides a highly encapsulated asynchronous calling interface
ThreadPoolExecutor: thread pool, providing asynchronous calling

submit(fn, *args, **kwargs) submits tasks asynchronously

#使用类来实现多线程通信
t1 = Mythread()
#创建一个包含2条线程的线程池
pool = ThreadPoolExecutor(max_workers=2)  # 定义两个线程
#向线程池提交2个任务
thread1 = pool.submit(t1.show_frame)
thread2 = pool.submit(t1.process_frame)

In this way, we can read data from the camera; the main thread outputs the data to the window, the background thread completes the computer vision target detection process, and shares the results with the main thread, and the main thread draws the results on the output image.

---

3. Complete code

from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from threading import Thread

import cv2 as cv
import torch

cap = cv.VideoCapture(0)


class Mythread(Thread):
    def __init__(self):
        Thread.__init__(self)
        cap = cv.VideoCapture(0)
        self.ret = {
    
    }
        self.frame = {
    
    }
        self.a = {
    
    }
        self.state = False
        self.thread_state = True
        self.project = Path(__file__).resolve().parent.parent
        self.path = self.project.joinpath("yolov5-master")
        self.pt_path = self.path.joinpath("yolov5s.pt")

    def process_frame(self):
        print("thread2 open!")
        model = torch.hub.load(self.path, "custom",
                               self.pt_path, source="local")
        while True:
            if self.ret:
                res = model(self.frame)
                res.render()  # 打标签
                self.a = res.ims[0]
                self.state = True
            elif self.thread_state is False:
                break
        print("t2 finished")

    def show_frame(self):
        print("thread1 open!")
        while True:
            self.ret, self.frame = cap.read()
            if self.state:
                cv.imshow("res", self.a)
                if cv.waitKey(10) & 0xFF == ord('q'):
                    break
        self.ret = False
        self.thread_state = False
        cap.release()
        cv.destroyAllWindows()
        print("t1 finished")


if __name__ == '__main__':
    t1 = Mythread()
    pool = ThreadPoolExecutor(max_workers=2)  # 定义两个线程
    thread1 = pool.submit(t1.show_frame)
    thread2 = pool.submit(t1.process_frame)

The above content is my personal summary. If you have any questions, please feel free to chat privately. Thank you! !