ultralytics / yolov5

YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite
https://docs.ultralytics.com
GNU Affero General Public License v3.0
50.44k stars 16.27k forks source link

How to write number of objects in image/video during inference #3959

Closed ghost closed 3 years ago

ghost commented 3 years ago

❔Question

Hi, that's great work. I want to know how the total number of objects are written during inference the image is attached below how can I count the total number of objects in the image in detect.py file for both image and video.

IMG_20200111_113830

Additional context

shershunov commented 3 years ago

Hi! The number is written on the command line.

Untitled-1

Or you can use PytorchHub.

from cv2 import waitKey, destroyAllWindows, imshow, VideoCapture, rectangle
from torch import hub

model = hub.load('ultralytics/yolov5', 'yolov5s')
cap = VideoCapture(0)
while True:
    a = 0
    ret, img = cap.read()
    result = model(img, size=640)
    repos = result.xyxy[0]
    for i in range(len(repos)):
        pos = repos[i].detach()
        rectangle(img, (int(pos[0]), int(pos[1])), (int(pos[2]), int(pos[3])), (255, 255, 255), 2)
        a += 1
    print('Detected objects:', a)
    imshow('Detect', img)
    if waitKey(1) & 0xFF == ord("q"):
        destroyAllWindows()
        break
ghost commented 3 years ago

Yes, They are written when running the inference but I want to write on the image/video, How can i do that?

shershunov commented 3 years ago

In PytorchHub, it will look like this:

from cv2 import rectangle, putText, LINE_AA, imread, imwrite
from torch import hub

model = hub.load('ultralytics/yolov5', 'yolov5s')

a = 0
img = imread('img.jpg')
result = model(img, size=640)
repos = result.xyxy[0]
for i in range(len(repos)):
    pos = repos[i].detach()
    rectangle(img, (int(pos[0]), int(pos[1])), (int(pos[2]), int(pos[3])), (255, 255, 255), 2)
    a += 1
putText(img, f'Detected objects: {a}', (10, 40), LINE_AA, 1, (255, 255, 255), 2)
imwrite('res.jpg', img)

res

ghost commented 3 years ago

Thanks for sharing, but this is for single I want it in detect.py file how can i modified that, i have only single class

ghost commented 3 years ago

while loading my model the output doesn't detect any thing

from cv2 import rectangle, putText, LINE_AA, imread, imwrite
from torch import hub

model = torch.hub.load('/content/yolov5', 'custom', path='/content/yolov5/runs/train/exp/weights/best.pt', source='local')

a = 0
img = imread('/content/00001IMG_00001_BURST20200111115109.jpg')
result = model(img, size=640)
repos = result.xyxy[0]
for i in range(len(repos)):
    pos = repos[i].detach()
    rectangle(img, (int(pos[0]), int(pos[1])), (int(pos[2]), int(pos[3])), (0, 0, 255), 2)
    a += 1
putText(img, f'Detected objects: {a}', (10, 40), LINE_AA, 1, (0, 0, 255), 2)
imwrite('res.jpg', img)
shershunov commented 3 years ago

The photo is saved to the folder where the Python file is located. Try it like this model = hub.load('ultralytics/yolov5', 'custom', path='/content/yolov5/runs/train/exp/weights/best.pt',)

shershunov commented 3 years ago

Describe in more detail, do you not have a photo saved or are the objects not marked in the photo?

ghost commented 3 years ago

Objects are not marked.

shershunov commented 3 years ago

Check your weights. Everything works fine for me)

ghost commented 3 years ago

Weights are working fine without this.

shershunov commented 3 years ago

You are using torch.hub even though hub has already been called

shershunov commented 3 years ago

Put the weights in the root folder of your Python file, and use this script. 👍

from cv2 import rectangle, putText, LINE_AA, imread, imwrite
from torch import hub

model = hub.load('ultralytics/yolov5', 'custom', path='best.pt',)
a = 0
img = imread('/content/00001IMG_00001_BURST20200111115109.jpg')
result = model(img, size=640)
repos = result.xyxy[0]
for i in range(len(repos)):
    pos = repos[i].detach()
    rectangle(img, (int(pos[0]), int(pos[1])), (int(pos[2]), int(pos[3])), (0, 0, 255), 2)
    a += 1
