yolo算法python代码_python通过yolo算法识别图片中的对象

1 场景

YOLO全称You Only Look Once，YOLO实现了自动驾驶汽车等前沿技术中使用的实时对象检测。

可以实现对图像中的对象进行识别和定位，如下效果：

1.png

这里不研究算法，只研究，在python环境下，如何使用yolo进行图像中物体的识别和定位。

2 官网

2.1 官方地址

YOLO使用相关网站如下：

github地址：

2.2 百度网盘

如下为作者在百度网盘中备份的相关文件(配置文件+演示文件)：

3 版本

python：3.6.3

4 依赖

(1)安装cv2依赖

pip install opencv-python

(2)安装numpy依赖

pip install numpy

4 代码

4.1 准备

(1)创建yolo配置文件

创建配置文件夹：cfg

将配置文件下载后，拷贝到cfg配置文件中：

(2)创建测试文件夹

测试文件夹：yoloSrc

输出文件夹：yoloRes

4.2 引入依赖

import cv2

import numpy as np

import os

import time

4.3 定义函数

def yolo_detect(pathIn='',

pathOut=None,

label_path='cfg/coco.names',

config_path='cfg/yolov3.cfg',

weights_path='cfg/yolov3.weights',

confidence_thre=0.5,

nms_thre=0.3,

jpg_quality=80):

'''

pathIn：原始图片的路径

pathOut：结果图片的路径

label_path：类别标签文件的路径

config_path：模型配置文件的路径

weights_path：模型权重文件的路径

confidence_thre：0-1，置信度(概率/打分)阈值，即保留概率大于这个值的边界框，默认为0.5

nms_thre：非极大值抑制的阈值，默认为0.3

jpg_quality：设定输出图片的质量，范围为0到100，默认为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')

# 载入图片并获取其维度

base_path = os.path.basename(pathIn)

img = cv2.imread(pathIn)

(H, W) = img.shape[:2]

# 加载模型配置和权重文件

print('从硬盘加载YOLO......')

net = cv2.dnn.readNetFromDarknet(config_path, weights_path)

# 获取YOLO输出层的名字

ln = net.getLayerNames()

ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]

# 将图片构建成一个blob，设置图片尺寸，然后执行一次

# YOLO前馈网络计算，最终获取边界框和相应概率

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('YOLO模型花费 {:.2f} 秒来预测一张图片'.format(end - start))

# 初始化边界框，置信度(概率)以及类别

boxes = []

confidences = []

classIDs = []

# 迭代每个输出层，总共三个

for output in layerOutputs:

# 迭代每个检测

for detection in output:

# 提取类别ID和置信度

scores = detection[5:]

classID = np.argmax(scores)

confidence = scores[classID]

# 只保留置信度大于某值的边界框

if confidence > confidence_thre:

# 将边界框的坐标还原至与原图片相匹配，记住YOLO返回的是

# 边界框的中心坐标以及边界框的宽度和高度

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)

# 确保至少一个边界框

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)

# 输出结果图片

if pathOut is None:

cv2.imwrite('with_box_' + base_path, img, [int(cv2.IMWRITE_JPEG_QUALITY), jpg_quality])

else:

cv2.imwrite(pathOut, img, [int(cv2.IMWRITE_JPEG_QUALITY), jpg_quality])

4.4 使用

if __name__ == '__main__':

yolo_detect('yoloSrc/trunk.jpg', 'yoloRes/trunk.jpg')

4.5 结果

输出日志：

从硬盘加载YOLO......

YOLO模型花费 1.75 秒来预测一张图片

(1)原文件

trunk.jpg

(2)输出结果

监测结果为truck，后面的0.999为对象匹配度。

trunk-2.jpg