Bob's Blog

Python、Web开发、测试框架、自动化平台、机器学习等

返回上页首页

用YOLO做实时目标检测



YOLO全称是you look only once,是一种对象检测算法,详情可参考2016年原作者的论文:https://arxiv.org/abs/1506.02640;主要是实时目标检测,应用到很多方面,常见的比如道路监控、安防检测、物品计数等等。可以对图片或者视频中的目标快速检测,预测类型,确定区域坐标范围。

YOLO将物体检测作为一个回归问题进行求解,输入图像经过一次推断,便能得到图像中所有物体的位置和其所属类别及相应的置信概率,用整张图作为网络的输入,直接在输出层回归bounding box的位置和bounding box所属的类别。而不是像rcnn或fast cnn之类的需要将检测结果分为两部分(分类和位置)求解。

最初实现yolo的是darknet,对应github地址是https://github.com/pjreddie/darknet,持续到yolov3。之后也出现了用tensorflow和pytorch实现的版本,比如推到yolov5的:https://github.com/ultralytics/yolov5。

对于darknet,原作者比较久没有更新过了,不过还有另一个在持续更新到yolov4,并增加了windows的支持,做了不少调优,一般提到darknet可以用这个:https://github.com/AlexeyAB/darknet。

这些大神,都提供了训练好的数据集,用于常见物品,所以最重要的的是权重文件(比如对于v3来说,需要用到3个文件:coco.names,yolov3.cfg,yolov3.weights,weights文件提供大概九千多物品的检测大概245M的样子,也提供了轻量级tiny的weights文件,包含的物品较少),只要有weights文件,其实也可以用opencv来加载darknet,不必带有全部代码;还可以预先对自定义数据做训练用作自定义的场景。

1. 直接用darknet来试试目标检测。

先下载代码:

git clone https://github.com/AlexeyAB/darknet.git

v3或v4的配置文件需要单独下载。

下载完成后,进入目录并构建:

cd darknet
make

在当前目录下会生成新文件,我们可以指定配置文件运行它来检测图片:

./darknet detector test cfg/coco.data cfg/yolov4.cfg cfg/yolov4.weights

会看到如下的信息,可以不停地输入图片路径,并返回目标类型和可能性:

也可以直接指定图像路径,并输出坐标范围:

./darknet detector test cfg/coco.data cfg/yolov3.cfg cfg/yolov3_config/cfg/yolov3.weights -ext_output ../work.jpg

也可以用来检测视频中的对象:

./darknet detector demo cfg/coco.data cfg/yolov3.cfg cfg/yolov3_config/cfg/yolov3.weights -ext_output ../sample.mp4

如果遇到"Demo needs OpenCV for webcam images"的错误提示,意味着需要启用opencv,以mac为例,先安装上opencv

brew install opencv

装成功一堆东西后,再修改darknet的makefile,将OPENCV=0改成OPENCV=1,然后重新运行make,等待完成即可重新运行命令了。

2. 用opencv来加载weights文件做目标检测。

效果是一样的,可以对图片或者对视频实时检测,只是在用opencv加载weights文件。这里就通过电脑摄像头来实时检测吧。

用后面的代码,保存并运行,即可看到如下效果,当前是用摄像头实时检测:

import time
import os
import numpy as np
import cv2

def yolo_detect(label_path='./yolov3_config/cfg/coco.names',
                config_path='./yolov3_config/cfg/yolov3.cfg',
                weights_path='./yolov3_config/cfg/yolov3.weights',
                confidence_thre=0.5,
                nms_thre=0.3,
                jpg_quality=80):

    LABELS = open(label_path).read().strip().split("\n")
    nclass = len(LABELS)

    np.random.seed(42)
    COLORS = np.random.randint(0, 255, size=(nclass, 3), dtype='uint8')

    cap = cv2.VideoCapture(0)
    cap.set(3, 640)
    cap.set(4, 480)
    print(cap.get(3), cap.get(4))
    net = cv2.dnn.readNetFromDarknet(config_path, weights_path)
    while cap.isOpened():
        ret, img = cap.read()
        user_input = cv2.waitKey(1)
        if user_input == ord('q'):
            break
        else:
            (H, W) = img.shape[:2]

            print('Load yolo weights...')
            ln = net.getLayerNames()
            ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]
            blob = cv2.dnn.blobFromImage(img, 1 / 255.0, (416, 416), swapRB=True, crop=False)
            net.setInput(blob)
            start = time.time()
            layerOutputs = net.forward(ln)
            end = time.time()

            print('Spend {:.2f} seconds to predict objects in current image...'.format(end - start))

            boxes = []
            confidences = []
            classIDs = []

            for output in layerOutputs:
                for detection in output:
                    scores = detection[5:]
                    classID = np.argmax(scores)
                    confidence = scores[classID]
                    if confidence > confidence_thre:
                        box = detection[0:4] * np.array([W, H, W, H])
                        (centerX, centerY, width, height) = box.astype("int")
                        x = int(centerX - (width / 2))
                        y = int(centerY - (height / 2))
                        boxes.append([x, y, int(width), int(height)])
                        confidences.append(float(confidence))
                        classIDs.append(classID)

            idxs = cv2.dnn.NMSBoxes(boxes, confidences, confidence_thre, nms_thre)

            print(boxes)
            print(confidences)
            print(classIDs)
            print(idxs)
            if len(idxs) > 0:
                for i in idxs.flatten():
                    (x, y) = (boxes[i][0], boxes[i][1])
                    (w, h) = (boxes[i][2], boxes[i][3])
                    color = [int(c) for c in COLORS[classIDs[i]]]
                    cv2.rectangle(img, (x, y), (x + w, y + h), color, 2)
                    text = '{}: {:.3f}'.format(LABELS[classIDs[i]], confidences[i])
                    (text_w, text_h), baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)
                    cv2.rectangle(img, (x, y - text_h - baseline), (x + text_w, y), color, -1)
                    cv2.putText(img, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 2)
                    print(text)
                    print((x, y), (x + w, y + h))

            cv2.imshow("show in real-time video capture", img)


if __name__ == '__main__':
    yolo_detect()

 

下一篇:  在终端输出中显示颜色
上一篇:  Django Restframework自定义修改局部的分页设置

共有0条评论

添加评论

暂无评论