putText(img, f'Detected objects: {a}', (10, 40), LINE_AA, 1, (0, 0, 255), 2)
imwrite('res.jpg', img)
shershunov commented 3 years ago

If you still need detect.py then here) 😊 car

"""Run inference with a YOLOv5 model on images, videos, directories, streams

Usage:
    $ python path/to/detect.py --source path/to/img.jpg --weights yolov5s.pt --img 640
"""

import argparse
import sys
import time
from pathlib import Path

import cv2
import torch
import torch.backends.cudnn as cudnn

FILE = Path(__file__).absolute()
sys.path.append(FILE.parents[0].as_posix())  # add yolov5/ to path

from models.experimental import attempt_load
from utils.datasets import LoadStreams, LoadImages
from utils.general import check_img_size, check_requirements, check_imshow, colorstr, non_max_suppression, \
    apply_classifier, scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path, save_one_box
from utils.plots import colors, plot_one_box
from utils.torch_utils import select_device, load_classifier, time_synchronized

@torch.no_grad()
def run(weights='yolov5s.pt',  # model.pt path(s)
        source='data/images',  # file/dir/URL/glob, 0 for webcam
        imgsz=640,  # inference size (pixels)
        conf_thres=0.25,  # confidence threshold
        iou_thres=0.45,  # NMS IOU threshold
        max_det=1000,  # maximum detections per image
        device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
        view_img=False,  # show results
        save_txt=False,  # save results to *.txt
        save_conf=False,  # save confidences in --save-txt labels
        save_crop=False,  # save cropped prediction boxes
        nosave=False,  # do not save images/videos
        classes=None,  # filter by class: --class 0, or --class 0 2 3
        agnostic_nms=False,  # class-agnostic NMS
        augment=False,  # augmented inference
        visualize=False,  # visualize features
        update=False,  # update all models
        project='runs/detect',  # save results to project/name
        name='exp',  # save results to project/name
        exist_ok=False,  # existing project/name ok, do not increment
        line_thickness=3,  # bounding box thickness (pixels)
        hide_labels=False,  # hide labels
        hide_conf=False,  # hide confidences
        half=False,  # use FP16 half-precision inference
        ):
    save_img = not nosave and not source.endswith('.txt')  # save inference images
    webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
        ('rtsp://', 'rtmp://', 'http://', 'https://'))

    # Directories
    save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
    (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir

    # Initialize
    set_logging()
    device = select_device(device)
    half &= device.type != 'cpu'  # half precision only supported on CUDA

    # Load model
    model = attempt_load(weights, map_location=device)  # load FP32 model
    stride = int(model.stride.max())  # model stride
    imgsz = check_img_size(imgsz, s=stride)  # check image size
    names = model.module.names if hasattr(model, 'module') else model.names  # get class names
    if half:
        model.half()  # to FP16

    # Second-stage classifier
    classify = False
    if classify:
        modelc = load_classifier(name='resnet50', n=2)  # initialize
        modelc.load_state_dict(torch.load('resnet50.pt', map_location=device)['model']).to(device).eval()

    # Dataloader
    if webcam:
        view_img = check_imshow()
        cudnn.benchmark = True  # set True to speed up constant image size inference
        dataset = LoadStreams(source, img_size=imgsz, stride=stride)
        bs = len(dataset)  # batch_size
    else:
        dataset = LoadImages(source, img_size=imgsz, stride=stride)
        bs = 1  # batch_size
    vid_path, vid_writer = [None] * bs, [None] * bs

    # Run inference
    if device.type != 'cpu':
        model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))  # run once
    t0 = time.time()
    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 /= 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        # Inference
        t1 = time_synchronized()
        pred = model(img,
                     augment=augment,
                     visualize=increment_path(save_dir / 'features', mkdir=True) if visualize else False)[0]

        # Apply NMS
        pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
        t2 = time_synchronized()

        # Apply Classifier
        if classify:
            pred = apply_classifier(pred, modelc, img, im0s)

        # Process detections
        for i, det in enumerate(pred):  # detections per image
            if webcam:  # batch_size >= 1
                p, s, im0, frame = path[i], f'{i}: ', im0s[i].copy(), dataset.count
            else:
                p, s, im0, frame = path, '', im0s.copy(), getattr(dataset, 'frame', 0)

            p = Path(p)  # to Path
            save_path = str(save_dir / p.name)  # img.jpg
            txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # img.txt
            s += '%gx%g ' % img.shape[2:]  # print string
            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
            imc = im0.copy() if save_crop else im0  # for save_crop
            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

                # Write results
                a = 0
                for *xyxy, conf, cls in reversed(det):
                    a += 1
                    if save_txt:  # Write to file
                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                        line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
                        with open(txt_path + '.txt', 'a') as f:
                            f.write(('%g ' * len(line)).rstrip() % line + '\n')

                    if save_img or save_crop or view_img:  # Add bbox to image
                        c = int(cls)  # integer class
                        label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
                        plot_one_box(xyxy, im0, label=label, color=colors(c, True), line_thickness=line_thickness)
                        if save_crop:
                            save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
                cv2.putText(im0, f'Detected objects: {a}', (10, 40), cv2.LINE_AA, 1, colors(c, True), 2)

            # Print time (inference + NMS)
            print(f'{s}Done. ({t2 - t1:.3f}s)')

            # Stream results
            if view_img:
                cv2.imshow(str(p), im0)
                cv2.waitKey(1)  # 1 millisecond

            # Save results (image with detections)
            if save_img:
                if dataset.mode == 'image':
                    cv2.imwrite(save_path, im0)
                else:  # 'video' or 'stream'
                    if vid_path[i] != save_path:  # new video
                        vid_path[i] = save_path
                        if isinstance(vid_writer[i], cv2.VideoWriter):
                            vid_writer[i].release()  # release previous video writer
                        if vid_cap:  # video
                            fps = vid_cap.get(cv2.CAP_PROP_FPS)
                            w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                            h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                        else:  # stream
                            fps, w, h = 30, im0.shape[1], im0.shape[0]
                            save_path += '.mp4'
                        vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
                    vid_writer[i].write(im0)

    if save_txt or save_img:
        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
        print(f"Results saved to {save_dir}{s}")

    if update:
        strip_optimizer(weights)  # update model (to fix SourceChangeWarning)

    print(f'Done. ({time.time() - t0:.3f}s)')

