yolo系列简化detect函数

yolo系列简化detect函数
1、封装类

import os.path
import argparse
import numpy as np
import torch
from numpy import random

from models.experimental import attempt_load
from utils.datasets import letterbox
from utils.general import non_max_suppression, scale_coords, xyxy2xywh, check_img_size
from utils.plots import plot_one_box
from utils.torch_utils import time_synchronized


class YOLO(object):
    _defaults = {
        # ----------------------------------------------------------------------------#
        # 事先训练完成的权重文件，比如yolov5s.pt,假如使用官方训练好的文件（比如yolov5s）,则会自动下载
        # ----------------------------------------------------------------------------#
        "weights": "weights/yolov7.pt",
        # ----------------------------------------------------------------------------#
        # 预测时的放缩后图片大小(因为YOLO算法需要预先放缩图片)
        # ----------------------------------------------------------------------------#
        "imgsz": 640,
        # ----------------------------------------------------------------------------#
        # use FP16 half-precision inference 是否使用半精度推理（节约显存）
        "half": False,
        # ----------------------------------------------------------------------------#
        # 置信度阈值, 高于此值的bounding_box才会被保留
        # ----------------------------------------------------------------------------#
        "conf_thres": 0.25,
        # ----------------------------------------------------------------------------#
        # IOU阈值,高于此值的bounding_box才会被保留
        # ----------------------------------------------------------------------------#
        "iou_thres": 0.45,
        # ----------------------------------------------------------------------------#
        # augmented inference
        # ----------------------------------------------------------------------------#
        "augment": None,
        # ----------------------------------------------------------------------------#
        # 过滤指定类的预测结果
        # ----------------------------------------------------------------------------#
        "classes": None,
        # ----------------------------------------------------------------------------#
        # 如为True,则为class-agnostic. 否则为class-specific
        # ----------------------------------------------------------------------------#
        "agnostic_nms": False,
        # ----------------------------------------------------------------------------#
        # txt的保存路径
        # ----------------------------------------------------------------------------#
        "txt_path": None,
        # ----------------------------------------------------------------------------#
        # 是否保存结果到txt
        # ----------------------------------------------------------------------------#
        "save_txt": False,
        # ----------------------------------------------------------------------------#
        # 是否保存置信度
        # ----------------------------------------------------------------------------#
        "save_conf": False,
        # ----------------------------------------------------------------------------#
        # 是否展示图片
        # ----------------------------------------------------------------------------#
        "view_img": True
        # ----------------------------------------------------------------------------#
    }

    @classmethod
    def get_defaults(cls, n):
        if n in cls._defaults:
            return cls._defaults[n]
        else:
            return "Unrecognized attribute name '" + n + "'"
    def __init__(self, **kwargs):
        self.__dict__.update(self._defaults)
        for name, value in kwargs.items():
            setattr(self, name, value)
            self._defaults[name] = value
        self.generate()
        self.show_config(**self._defaults)

    def generate(self):
        assert os.path.exists(self.weights), "weights path does not exist..."
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.model = attempt_load(self.weights, map_location=self.device)  # load FP32 model
        self.model.to(self.device)
        self.model.eval()
        # Get names and colors
        self.names = self.model.module.names if hasattr(self.model, 'module') else self.model.names
        self.colors = [[random.randint(0, 255) for _ in range(3)] for _ in self.names]

    def show_config(self, **kwargs):
        print('Configurations:')
        print('-' * 70)
        print('|%25s | %40s|' % ('keys', 'values'))
        print('-' * 70)
        for key, value in kwargs.items():
            print('|%25s | %40s|' % (str(key), str(value)))
        print('-' * 70)

    def detect(self, img):
        im0 = img
        stride = int(self.model.stride.max())  # model stride
        imgsz = check_img_size(self.imgsz, s=stride)  # check img_size
        img = letterbox(im0, imgsz, stride=stride)[0]
        # Convert
        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB,to 3x416x416
        img = np.ascontiguousarray(img)

        img = torch.from_numpy(img).to(self.device)
        img = img.half() if self.half else img.float()  # uint8 to fp16/32
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        # Inference
        t1 = time_synchronized()
        pred = self.model(img, augment=self.augment)[0]

        # Apply NMS
        pred = non_max_suppression(pred, self.conf_thres, self.iou_thres, classes=self.classes,
                                   agnostic=self.agnostic_nms)
        t2 = time_synchronized()
        for i, det in enumerate(pred):  # detections per image
            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
            if len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
                # for c in det[:, -1].unique():
                #     n = (det[:, -1] == c).sum()  # detections per class
                # Write results
                for *xyxy, conf, cls in reversed(det):
                    conf = conf.item()
                    cls = int(cls.item())
                    # x0 = int(xyxy[0].item())
                    # y0 = int(xyxy[1].item())
                    # x1 = int(xyxy[2].item())
                    # y1 = int(xyxy[3].item())
                    if self.save_txt:  # Write to file
                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                        line = (cls, *xywh, conf) if self.save_conf else (cls, *xywh)  # label format
                        with open(self.txt_path + '.txt', 'a') as f:
                            f.write(('%g ' * len(line)).rstrip() % line + '\n')
                    if self.view_img:  # Add bbox to image
                        label = f'{self.names[int(cls)]} {conf:.2f}'
                        plot_one_box(xyxy, im0, label=label, color=self.colors[int(cls)], line_thickness=3)
        return im0

if __name__ == '__main__':
    yolo = YOLO()

2、调用

    if mode == "predict":
        while True:
            img = input('Input image filename:')
            try:
                image = cv2.imread(img)
            except:
                print('Open Error! Try again!')
                continue
            else:
                with torch.no_grad():
                    r_image = yolo.detect(image)
                    cv2.imshow('0', r_image)
                    cv2.waitKey(0)

    elif mode == "video":
        capture = cv2.VideoCapture(video_path)
        if video_save_path!="":
            fourcc  = cv2.VideoWriter_fourcc(*'XVID')
            size    = (int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)), int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)))
            out     = cv2.VideoWriter(video_save_path, fourcc, video_fps, size)

        ref, frame = capture.read()
        if not ref:
            raise ValueError("未能正确读取摄像头（视频），请注意是否正确安装摄像头（是否正确填写视频路径）。")

        fps = 0.0
        while(True):
            t1 = time.time()
            # 读取某一帧
            ref, frame = capture.read()
            if not ref:
                break
            # 格式转变，BGRtoRGB
            # frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)
            # 转变成Image
            # frame = Image.fromarray(np.uint8(frame))
            # 进行检测
            frame = yolo.detect(frame)
            # RGBtoBGR满足opencv显示格式
            # frame = cv2.cvtColor(frame,cv2.COLOR_RGB2BGR)
            
            fps  = ( fps + (1./(time.time()-t1))) / 2.0
            print("fps= %.2f"%(fps))
            frame = cv2.putText(frame, "fps= %.2f"%(fps), (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            
            cv2.imshow("video",frame)
            c= cv2.waitKey(1) & 0xff 
            if video_save_path!="":
                out.write(frame)

            if c==27:
                capture.release()
                break

        print("Video Detection Done!")
        capture.release()
        if video_save_path!="":
            print("Save processed video to the path :" + video_save_path)
            out.release()
        cv2.destroyAllWindows()