python使用yolov3/yolov3-tiny训练好的权重文件.weights进行行人检测,批量测试自定义文件夹下的图片并输出至指定文件夹
目录
python使用yolov3/yolov3-tiny训练好的权重文件.weights进行行人检测,批量测试自定义文件夹下的图片并输出至指定文件夹
一、写在开头
最近在做毕业设计,研究的是目标识别与追踪,前段时间打算只用opencv识别个简单的目标就算了,但参考着论文硬着头皮撸了一星期的图片预处理,到了后面的识别部分,实在做不下去了,太南了!也不知道怎么具体去改其中的一些参数(数学不太好,害...),所以经过多方的baidu,用了yolov3跑训练模型,昨天刚跑完个只检测person的demo,买的东西还没有来到,也不知道实际效果怎么样,只是网上找的一些照片随便测试了一下,但效果其实一般般,看过几天实际效果怎么样吧,不太好的话就自己再做一部分数据加入到训练集中去,效果应该会好很多。所以,在此先记录下,如何使用已经训练好的权重文件来批量测试自定义文件夹下的照片并保存至指定文件夹。之前也找了这方面的教程,发现多数是去修改的.c的配置文件,我又懒,记性又不好,修改了以后一段时间我铁定忘记再改回来。所以,也就有了这篇blog。
二、已有的环境条件
1. pycharm--python
我在win10的pycharm中进行的测试,Ubuntu或其他环境也一样的。
2. opencv3.4
我win10和Ubuntu上都是用的opencv3.4,保持版本一致,可能会避免一些没必要的麻烦。
3. 用yolov3训练好了自己的权重文件.weights
由于我是深度学习方面的新手,所以暂时没有制作自己的数据集,直接用的VOC的数据集。在此程序中,需要用到训练好的三个文件,分别是voc.names,yolo3-tiny_person.cfg,yolov3-tiny_person800.weights。如下图所示,三个文件在yolo_voc文件夹下,
三、文件目录结构
如下所示:
py_person_yolov3
|
│
├─yolo_voc
│ voc.names
│ yolov3-tiny_person.cfg
│ yolov3-tiny_person800.weights
│
├─test-img1 # 要测试的图片文件夹,测试的图片都放在该文件夹下
│ 1.jpg
│ 2.jpg
│ 3.jpg
│ 4.jpg
│ 5.jpg
│ 6.jpg
│ 7.jpg
│ 8.jpg
│ 9.jpg
| 10.jpg
├─detect_person_fromDir.py
四、批量测试图片测试程序
detect_person_fromDir.py程序如下:
# -*- coding: utf-8 -*-
# 载入所需库
# author:cpz
# date:2020-04-10
"""
该程序在程序detect_person.py的基础上,改进了只能对一张照片进行目标检测,
可以实现对指定文件夹下的所有图片进行目标检测。
加入命令行参数--inDir——即待检测的所有图片的文件夹;
--outDir——即检测完之后的图片的保存路径。
"""
import cv2
import numpy as np
import os
import time
import argparse
def yolo_detect(inDir,
outDir,
label_path='./yolo_voc/voc.names',
config_path='./yolo_voc/yolov3-tiny_person.cfg',
weights_path='./yolo_voc/yolov3-tiny_person800.weights',
confidence_thre=0.3,
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,越大质量越好
'''
# 输出结果图片--->用绝对路径,用了一次相对路径,发现并没有结果
# os.getcwd()——用于返回当前工作目录
currentDir = os.getcwd()
print("当前工作路径为:%s" % (currentDir))
# 如果输出文件夹在命令行中未输入
# 则先创建一个文件夹:"outDir-from+输入照片的文件夹名字"
if outDir is None:
cDir = "outDir-from_" + inDir
# 判断是否已经存在了文件夹----> 存在的话就跳过,什么也不做
if os.path.exists(cDir):
print("由于未输入输出图片保存的文件夹,且已经存在了这个文件夹 %s" %(cDir))
pass
else:
os.mkdir(cDir)
print("由于未输入输出图片保存的文件夹,则创建了文件夹 %s" %(cDir))
# 加载类别标签文件
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')
"cpz在2020-4-10新改的,为了对一个目录中所有的图片进行检测"
# 试试只加载一次模型行不行
# 加载模型配置和权重文件
print('从硬盘加载YOLO......')
net = cv2.dnn.readNetFromDarknet(config_path, weights_path)
# 获取YOLO输出层的名字
ln = net.getLayerNames()
ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]
# os.listdir()——用于返回指定的文件夹包含的文件或文件夹的名字的列表
img_list = os.listdir(inDir)
for i in range(0, len(img_list)):
# os.path.join()——连接两个或更多的路径名组件
img_path = os.path.join(inDir, img_list[i])
# 载入图片并获取其维度
base_path = os.path.basename(img_path) # base_path为照片的名字(有扩展名)
img = cv2.imread(img_path)
(H, W) = img.shape[:2]
# 将图片构建成一个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)
# 要输出的图片的名称(含有扩展名)
outImg_name = 'with_box_' + base_path
outImg_path = os.path.join(currentDir, cDir, outImg_name)
cv2.imwrite(outImg_path, img, [int(cv2.IMWRITE_JPEG_QUALITY), jpg_quality])
# 输入了要输出的照片的文件夹
if __name__ == '__main__':
ap = argparse.ArgumentParser(description='Input the image to detect person, and the path to save the detected image.')
ap.add_argument("-i", "--inDir", required=True, help="path dir of input image")
ap.add_argument("-o", "--outDir", required=False, help="path to save detected image, default: ./outDir-from+inDir/")
args = vars(ap.parse_args())
# 从命令行输入中获得输入图片的路径
pathIn = args["inDir"]
# 从命令行输入中获得检测完之后的结果图片的保存路径
pathOut = args["outDir"]
# 进行检测咯!
yolo_detect(inDir=pathIn, outDir=pathOut)
五、进行测试
1. 打开pycharm的Terminal,并打开当前程序所在的文件夹,如下所示:
2. 运行程序,并输入所必需的参数--inDir:
python detect_person_fromDir.py --inDir test-img1
3. 运行结果:
4. 打开输出文件夹 outDir-from_test-img1,下面这两张是我挑的效果比较好的两张(先欺骗下自己吧):
六、写在最后
经过一段时间的学习,对yolov3的原理还每细看,这是我下面准备要做的。还有,感谢写blog分享经验的各位,用自己的一段时间写的blog真的可以帮到很多人,这也是我blog名字的由来。最后的最后,该程序是我在Ubuntu上搜到的一个程序,并做了进一步的改进(yolo_detect()函数是主要出自该博主),等我下次打开Ubuntu系统的时候我记得把原链接贴上,非常感谢这位博主。有找到的同学可以评论下,谢谢啦,就先写到这里了。