def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
    parser.add_argument('--source', type=str, default='data/images', help='file/dir/URL/glob, 0 for webcam')
    parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold')
    parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--view-img', action='store_true', help='show results')
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
    parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
    parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--visualize', action='store_true', help='visualize features')
    parser.add_argument('--update', action='store_true', help='update all models')
    parser.add_argument('--project', default='runs/detect', help='save results to project/name')
    parser.add_argument('--name', default='exp', help='save results to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
    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='use FP16 half-precision inference')
    opt = parser.parse_args()
    return opt

def main(opt):
    print(colorstr('detect: ') + ', '.join(f'{k}={v}' for k, v in vars(opt).items()))
    check_requirements(exclude=('tensorboard', 'thop'))
    run(**vars(opt))

if __name__ == "__main__":
    opt = parse_opt()
    main(opt)
ghost commented 3 years ago

Thanks for this, already done the same thing :), How can I apply counting on video do you have any idea about it?

shershunov commented 3 years ago

The video display will also work.

github-actions[bot] commented 3 years ago

👋 Hello, this issue has been automatically marked as stale because it has not had recent activity. Please note it will be closed if no further activity occurs.

Access additional YOLOv5 🚀 resources:

Access additional Ultralytics ⚡ resources:

Feel free to inform us of any other issues you discover or feature requests that come to mind in the future. Pull Requests (PRs) are also always welcomed!

Thank you for your contributions to YOLOv5 🚀 and Vision AI ⭐!

shiceQ commented 2 years ago

如果您仍然需要detect.py,那么在这里)😊 车

"""使用 YOLOv5 模型对图像、视频、目录、流进行推理

用法:
    $ python path/to/detect.py --source path/to/img.jpg --weights yolov5s.pt --img 640 
"""

import  argparse 
import  sys 
import  time 
from  pathlib  import  Path

进口 cv2
进口 火炬
进口 火炬。后端。cudnn 作为 cudnn

文件 = 路径(__file__)。绝对()
系统。路径。append ( FILE . parents [ 0 ]. as_posix ())   # 将 yolov5/ 添加到路径

从 模型。来自utils的实验性 导入 尝试加载
。数据集从utils导入LoadStreams、LoadImages。一般导入check_img_size,check_requirements,check_imshow,colorstr,non_max_suppression, \
     apply_classifier,scale_coords,xyxy2xywh,strip_optimizer,set_logging,increment_path,save_one_box   

