yolov5 文件函数detect分析,有人看就分享出来


def detect(save_img=False):
    #save_img=False
    # 获取输出文件夹,输入源,权重,参数等参数
    out, source, weights, view_img, save_txt, imgsz = \
        opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
    #webcam获取source的信息返回true表示是?
    webcam = source.isnumeric() or source.startswith('rtsp') or source.startswith('http') or source.endswith('.txt')
    #print(webcam) webcam web或者摄像头


    # Initialize
    set_logging()
    device = select_device(opt.device)

    #out默认输出目录为inference / output,存在就删除输出目录,再建目录
    if os.path.exists(out):
        shutil.rmtree(out)  # delete output folder
    os.makedirs(out)  # make new output folder
    half = device.type != 'cpu'  # half precision only supported on CUDA

    # Load model
    print("-----weights:",weights)

    model = attempt_load(weights, map_location=device)  # load FP32 model,此处加载模型该模型在程序同级目录下
#---装载------------------------------------------------------------

    imgsz = check_img_size(imgsz, s=model.stride.max())  # check img_size
    if half:
        model.half()  # to FP16

    # Second-stage classifier# 设置第二次分类,默认不使用
    classify = False
    if classify:
        modelc = load_classifier(name='resnet101', n=2)  # initialize
        modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model'])  # load weights
        modelc.to(device).eval()

    # Set Dataloader# 通过不同的输入源来设置不同的数据加载方式
    vid_path, vid_writer = None, None
    if webcam:
        view_img = True
        #如果网络的输入数据维度或类型上变化不大,也就是每次训练的图像尺寸都是一样的时候,设置 torch.backends.cudnn.benchmark = true 可以增加运行效率;
        cudnn.benchmark = True  # set True to speed up constant image size inference可加速图像预测
        dataset = LoadStreams(source, img_size=imgsz)
        #获取视频信息,线程抓取图片dataset类中imgs[0]是0个摄像头的图片,LoadStreams是迭代类---》dataset是一个迭代器
        #dataset获取的数据是(sources, img, img0, None)
        # img0是原始数据源获取的图片列表,img是处理后的图片列表,-----图像可能是是与设备同宽, 在上下添加黑边的显示模式;

    else:
        save_img = True
        dataset = LoadImages(source, img_size=imgsz) #?????????????????????????

    #print("dataset=",dataset)

    # Get names and colors
    names = model.module.names if hasattr(model, 'module') else model.names
    colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]

    # Run inference# 进行一次前向推理,测试程序是否正常
    t0 = time.time()
    img = torch.zeros((1, 3, imgsz, imgsz), device=device)  # init img
    _ = model(img.half() if half else img) if device.type != 'cpu' else None  # run once 测试程序是否正常(非cpu)


