yolov3:
# -*- coding: utf-8 -*
import numpy as np
import cv2 as cv
import os
import time
yolo_dir = '/home/ubuntu/model/yolo4-tiny' # YOLO文件路径
weightsPath = os.path.join(yolo_dir, 'yolov3-test.weights') # 权重文件
configPath = os.path.join(yolo_dir, 'yolov3-test.cfg') # 配置文件
labelsPath = os.path.join(yolo_dir, 'test.names') # label名称
imgPath = os.path.join(yolo_dir, 'test.jpg') # 测试图像
CONFIDENCE = 0.25 # 过滤弱检测的最小概率
THRESHOLD = 0.2 # 非最大值抑制阈值
#批量图片检测和保存位置
test_dir = '/home/ubuntu/model/yolo4-tiny/imgs'
save_dir = '/home/ubuntu/model/yolo4-tiny/imgsResultV3'
# 加载网络、配置权重
net = cv.dnn.readNetFromDarknet(configPath, weightsPath) # # 利用下载的文件
#批量加载图像
pics = os.listdir(test_dir)
start = time.time()
minTime = 10;
maxTime = 0;
for im in pics:
# 加载图片、转为blob格式、送入网络输入层
s = time.time()
imgPath = os.path.join(test_dir,im)
img = cv.imread(imgPath)
blobImg = cv.dnn.blobFromImage(img, 1.0/255.0, (416, 416), None, True, False) # # net需要的输入是blob格式的,用blobFromImage这个函数来转格式
net.setInput(blobImg) # # 调用setInput函数将图片送入输入层
# 获取网络输出层信息(所有输出层的名字),设定并前向传播
outInfo = net.getUnconnectedOutLayersNames() # # 前面的yolov3架构也讲了,yolo在每个scale都有输出,outInfo是每个scale的名字信息,供net.forward使用
layerOutputs = net.forward(outInfo) # 得到各个输出层的、各个检测框等信息,是二维结构。
# 拿到图片尺寸
(H, W) = img.shape[:2]
# 过滤layerOutputs
# layerOutputs的第1维的元素内容: [center_x, center_y, width, height, objectness, N-class score data]
# 过滤后的结果放入:
boxes = [] # 所有边界框(各层结果放一起)
confidences = [] # 所有置信度
classIDs = [] # 所有分类ID
# # 1)过滤掉置信度低的框框
for out in layerOutputs: # 各个输出层
for detection in out: # 各个框框
# 拿到置信度
scores = detection[5:] # 各个类别的置信度
classID = np.argmax(scores) # 最高置信度的id即为分类id
confidence = scores[classID] # 拿到置信度
# 根据置信度筛查
if confidence > CONFIDENCE:
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)
# # 2)应用非最大值抑制(non-maxima suppression,nms)进一步筛掉
idxs = cv.dnn.NMSBoxes(boxes, confidences, CONFIDENCE, THRESHOLD) # boxes中,保留的box的索引index存入idxs
# 得到labels列表
with open(labelsPath, 'rt') as f:
labels = f.read().rstrip('\n').split('\n')
# 应用检测结果
np.random.seed(42)
COLORS = np.random.randint(0, 255, size=(len(labels), 3), dtype="uint8") # 框框显示颜色,每一类有不同的颜色,每种颜色都是由RGB三个值组成的,所以size为(len(labels), 3)
if len(idxs) > 0:
for i in idxs.flatten(): # indxs是二维的,第0维是输出层,所以这里把它展平成1维
(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]]]
cv.rectangle(img, (x, y), (x+w, y+h), color, 2) # 线条粗细为2px
text = "{}: {:.4f}".format(labels[classIDs[i]], confidences[i])
cv.putText(img, text, (x, y-5), cv.FONT_HERSHEY_SIMPLEX, 0.5, color, 2) # cv.FONT_HERSHEY_SIMPLEX字体风格、0.5字体大小、粗细2px
print(imgPath+"检测耗时: {:.6f} seconds".format(time.time() - s)) # # 可以打印下信息
times = time.time() - s
if(minTime>times):
minTime = times
if(maxTime<times):
maxTime = times
cv.imwrite(os.path.join(save_dir,im),img)
print("yolov3检测100张图总共耗时:%.6f秒" % (time.time() - start))
print("耗时最长为:%.6f秒" % (maxTime))
print("耗时最短为:%.6f秒" % (minTime))
yolov4:
# -*- coding: utf-8 -*
import cv2 as cv
import time
import os
#置信度和nms设置
confThreshold = 0.25
nmsThreshold = 0.2
class_names = []#初始化一个列表以存储类名
COLORS = [(0, 255, 255),(0,0,255), (255, 255, 0),(0, 255, 0), (255, 0, 0)]#颜色
color=(0,255,0)
color1=(0,0,255)
c=(0, 0, 0)
#模型参数
weights="/home/ubuntu/model/yolo4-tiny/yolov4-tiny.weights"
cfg="/home/ubuntu/model/yolo4-tiny/yolov4-tiny.cfg"
m="/home/ubuntu/model/yolo4-tiny/coco.names"
#批量图片检测和保存位置
test_dir = '/home/ubuntu/model/yolo4-tiny/imgs'
save_dir = '/home/ubuntu/model/yolo4-tiny/imgsResultV4'
# 网络设置
net = cv.dnn_DetectionModel(cfg, weights)
#GPU运行
net.setPreferableBackend(cv.dnn.DNN_BACKEND_CUDA)
net.setPreferableTarget(cv.dnn.DNN_TARGET_CUDA)
net.setInputSize(416, 416)
net.setInputScale(1.0 / 255)
net.setInputSwapRB(True)
with open(m, "r") as f:
class_names = [cname.strip() for cname in f.readlines()]
#批量加载图像路径
pics = os.listdir(test_dir)
start = time.time()
minTime = 10;
maxTime = 0;
inImgTime=0;
testImgTime=0;
for im in pics:
s = time.time()
#加载图片
img = os.path.join(test_dir,im)
image = cv.imread(img)
sTest = time.time()
inImg = sTest-s
inImgTime +=inImg
#模型检测
classes, confidences, boxes = net.detect(image, confThreshold, nmsThreshold)
sTestEnd = time.time()
testImg = sTestEnd-sTest
testImgTime+= testImg
count=0
#循环画框
for (classid, score, box) in zip(classes, confidences, boxes):
label = "%s(%2.0f%%)" % (class_names[classid[0]], score*100)#标签置信度
labelSize, baseLine = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, 0.5, 1)
left, top, width, height = box
cv.rectangle(image, box, color1, 1)
cv.rectangle(image, (left-1, top - labelSize[1]-5), (left + labelSize[0], top), color, cv.FILLED)
cv.putText(image, label, (left, top-5), cv.FONT_HERSHEY_SIMPLEX, 0.5, c,2)
count+=1
times = time.time() - s
if(minTime>times):
minTime = times
if(maxTime<times):
maxTime = times
print(img+"加载图片耗时:%.6f秒" % inImg)
print(img+"检测图片耗时:%.6f秒" % testImg)
print(img+"检测耗时:%.3f秒" % times)
cv.putText(image,"Num: %s" % str(count),(25,50),cv.FONT_HERSHEY_SIMPLEX,1,[255,255,0],2)
cv.imwrite(os.path.join(save_dir,im),image)#用原始图像名im保存转换后的图片
print("yolov4检测6张图总共耗时:%.3f秒" % (time.time() - start))
print("耗时最长为:%.3f秒" % (maxTime))
print("耗时最短为:%.3f秒" % (minTime))
print("加载图片总共耗时:%.6f秒" % (inImgTime))
print("检测图片总共耗时:%.6f秒" % (testImgTime))