从 utils。从utils绘制 导入 颜色和plot_one_box
。torch_utils导入select_device,load_classifier,time_synchronized   

@火炬。no_grad () 
def  run ( weights = 'yolov5s.pt' ,   #model.pt path(s) 
        source = 'data/images' ,   # file/dir/URL/glob, 0 for webcam 
        imgsz = 640 ,   # inference size (像素) 
        conf_thres = 0.25 ,   # 置信阈值
        iou_thres = 0.45 ,   # NMS IOU 阈值
        max_det = 1000 ,   # 每个图像
        设备的最大检测数= '' ,  # cuda 设备,即 0 或 0,1,2,3 或 cpu 
        view_img = False ,   # 显示结果
        save_txt = False ,   # 将结果保存到 *.txt 
        save_conf = False ,   # 在 --save-txt 标签中保存置信度
        save_crop = False ,   # 保存裁剪的预测框
        nosave = False ,   # 不保存图像/视频
        classes = None ,   # 按类别过滤:--class 0 或 --class 0 2 3 
        agnostic_nms = False ,   # 与类别无关的 NMS
        增强= False ,   # 增强推理
        Visualize = False ,   # 可视化特征
        update = False ,   # 更新所有模型
        project = 'runs/detect' ,   # 将结果保存到项目/名称
        name = 'exp' ,   # 将结果保存到项目/名称
        exists_ok = False ,   # 现有项目/名称 ok, 不要增加
        line_thickness = 3 ,   # 边界框厚度(像素)
        hide_labels = False ,   # 隐藏标签
        hide_conf = False ,   # 隐藏置信度
        half = False ,   # 使用 FP16 半精度推理
        ):
    save_img  = 不是 nosave 也不 是source 。endswith ( '.txt' )   # 保存推理图像webcam = source . isnumeric ()或source。以( '.txt' )或source结尾。较低()。以(

        ( 'rtsp://' , 'rtmp://' , 'http://' , 'https://' ))

    # 目录
    save_dir  =  increment_path ( Path ( project ) /  name , exists_ok = exists_ok )   # increment run 
    ( save_dir  /  'labels'  if  save_txt  else  save_dir )。mkdir ( parents = True , exist_ok = True )   # 制作目录

    # 初始化
    set_logging ()
     device  =  select_device ( device )
     half &= device . type  !=  'cpu'   # 半精度仅在 CUDA 上支持

    # 加载模型
    model  =  attempt_load ( weights , map_location = device )   # 加载 FP32 模型
    stride =  int  ( model.stride.max ( ) )   # model stride imgsz = check_img_size ( imgsz , s = stride ) #  检查图像尺寸名称= model . 模块。名称if hasattr ( model , 'module' )

        其他 型号。names   # 获取类名
    if  half :
         model . half ()   # 到 FP16

    #第二阶段分类器
    classify  =  False 
    if 分类:
         modelc  =  load_classifier ( name = 'resnet50' , n = 2 )   #初始化
        modelc。load_state_dict ( torch.load ( ' resnet50.pt ' , map_location = device ) [ 'model' ])。到(设备)。评估()

    # Dataloader 
    if  webcam :
         view_img  =  check_imshow ()
         cudnn . benchmark  =  True   # 设置 True 以加快恒定图像大小推理
        dataset  =  LoadStreams ( source , img_size = imgsz , stride = stride )
         bs  =  len ( dataset )   # batch_size 
    else :
         dataset  =  LoadImages ( source , img_size= imgsz , stride = stride )
         bs  =  1   # batch_size 
    vid_path , vid_writer  = [ None ] *  bs , [ None ] *  bs

    # 运行推理
    if  device . type  !=  'cpu' :
         model ( torch . zeros ( 1 , 3 , imgsz , imgsz ). to ( device ). type_as ( next ( model . parameters ())))   # 运行一次
    t0  =  time . time ()
     for  path , img , im0s , vid_cap 在 数据集中:
         img  = 火炬。from_numpy ( img )。到(设备)
         img  =  img。half () if  half  else  img . float ()   # uint8 to fp16/32 
        img  /=  255.0   # 0 - 255 to 0.0 - 1.0 
        if  img . ndimension () ==  3 :
             img  =  img。解压( 0 )

        # Inference 
        t1  =  time_synchronized ()
         pred  =  model ( img ,
                      augment = augment ,
                      visible = increment_path ( save_dir / 'features' , mkdir = True ) if visible else False )[ 0 ]     

        # 应用 NMS 
        pred  =  non_max_suppression ( pred , conf_thres , iou_thres , classes , agnostic_nms , max_det = max_det )
         t2  =  time_synchronized ()

        # 应用分类器
        if 分类:
             pred  =  apply_classifier ( pred , modelc , img , im0s )

        # 处理 i 的检测
        , det in enumerate  (  pred ) : #  每个图像的检测
            如果 网络摄像头:   # batch_size >= 1 
                p , s , im0 , frame  =  path [ i ], f' { i } : ' , im0s [ i ] . 复制(),数据集。计算
            其他:
                 p,s,im0,框架 = 路径,'',im0s。复制(), getattr (数据集, 'frame' , 0 )

            p  =  Path ( p )   # 到路径
            save_path  =  str ( save_dir  /  p .name ) # img.jpg txt_path = str ( save_dir / ' labels' / p . stem )   + ( '' if dataset .mode == ' image ' else f'_ { frame } ' )   # img.txt s += '%gx%g ' % img .

                shape [ 2 :]   # 打印字符串
            gn  =  torch . 张量( im0 . shape )[[ 1 , 0 , 1 , 0 ]]   # 归一化增益 whwh 
            imc  =  im0 . copy () if  save_crop  else  im0   # for save_crop 
            if  len ( det ):
                 # rescale box from img_size to im0 size 
                det [:, : 4 ] =  scale_coords (图像。形状[ 2 :],det [:, : 4 ],im0。形状)。圆形()

                #在det [:, - 1 ]中打印c的结果。 唯一():
                     n = ( det [:, - 1 ] == c )。sum ()   # 每个类的检测s += f" { n } { names [ int ( c )] } { 's' * ( n > 1 ) } , " # 添加到字符串    

                # 将结果
                a  =  0 
                for  * xyxy , conf , cls  in  reversed ( det ):
                     a  +=  1 
                    if  save_txt :   # 写入文件
                        xywh  = ( xyxy2xywh ( torch . tensor ( xyxy ) .view ( 1 , 4 )) /  gn )。查看( - 1 )。tolist ()   # 归一化 xywh
                        line  = ( cls , * xywh , conf ) if  save_conf  else ( cls , * xywh )   # 标签格式
                        with  open ( txt_path  +  '.txt' , 'a' ) as  f :
                             f . 写(( '%g'  *  len ( line )). rstrip () %  line  +  ' \n ' )

                    if  save_img  or  save_crop  or  view_img :   # 添加 bbox 到图像
                        c  =  int ( cls )   # integer class 
                        label  =  None  if  hide_labels  else ( names [ c ] if  hide_conf  else  f' { names [ c ] }  { conf :.2f } ' )
                         plot_one_box ( xyxy , im0 ,标签=标签,颜色=颜色(c,True),line_thickness = line_thickness)
                        如果 save_crop:
                             save_one_box(xyxy,imc,file = save_dir  /  'crops'  /  names [ c ] /  f' { p . stem } .jpg',BGR =真)
                 cv2. putText ( im0 , f'检测到的对象: { a } ' , ( 10 , 40 ), cv2 . LINE_AA , 1 , colors ( c , True ), 2 )

            # 打印时间(推理 + NMS)
            print ( f' { s } Done. ( { t2  -  t1 :.3f } s)' )

            # 流结果
            if  view_img :
                 cv2 . imshow ( str ( p ), im0 )
                 cv2。waitKey ( 1 )   # 1毫秒

            # 保存结果(带有检测的图像)
            if  save_img :
                 if  dataset . 模式 ==  '图像' :
                     cv2。imwrite ( save_path , im0 )
                 else :   # 'video' or 'stream' 
                    if  vid_path [ i ] !=  save_path :   # new video 
                        vid_path [ i ] =  save_path 
                        if  isinstance ( vid_writer [ i ], cv2. VideoWriter ):
                             vid_writer [ i ]。release ()   # 释放之前的视频作者
                        if  vid_cap :   # video 
                            fps  =  vid_cap . 获取( cv2.CAP_PROP_FPS ) w = int ( vid_cap.get ( cv2.CAP_PROP_FRAME_WIDTH ) )
                             h = int ( vid_cap.get ( cv2 _ _
                             _ _ _ _ _    . CAP_PROP_FRAME_HEIGHT ))
                         else :   # 流
                            fps , w , h  =  30 , im0 . 形状[ 1 ],im0。形状[ 0 ]
                            保存路径 +=  '.mp4' 
                        vid_writer [ i ] =  cv2。VideoWriter ( save_path , cv2 . VideoWriter_fourcc ( * 'mp4v' ), fps, ( w , h ))
                     vid_writer [ i ]。写(im0)

    if  save_txt  or  save_img :
         s  =  f" \n { len ( list ( save_dir . glob ( 'labels/*.txt' ))) }标签保存到{ save_dir  /  'labels' } "  if  save_txt  else  '' 
        print ( f "结果保存到{ save_dir } { s } " )

    if  update :
         strip_optimizer ( weights )   # 更新模型(修复 SourceChangeWarning)

    打印(f'完成。({时间.时间()-t0 :.3f } s)')

