YOlO实现目标检测-单目标

前言

在配置了YOLO和Pytorch后，利用yolo自带的yolov5s.pt轻量化模型实现目标检测

代码实现

图片实现，多目标：

import time
import torch
import cv2 as cv


class MultipleTarget:

    def __init__(self):
        """
        初始化
        """
        # 加载训练模型
        self.model = torch.hub.load('./yolov5', 'custom', path='./weight/yolov5s.pt', source='local')
        # 设置阈值
        self.model.conf = 0.52  # confidence threshold (0-1)
        self.model.iou = 0.45  # NMS IoU threshold (0-1)
        # # 加载摄像头
        # self.cap = cv.VideoCapture(0)

    def draw(self, list1, image_temp):

        for temp in list1:
            name = temp[6]      # 取出标签名
            temp = temp[:4].astype('int')   # 转成int
            # print(temp)
            # print('***************')
            # cv.rectangle(image_temp, (int(temp[0]), int(temp[1])), (int(temp[2]), int(temp[3])),
            #              (0, 0, 255), 3)  # 框出识别物体
            cv.rectangle(image_temp, (temp[0], temp[1]), (temp[2], temp[3]), (0, 0, 255), 3)  # 框出识别物体
            cv.putText(image_temp, name, (int(temp[0]-10), int(temp[1]-10)), cv.FONT_ITALIC, 1, (0, 255, 0), 2)

    def detect(self, image_temp):
        """
        目标检测
        """
        img = image_temp
        # Inference
        results = self.model(img)
        # Results
        # 显示相关结果信息，如：图片信息、识别速度...
        results.print()  # or .show(), .save(), .crop(), .pandas(), etc.
        pd1 = results.xyxy[0]  # 识别结果用tensor保存，包含标签、坐标范围、IOU
        pd = results.pandas().xyxy[0]   # tensor 转为 DataFrame(Pandas中的数据类型)

        person_list = pd[pd['name'] == 'person'].to_numpy()
        bus_list = pd[pd['name'] == 'bus'].to_numpy()

        self.draw(person_list, img)
        self.draw(bus_list, img)

        # 控制台显示
        # print(pd1)
        # print(pd)
        print(person_list)
        print(results.xyxy[0])
        print('----------------------------------------------------------')
        print(results.pandas().xyxy[0])
        # 窗口显示s
        cv.imshow('results', img)
        cv.waitKey(0)
        cv.destroyAllWindows()


image = cv.imread(r'yolov5/data/images/bus.jpg')
test = MultipleTarget()
test.detect(image)

摄像头获取，多目标：

import time
import torch
import cv2 as cv


class MultipleTarget:

    def __init__(self):
        """
        初始化
        """
        # 加载训练模型
        self.model = torch.hub.load('./yolov5', 'custom', path='./weight/yolov5s.pt', source='local')
        # 设置阈值
        self.model.conf = 0.52  # confidence threshold (0-1)
        self.model.iou = 0.45  # NMS IoU threshold (0-1)
        # 加载摄像头
        self.cap = cv.VideoCapture(0)
        if not self.cap.isOpened():
            print("Cannot open camera")
            exit()

    def draw(self, list_temp, image_temp):
        for temp in list_temp:
            name = temp[6]      # 取出标签名
            temp = temp[:4].astype('int')   # 转成int加快计算
            cv.rectangle(image_temp, (temp[0], temp[1]), (temp[2], temp[3]), (0, 0, 255), 3)  # 框出识别物体
            cv.putText(image_temp, name, (int(temp[0]-10), int(temp[1]-10)), cv.FONT_ITALIC, 1, (0, 255, 0), 2)

    def detect(self):
        """
        目标检测
        """
        while True:
            ret, frame = self.cap.read()
            # 如果正确读取帧，ret为True
            if not ret:
                print("Can't receive frame (stream end?). Exiting ...")
                break
            frame = cv.flip(frame, 1)

            # FPS计算time.start
            start_time = time.time()

            # Inference
            results = self.model(frame)
            pd = results.pandas().xyxy[0]
            # 取出对应标签的list
            person_list = pd[pd['name'] == 'person'].to_numpy()
            bus_list = pd[pd['name'] == 'bus'].to_numpy()
            # 框出物体
            self.draw(person_list, frame)
            self.draw(bus_list, frame)
            # end_time
            end_time = time.time()
            fps = 1 / (end_time - start_time)

            # 控制台显示
            results.print()  # or .show(), .save(), .crop(), .pandas(), etc.
            print(results.xyxy[0])  # img1 predictions (tensor)
            print('----------------')
            print(results.pandas().xyxy[0])  # img1 predictions (pandas)

            # FPS显示
            cv.putText(frame, 'FPS:' + str(int(fps)), (30, 50), cv.FONT_ITALIC, 1, (0, 255, 0), 2)

            cv.imshow('results', frame)
            cv.waitKey(10)
            if cv.waitKey(10) & 0xFF == ord('q'):
                break

        self.cap.release()
        cv.destroyAllWindows()


test = MultipleTarget()
test.detect()

赏

谢谢观看