|
|
from __future__ import division
|
|
|
|
|
|
import time
|
|
|
import torch
|
|
|
import torch.nn as nn
|
|
|
from torch.autograd import Variable
|
|
|
import numpy as np
|
|
|
import cv2
|
|
|
from util import *
|
|
|
import argparse
|
|
|
import os
|
|
|
import os.path as osp
|
|
|
from darknet import Darknet
|
|
|
import pickle as pkl
|
|
|
import pandas as pd
|
|
|
import random
|
|
|
|
|
|
# 命令行参数
|
|
|
def arg_parse():
|
|
|
|
|
|
parser = argparse.ArgumentParser(description='YOLO v3 Detection Module')
|
|
|
|
|
|
# images（用于指定输入图像或图像目录）
|
|
|
parser.add_argument("--images", dest = 'images', help =
|
|
|
"Image / Directory containing images to perform detection upon",
|
|
|
default = "imgs", type = str)
|
|
|
# det（保存检测结果的目录）
|
|
|
parser.add_argument("--det", dest = 'det', help =
|
|
|
"Image / Directory to store detections to",
|
|
|
default = "det", type = str)
|
|
|
# batch大小
|
|
|
parser.add_argument("--bs", dest = "bs", help = "Batch size", default = 1)
|
|
|
# objectness置信度
|
|
|
parser.add_argument("--confidence", dest = "confidence", help = "Object Confidence to filter predictions", default = 0.5)
|
|
|
# NMS阈值
|
|
|
parser.add_argument("--nms_thresh", dest = "nms_thresh", help = "NMS Threshhold", default = 0.4)
|
|
|
# cfg（替代配置文件）
|
|
|
parser.add_argument("--cfg", dest = 'cfgfile', help =
|
|
|
"Config file",
|
|
|
default = "cfg/yolov3.cfg", type = str)
|
|
|
parser.add_argument("--weights", dest = 'weightsfile', help =
|
|
|
"weightsfile",
|
|
|
default = "yolov3.weights", type = str)
|
|
|
# reso（输入图像的分辨率，可用于在速度与准确度之间的权衡）
|
|
|
parser.add_argument("--reso", dest = 'reso', help =
|
|
|
"Input resolution of the network. Increase to increase accuracy. Decrease to increase speed",
|
|
|
default = "416", type = str)
|
|
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
|
|
args = arg_parse()
|
|
|
images = args.images
|
|
|
batch_size = int(args.bs)
|
|
|
confidence = float(args.confidence)
|
|
|
nms_thesh = float(args.nms_thresh)
|
|
|
start = 0
|
|
|
CUDA = torch.cuda.is_available()
|
|
|
num_classes = 80 # COCO数据集中目标的名称
|
|
|
classes = load_classes("data/coco.names")
|
|
|
|
|
|
# 初始化网络，加载权重
|
|
|
print("正在加载网络QAQ")
|
|
|
model = Darknet(args.cfgfile)
|
|
|
model.load_weights(args.weightsfile)
|
|
|
print("网络加载成功QvQ")
|
|
|
|
|
|
model.net_info["height"] = args.reso
|
|
|
inp_dim = int(model.net_info["height"])
|
|
|
assert inp_dim % 32 == 0
|
|
|
assert inp_dim > 32
|
|
|
|
|
|
# GPU加速
|
|
|
if CUDA:
|
|
|
model.cuda()
|
|
|
|
|
|
# 模型评估
|
|
|
model.eval()
|
|
|
|
|
|
# 从磁盘读取图像或从目录读取多张图像，图像路径imlist
|
|
|
read_dir = time.time() # 测量时间的检查点
|
|
|
# 检测阶段
|
|
|
try:
|
|
|
imlist = [osp.join(osp.realpath('.'), images, img) for img in os.listdir(images)]
|
|
|
except NotADirectoryError:
|
|
|
imlist = []
|
|
|
imlist.append(osp.join(osp.realpath('.'), images))
|
|
|
except FileNotFoundError:
|
|
|
print("没有找到{}文件或目录QwQ".format(images))
|
|
|
exit()
|
|
|
|
|
|
# 如果没有保存检测结果的目录，就创建一个
|
|
|
if not os.path.exists(args.det):
|
|
|
os.makedirs(args.det)
|
|
|
|
|
|
# 用OpenCV加载多张图片图像
|
|
|
load_batch = time.time()
|
|
|
loaded_ims = [cv2.imread(x) for x in imlist]
|
|
|
# 转成PyTorch图像格式
|
|
|
im_batches = list(map(prep_image, loaded_ims, [inp_dim for x in range(len(imlist))]))
|
|
|
# 包含原始图像的维度的列表
|
|
|
im_dim_list = [(x.shape[1], x.shape[0]) for x in loaded_ims]
|
|
|
im_dim_list = torch.FloatTensor(im_dim_list).repeat(1,2)
|
|
|
|
|
|
# 创建batch
|
|
|
leftover = 0
|
|
|
if (len(im_dim_list) % batch_size):
|
|
|
leftover = 1
|
|
|
if batch_size != 1:
|
|
|
num_batches = len(imlist) // batch_size + leftover
|
|
|
im_batches = [torch.cat((im_batches[i*batch_size : min((i+1)*batch_size,
|
|
|
len(im_batches))])) for i in range(num_batches)]
|
|
|
|
|
|
write = 0
|
|
|
|
|
|
if CUDA:
|
|
|
im_dim_list = im_dim_list.cuda()
|
|
|
|
|
|
start_det_loop = time.time()
|
|
|
for i, batch in enumerate(im_batches):
|
|
|
# 载入图片
|
|
|
start = time.time()
|
|
|
if CUDA:
|
|
|
batch = batch.cuda()
|
|
|
with torch.no_grad():
|
|
|
prediction = model(Variable(batch), CUDA)
|
|
|
prediction = write_results(prediction, confidence, num_classes, nms_conf=nms_thesh)
|
|
|
|
|
|
end = time.time()
|
|
|
|
|
|
if type(prediction) == int:
|
|
|
for im_num, image in enumerate(imlist[i*batch_size: min((i + 1)*batch_size, len(imlist))]):
|
|
|
im_id = i*batch_size + im_num
|
|
|
print("{0:20s} 预测用时{1:6.3f} 秒".format(image.split("/")[-1], (end - start)/batch_size))
|
|
|
print("{0:20s} {1:s}".format("检测到的对象：", " "))
|
|
|
print("----------------------------------------------------------")
|
|
|
continue
|
|
|
|
|
|
prediction[:,0] += i*batch_size # 将batch索引转换成imlist索引
|
|
|
|
|
|
if not write: # 初始化output
|
|
|
output = prediction
|
|
|
write = 1
|
|
|
else:
|
|
|
output = torch.cat((output, prediction))
|
|
|
|
|
|
for im_num, image in enumerate(imlist[i*batch_size:min((i+1)*batch_size, len(imlist))]):
|
|
|
im_id = i*batch_size + im_num
|
|
|
objs = [classes[int(x[-1])] for x in output if int(x[0]) == im_id]
|
|
|
print("{0:20s} 预测用时{1:6.3f} 秒".format(image.split("/")[-1], (end - start)/batch_size))
|
|
|
print("{0:20s} {1:s}".format("检测到的对象:", " ".join(objs)))
|
|
|
print("----------------------------------------------------------------")
|
|
|
|
|
|
if CUDA:
|
|
|
torch.cuda.synchronize() # 保证CUDA核与CPU同步
|
|
|
|
|
|
# 在图像上绘制边界框
|
|
|
try:
|
|
|
output
|
|
|
except NameError:
|
|
|
print("不存在检测结果TAT")
|
|
|
exit()
|
|
|
|
|
|
# 输出边界框对应网络输入大小，需要将边界框属性转换到图像的原始尺寸
|
|
|
im_dim_list = torch.index_select(im_dim_list, 0, output[:,0].long())
|
|
|
scaling_factor = torch.min(inp_dim/im_dim_list,1)[0].view(-1,1)
|
|
|
|
|
|
output[:,[1,3]] -= (inp_dim - scaling_factor*im_dim_list[:,0].view(-1,1))/2
|
|
|
output[:,[2,4]] -= (inp_dim - scaling_factor*im_dim_list[:,1].view(-1,1))/2
|
|
|
output[:,1:5] /= scaling_factor
|
|
|
|
|
|
for i in range(output.shape[0]):
|
|
|
output[i, [1,3]] = torch.clamp(output[i, [1,3]], 0.0, im_dim_list[i,0])
|
|
|
output[i, [2,4]] = torch.clamp(output[i, [2,4]], 0.0, im_dim_list[i,1])
|
|
|
|
|
|
output_recast = time.time()
|
|
|
|
|
|
# 随机选择的颜色用于绘制边界框
|
|
|
class_load = time.time()
|
|
|
colors = pkl.load(open("pallete", "rb"))
|
|
|
|
|
|
# 开始绘制边界框
|
|
|
draw = time.time()
|
|
|
|
|
|
# 绘制边界框:从colors中随机选颜色绘制矩形框
|
|
|
# 边界框左上角创建一个填充后的矩形，写入该框位置检测到的目标的类别
|
|
|
def write(x, results):
|
|
|
|
|
|
c1 = tuple(x[1:3].int())
|
|
|
c2 = tuple(x[3:5].int())
|
|
|
img = results[int(x[0])]
|
|
|
cls = int(x[-1])
|
|
|
color = random.choice(colors)
|
|
|
label = "{0}".format(classes[cls])
|
|
|
cv2.rectangle(img, c1, c2, color, 1)
|
|
|
t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 1, 1)[0]
|
|
|
c2 = c1[0] + t_size[0] + 3, c1[1] + t_size[1] + 4
|
|
|
cv2.rectangle(img, c1, c2, color, -1) # -1表示填充的矩形
|
|
|
cv2.putText(img, label, (c1[0], c1[1] + t_size[1] + 4), cv2.FONT_HERSHEY_PLAIN, 1, [225,225,225], 1)
|
|
|
|
|
|
return img
|
|
|
|
|
|
list(map(lambda x:write(x, loaded_ims), output))
|
|
|
|
|
|
# 保存检测结果图像,det_图像名
|
|
|
det_names = pd.Series(imlist).apply(lambda x: "{}/det_{}".format(args.det, x.split("/")[-1]))
|
|
|
# 将带有检测结果的图像写入det_names中的地址
|
|
|
list(map(cv2.imwrite, det_names, loaded_ims))
|
|
|
end = time.time()
|
|
|
|
|
|
# 显示输出时间的总结
|
|
|
print("总结")
|
|
|
print("----------------------------------------------------------------")
|
|
|
print("{:25s} {}".format("任务", "所用时间(s)"))
|
|
|
print()
|
|
|
print("{:25s} {:2.3f}".format("读入目录", load_batch - read_dir))
|
|
|
print("{:25s} {:2.3f}".format("加载batch", start_det_loop - load_batch))
|
|
|
print("{:25s} {:2.3f}".format("检测(" + str(len(imlist)) + "张图)", output_recast - start_det_loop))
|
|
|
print("{:25s} {:2.3f}".format("输出处理", class_load - output_recast))
|
|
|
print("{:25s} {:2.3f}".format("绘制边界框", end - draw))
|
|
|
print("{:25s} {:2.3f}".format("平均检测时间", (end - load_batch)/len(imlist)))
|
|
|
print("----------------------------------------------------------------")
|
|
|
|
|
|
torch.cuda.empty_cache() |