def  parse_opt ():
    解析器 =  argparse。ArgumentParser ()
    解析器。add_argument ( '--weights' , nargs = '+' , type = str , default = 'yolov5s.pt' , help = 'model.pt path(s)' )
    解析器。add_argument ( '--source' , type = str , default = 'data/images' , help ='file/dir/URL/glob, 0 for webcam' )
    解析器。add_argument ( '--imgsz' , '--img' , '--img-size' , type = int , default = 640 , help = 'inference size (pixels)' )
    解析器。add_argument ( '--conf-thres' , type = float , default = 0.25 , help = 'confidence threshold' )
    解析器。add_argument( '--iou-thres' , type = float , default = 0.45 , help = 'NMS IoU threshold' )
    解析器。add_argument ( '--max-det' , type = int , default = 1000 , help = 'maximum detections per image' )
    解析器。add_argument ( '--device' , default = '' , help = 'cuda device, ie 0 or 0,1,2,3 or cpu')
    解析器。add_argument ( '--view-img' , action = 'store_true' , help = 'show results' )
    解析器。add_argument ( '--save-txt' , action = 'store_true' , help = 'save results to *.txt' )
    解析器。add_argument ( '--save-conf' , action = 'store_true' , help = '在 --save-txt 标签中保存置信度'. add_argument ( '--save-crop' , action = 'store_true' , help = '保存裁剪的预测框' )
    解析器。add_argument ( '--nosave' , action = 'store_true' , help = '不保存图像/视频' )
    解析器。add_argument ( '--classes' , nargs = '+' , type = int , help = '按类过滤: --class 0,)
    解析器。add_argument ( '--agnostic-nms' , action = 'store_true' , help = 'class-agnostic NMS' )
    解析器。add_argument ( '--augment' , action = 'store_true' , help = 'augmented inference' )
    解析器。add_argument ( '--visualize' , action = 'store_true' , help = 'visualize features'. add_argument ( '--update' , action = 'store_true' , help = 'update all models' )
    解析器。add_argument ( '--project' , default = 'runs/detect' , help = 'save results to project/name' )
    解析器。add_argument ( '--name' , default = 'exp' , help = 'save results to project/name' )
    解析器。( '--exist-ok' , action = 'store_true' , help = 'existing project/name ok, do not increase' )
    解析器。add_argument ( '--line-thickness' , default = 3 , type = int , help = 'bounding box thickness (pixels)' )
    解析器。add_argument ( '--hide-labels' ,默认= False , action = 'store_true' , help ='隐藏标签')
    解析器。add_argument ( '--hide-conf' , default = False , action = 'store_true' , help = 'hide confidences' )
    解析器。add_argument ( '--half' , action = 'store_true' , help = '使用 FP16 半精度推理' )
     opt  =  parser。parse_args ()
    返回 选项

def  main ( opt ):
     print ( colorstr ( 'detect: ' ) +  ', ' . join ( f' { k } = { v } '  for  k , v  in  vars ( opt ). items ()))
     check_requirements ( exclude = ( 'tensorboard' , 'thop ' ))
    运行( ** vars ( opt ))

如果 __name__  ==  "__main__" :
     opt  =  parse_opt ()
     main ( opt )

视频显示也将起作用。

how to use on 6.0