import os           # 导入系统环境变量
import sys          # 导入 python环境变量相关的函数
import cv2
import torch        # 导入机器学习库 pytorch
import torch.backends.cudnn as cudnn    # 导入nvdia cudnn 相关库
import warnings
warnings.filterwarnings("ignore")

from pathlib import Path
from models.experimental import attempt_load    # attempt_load用于加载模型权重文件并构建模型(可以构造普通模型或者集成模型)
from utils.datasets import LoadImages        # cv2图片、视频读取函数
from utils.general import non_max_suppression, scale_coords, set_logging
from utils.plots import Annotator, colors
from utils.torch_utils import select_device


# 1. 路径处理
FILE = Path(__file__).resolve()     # 解析当前文件的绝对路径
ROOT = FILE.parents[0]              # 解析文件root 目录
if str(ROOT) not in  sys.path:      # 如果当前目前不在python 环境变量中,则添加环境变量
    sys.path.append(str(ROOT))      
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # 计算当前路径距根目录的相对路径

source = str(ROOT / "data/images/4.jpeg")          # 要检测的文件路径 

set_logging()   # 日志

# 2. 配置显卡
device = select_device(device='')           # 使用默认的CUDA 显卡
half = device.type != 'cpu'                 # 半精度浮点数仅GPU支持

# 3. 加载 yolov5s.pt 模型,开始建模
weights = str(ROOT / 'yolov5s.pt')               # 模理权重文件路径
model = attempt_load(weights, map_location=device)
stride = int(model.stride.max())                # 获取当前模型中支持的最大模型数量
names = model.modules.names if hasattr(model, 'module') else model.names    # 获取当前模型中各模型的name

if half:                # 如果支持半精度浮点
    model.half()        # 修改所有参数和buffer 为half半精度浮点类型

imgsz = [640,640]

# 4. 加载图片、视频资源
dataset = LoadImages(source, img_size=640,stride=stride, auto=True)

# 5. 开始检测
if device.type != 'cpu':
    model(torch.zeros(1,3,*imgsz).to(device).type_as(next(model.parameters())))

for path, img, im0s, vid_cap in dataset:
    img = torch.from_numpy(img).to(device)
    img = img.half() if half else img.float()  # uint8 to fp16/32
    img = img / 255.0
    if len(img.shape) == 3:
            img = img[None]

    # 6. 配置特征图可视化
    pred = model(img, augment=False, visualize=False)[0]

    # NMS
    pred = non_max_suppression(pred, 0.5, 0.45, None, False, multi_label=False, labels=(), max_det=300)

    # 7.开始预测    
    for i, det in enumerate(pred):
        p, s, im0, frame = Path(path), '', im0s.copy(), getattr(dataset, 'frame', 0)
        s += '%gx%g, ' % img.shape[2:]                # print string

        annotator = Annotator(im0, line_width=3, example=str(names))

        if len(det):
            # Rescale boxes from img_size to im0 size
            det[:, :4]  = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

            # Print results
            for c in det[:, -1].unique():
                n = (det[:, -1] == c).sum()  # detections per class
                s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
            print(s)

            # Write results
            for *xyxy, conf, cls in reversed(det):
                c = int(cls)  # integer class
                label = (f'{names[c]} {conf:.2f}')
                annotator.box_label(xyxy, label, color=colors(c, True))
                print(label,xyxy)

        # Stream results
        im0 = annotator.result()
        cv2.imshow(str(p), im0)
        cv2.waitKey(5000)  # 1 millisecond

运行结果:

PS C:\Users\ciellee\Desktop\Yolo\yolov5-6.0> & C:/ProgramData/Miniconda3/envs/pytorch/python.exe c:/Users/ciellee/Desktop/Yolo/yolov5-6.0/detect_test.py
YOLOv5  2021-10-12 torch 1.10.2 CPU

Fusing layers... 
Model Summary: 224 layers, 7266973 parameters, 0 gradients
image 1/1 C:\Users\ciellee\Desktop\Yolo\yolov5-6.0\data\images\4.jpeg: 384x640, 5 persons, 4 cars, 1 traffic light, 
person 0.54 [tensor(61.), tensor(158.), tensor(145.), tensor(382.)]
traffic light 0.59 [tensor(1225.), tensor(59.), tensor(1242.), tensor(88.)]
person 0.65 [tensor(728.), tensor(173.), tensor(775.), tensor(371.)]
person 0.70 [tensor(673.), tensor(166.), tensor(750.), tensor(366.)]
car 0.75 [tensor(1133.), tensor(198.), tensor(1218.), tensor(268.)]
car 0.75 [tensor(1027.), tensor(195.), tensor(1102.), tensor(273.)]
car 0.79 [tensor(367.), tensor(206.), tensor(483.), tensor(328.)]
car 0.86 [tensor(128.), tensor(172.), tensor(385.), tensor(357.)]
person 0.88 [tensor(436.), tensor(0.), tensor(744.), tensor(710.)]
person 0.91 [tensor(811.), tensor(26.), tensor(1054.), tensor(613.)]
PS C:\Users\ciellee\Desktop\Yolo\yolov5-6.0> 

在这里插入图片描述

Logo

为开发者提供学习成长、分享交流、生态实践、资源工具等服务,帮助开发者快速成长。

更多推荐