#-----------------摄像头从此处开始反复循环-dataset为迭代器类--------------------------------
    for path, img, im0s, vid_cap in dataset: #见上面
        #path 如果是摄像头就是摄像头编号,为保存识别准备文件名,例如:在inference\output文件夹中0.txt
        #把图片转为tensor,放到device上去计算
        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的shape是(1, num_boxes, 5+num_class)
                h,w为传入网络图片的长和宽,注意dataset在检测时使用了矩形推理,所以这里h不一定等于w
                num_boxes = h/32 * w/32 + h/16 * w/16 + h/8 * w/8
                pred[..., 0:4]为预测框坐标
                预测框坐标为xywh(中心点+宽长)格式
                pred[..., 4]为objectness置信度
                pred[..., 5:-1]为分类结果,5:-1  5到倒数第一个
                """
        pred_all = model(img, augment=opt.augment)

        #print("pred_all[1].len=",len(pred_all[1]))
        #pred_all结构为长度为2的tuple
        #pred_all[0]=(tensor[1, 20160, 7])  20160=3*(3个网格点数),3个网格点数=64*80+32*40+16*4)
        # pred_all[0]为所有网格点预测结果汇总,pred_all[1]为原始预测结果

        #pred_all[1]=(tensor1[shape为1, 3, 64, 80, 7]),tensor2([1, 3, 64/2, 80/2, 7]),tensor3([1, 3, 64/4, 80/4, 7])


        pred = pred_all[0]  #取出3个网格预测结果,torch.Size([1, 20160, 7])
        # print("-------------",pred.shape)
        # input("pred")
        # Apply NMS

        """
                pred:前向传播的输出
                conf_thres:置信度阈值
                iou_thres:iou阈值
                classes:是否只保留特定的类别
                agnostic:进行nms是否也去除不同类别之间的框
                经过nms之后,预测框格式:xywh-->xyxy(左上角右下角)
                pred是一个列表list[torch.tensor],长度为batch_size
                每一个torch.tensor的shape为(num_boxes, 6),6的内容为box+置信度+类别
                
                
                """
        # 下面的pred.shape=[1,tensor(x,6),...]  x是预测框个数,6的内容为box位置(4个)+置信度+类别(该函数取出预测结果)
        pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
        # print("------------pred\n",pred,"\n")
        # input("----------------")
        #pred经过非极大抑制后去掉了多余预测结果结构变为 [tensor[[x1,y1,x2,y2,概率,类别],[...],...]]

        t2 = time_synchronized()

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

        #mm=list(enumerate(pred))
        #print("-----------mm\n",mm)
        # Process detections
        for i, det in enumerate(pred):  # detections per image, pred  是所有预测结果,[[x0,y0,x1,y1,置信度,类别],....]
        #det是每张图的多个预测结果,pred是多张图的预测结果
        #此循环,如果是视频,只循环一次,pred就是一张图预测结果

            # print(len(pred))
            # input("pred len---------------")


            if webcam:  # batch_size >= 1#是摄像头或者视频
                #如果是视频 path=["0"]或者["1"],im0s=[一张原始图片]
                # print("---path[i]:",i,path[i])
                p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
            else:#如果是图片path=[多张图片路径],im0s多张图片数据
                p, s, im0 = path, '', im0s

            save_path = str(Path(out) / Path(p).name)#保存输出结果的路径
            #print("save:=",save_path)

            # print("dataset.mode=____",dataset.mode)

            txt_path = str(Path(out) / Path(p).stem) + ('_%g' % dataset.frame if dataset.mode == 'video' else '')
            s += '%gx%g ' % img.shape[2:]  # print string
            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh ,gn=[w,h,w,h]图片的宽高
            if det is not None and len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                # print(det)
                # input("det")


                # Print results,,, c取出来的就是类别,而且唯一,下面输出各类别的预测数
                for c in det[:, -1].unique():  #det.shape=(n,6),第二维度最后一个数(是类别),unique()去除det倒数第一行的重复元素,且排序
                    n = (det[:, -1] == c).sum()  # detections per class每类的检测数量
                    s += '%g %ss, ' % (n, names[int(c)])  # add to string

                # Write results#此处循环画框--------------------------------------
                for *xyxy, conf, cls in reversed(det):  #?reversed(det)啥意思
                    # print(reversed(det))
                    # input("---------reversed(det-------")
                    if save_txt:  # Write to file
                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                        with open(txt_path + '.txt', 'a') as f:
                            f.write(('%g ' * 5 + '\n') % (cls, *xywh))  # label format

                    if save_img or view_img:  # Add bbox to image
                        label = '%s %.2f' % (names[int(cls)], conf)
                        #图片上标记框框

                        #img0是视频获取的图片

                        plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)

                    #---------------------------------------------
                    myxyxy=(xyxy[0].item(), xyxy[1].item(), xyxy[2].item(), xyxy[3].item())
                    myp=conf.item()
                    myclass=cls.item()
                    print("类别、置信度、坐标::",myclass,myp,myxyxy)
                    #-------------------------------输出坐标------------------------------------------------------------

            # Print time (inference + NMS)
            print('%sDone. (%.3fs)' % (s, t2 - t1))



            # Stream results
            if view_img:
                #此处显示视频流和框的结果
                cv2.imshow(p, im0)
                if cv2.waitKey(1) == ord('q'):  # q to quit
                    raise StopIteration

            # Save results (image with detections)
            if save_img:#如果保存结果
                if dataset.mode == 'images':  #如果是图
                    cv2.imwrite(save_path, im0)
                else:#如果不是图,是视频
                    if vid_path != save_path:  # new video
                        vid_path = save_path
                        if isinstance(vid_writer, cv2.VideoWriter):
                            vid_writer.release()  # release previous video writer

                        fourcc = 'mp4v'  # output video codec
                        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))
                        vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h))
                    vid_writer.write(im0)

    if save_txt or save_img:
        print('Results saved to %s' % Path(out))
        if platform.system() == 'Darwin' and not opt.update:  # MacOS
            os.system('open ' + save_path)

    print('Done. (%.3fs)' % (time.time() - t0))

你可能感兴趣的:(python,深度学习,人工智能,计算机视觉,网络)