YOLOv6 是美团视觉智能部研发的一款目标检测框架,致力于工业应用。本框架同时专注于检测的精度和推理效率,在工业界常用的尺寸模型中:YOLOv6-nano 在 COCO 上精度可达 35.0% AP,在 T4 上推理速度可达 1242 FPS;YOLOv6-s 在 COCO 上精度可达 43.1% AP,在 T4 上推理速度可达 520 FPS。在部署方面,YOLOv6 支持 GPU(TensorRT)、CPU(OPENVINO)、ARM(MNN、TNN、NCNN)等不同平台的部署,极大地简化工程部署时的适配工作。
YOLOv6 GitHub网址:美团/YOLOv6:YOLOv6:专用于工业应用的单级物体检测框架。 (github.com)
记得把对应版本的模型也下载,我下的是YOLOv6-s
终端输入 pip install requirements.txt,YOLOv6比V5多了一个addict库,也可以只下载一个addict
打开程序,找到文件夹tools->infer.py
默认的路径需要改一下,否则会显示找不到文件,不想动手的就直接复制下面的吧
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
import argparse
import os
import sys
import os.path as osp
import torch
ROOT = os.getcwd()
if str(ROOT) not in sys.path:
sys.path.append(str(ROOT))
from yolov6.utils.events import LOGGER
from yolov6.core.inferer import Inferer
def get_args_parser(add_help=True):
parser = argparse.ArgumentParser(description='YOLOv6 PyTorch Inference.', add_help=add_help)
parser.add_argument('--weights', type=str, default='../weights/yolov6s.pt', help='model path(s) for inference.')
parser.add_argument('--source', type=str, default='../data/images', help='the source path, e.g. image-file/dir.')
parser.add_argument('--yaml', type=str, default='../data/coco.yaml', help='data yaml file.')
parser.add_argument('--img-size', type=int, default=640, help='the image-size(h,w) in inference size.')
parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold for inference.')
parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold for inference.')
parser.add_argument('--max-det', type=int, default=1000, help='maximal inferences per image.')
parser.add_argument('--device', default='0', help='device to run our model i.e. 0 or 0,1,2,3 or cpu.')
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt.')
parser.add_argument('--save-img', action='store_false', help='save visuallized inference results.')
parser.add_argument('--classes', nargs='+', type=int, help='filter by classes, e.g. --classes 0, or --classes 0 2 3.')
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS.')
parser.add_argument('--project', default='runs/inference', help='save inference results to project/name.')
parser.add_argument('--name', default='exp', help='save inference results to project/name.')
parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels.')
parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences.')
parser.add_argument('--half', action='store_true', help='whether to use FP16 half-precision inference.')
args = parser.parse_args()
LOGGER.info(args)
return args
@torch.no_grad()
def run(weights=osp.join(ROOT, 'yolov6s.pt'),
source=osp.join(ROOT, 'data/images'),
yaml=None,
img_size=640,
conf_thres=0.25,
iou_thres=0.45,
max_det=1000,
device='',
save_txt=False,
save_img=True,
classes=None,
agnostic_nms=False,
project=osp.join(ROOT, 'runs/inference'),
name='exp',
hide_labels=False,
hide_conf=False,
half=False,
):
""" Inference process
This function is the main process of inference, supporting image files or dirs containing images.
Args:
weights: The path of model.pt, e.g. yolov6s.pt
source: Source path, supporting image files or dirs containing images.
yaml: Data yaml file, .
img_size: Inference image-size, e.g. 640
conf_thres: Confidence threshold in inference, e.g. 0.25
iou_thres: NMS IOU threshold in inference, e.g. 0.45
max_det: Maximal detections per image, e.g. 1000
device: Cuda device, e.e. 0, or 0,1,2,3 or cpu
save_txt: Save results to *.txt
save_img: Save visualized inference results
classes: Filter by class: --class 0, or --class 0 2 3
agnostic_nms: Class-agnostic NMS
project: Save results to project/name
name: Save results to project/name, e.g. 'exp'
line_thickness: Bounding box thickness (pixels), e.g. 3
hide_labels: Hide labels, e.g. False
hide_conf: Hide confidences
half: Use FP16 half-precision inference, e.g. False
"""
# create save dir
save_dir = osp.join(project, name)
if (save_img or save_txt) and not osp.exists(save_dir):
os.makedirs(save_dir)
else:
LOGGER.warning('Save directory already existed')
if save_txt:
os.mkdir(osp.join(save_dir, 'labels'))
# Inference
inferer = Inferer(source, weights, device, yaml, img_size, half)
inferer.infer(conf_thres, iou_thres, classes, agnostic_nms, max_det, save_dir, save_txt, save_img, hide_labels, hide_conf)
if save_txt or save_img:
LOGGER.info(f"Results saved to {save_dir}")
def main(args):
run(**vars(args))
if __name__ == "__main__":
args = get_args_parser()
main(args)
创建一个weights文件夹,将在官网下载好的yolov6s.pt模型放进去
再找到文件夹yolov6->core->inferer.py文件中168行,在路径加个点,果然初版还是不够完善啊!
这里就配置完了,运行infer.py文件
测试效果路径在tools文件夹里
我用YOLOv5测试对比一下:
YOLOv6的精度和置信度确实比YOLOv5要好一些,但是误检率太高,并且版本维护更新速度太慢,不适合用于工业领域,自己测着玩还行,下次再更新训练自己的数据集。