darknet模型转onnx后进行多进程推理
darknet转换为onnx后进行多进程推理脚本
from re import I
import onnxruntime
import numpy as np
import cv2
import os,time,logging,multiprocessing
label=["background","left","upperleft", "up", "upperright", "fen", "right", "he", "downleft", "downright", "down"]
def sigmoid(x):
s = 1 / (1 + np.exp(-1*x))
return s
def process_image(img_path,model_h,model_w):
img = cv2.imread(img_path)
image = cv2.resize(img, (model_h, model_w))
image = image[:,:,::-1].transpose((2,0,1))
image = image[np.newaxis,:,:,:]/255
image = np.array(image,dtype=np.float32)
return img,image
def getMaxClassScore(class_scores):
class_score = 0
class_index = 0
for i in range(len(class_scores)):
if class_scores[i] > class_score:
class_index = i+1
class_score = class_scores[i]
return class_score,class_index
def getBBox(feat,anchors,image_shape,confidence_threshold,classes):
box = []
total=classes+5
for i in range(len(anchors)):
for cx in range(feat.shape[0]):
for cy in range(feat.shape[1]):
tx = feat[cx][cy][0 + total * i]
ty = feat[cx][cy][1 + total * i]
tw = feat[cx][cy][2 + total * i]
th = feat[cx][cy][3 + total * i]
cf = feat[cx][cy][4 + total * i]
cp = feat[cx][cy][5 + total * i:total + total * i]
bx = (sigmoid(tx) + cx)/feat.shape[0]
by = (sigmoid(ty) + cy)/feat.shape[1]
bw = anchors[i][0]*np.exp(tw)/image_shape[0]
bh = anchors[i][1]*np.exp(th)/image_shape[1]
b_confidence = sigmoid(cf)
b_class_prob = sigmoid(cp)
b_scores = b_confidence*b_class_prob
b_class_score,b_class_index = getMaxClassScore(b_scores)
if b_class_score > confidence_threshold:
box.append([bx,by,bw,bh,b_class_score,b_class_index])
return box
def donms(boxes,nms_threshold,iou_min_area):
b_x = boxes[:, 0]
b_y = boxes[:, 1]
b_w = boxes[:, 2]
b_h = boxes[:, 3]
scores = boxes[:,4]
areas = (b_w)*(b_h)
order = scores.argsort()[::-1]
# print(order,'----------------order-----------------')
keep = []
while order.size > 0:
i = order[0]
keep.append(i)
xx1 = np.maximum(b_x[i], b_x[order[1:]])
yy1 = np.maximum(b_y[i], b_y[order[1:]])
xx2 = np.minimum(b_x[i] + b_w[i], b_x[order[1:]] + b_w[order[1:]])
yy2 = np.minimum(b_y[i] + b_h[i], b_y[order[1:]] + b_h[order[1:]])
w = np.maximum(0.0, xx2 - xx1)
h = np.maximum(0.0, yy2 - yy1)
inter = w * h
if iou_min_area:
ovr = inter / np.minimum(areas[i], areas[order[1:]])
else:
ovr = inter / (areas[i] + areas[order[1:]] - inter)
inds = np.where(ovr <= nms_threshold)[0]
order = order[inds + 1]
final_boxes = [boxes[i] for i in keep]
return final_boxes
def drawBox(img_path,boxes,img,model_h,model_w):
img_size = img.shape
img_h=img_size[0]
img_w=img_size[1]
m_scale_w = (float)(img_w) /(float)(model_w)
m_scale_h = (float)(img_h) /(float)(model_h)
img_name = os.path.basename(img_path)
for box in boxes:
x1 = int((box[0]-box[2]/2)*(float)(model_w)*m_scale_w)
y1 = int((box[1]-box[3]/2)*(float)(model_h)*m_scale_h)
x2 = int((box[0]+box[2]/2)*(float)(model_w)*m_scale_w)
y2 = int((box[1]+box[3]/2)*(float)(model_h)*m_scale_h)
cv2.rectangle(img,(x1,y1),(x2,y2),(0,255,0),2)
cv2.putText(img, label[int(box[5])]+":"+str(round(box[4],3)), (x1+5,y1+10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
print('----------draw result images start----------------------------',img_name)
cv2.imwrite("result/"+img_name,img)
def getBoxes(prediction,confidence_threshold,nms_threshold,model_h,model_w,classes,anchors):
boxes = []
for i in range(len(prediction)):
feature_map = prediction[i][0].transpose((2, 1, 0))
box = getBBox(feature_map, anchors[i], [model_h, model_w], confidence_threshold,classes)
boxes.extend(box)
print('the len of result:\n',len(boxes))
boxes = donms(np.array(boxes),nms_threshold,False)
print('the len of result:\n',len(boxes))
return boxes
def worker_proc(gpu_id, process_id, task_queue):
model_h=608
model_w=608
classes=10
confidence_th=0.1
nms_th=0.6
anchors = [[(110.5563,261.7314), (259.1454,131.4648), (278.3166,279.8123)], [(125.9803,94.2299), (85.8488,157.2515), (55.6883,253.7204)],[(33.6239,37.5227), (46.3808,75.3617), (82.3125,94.6357)]]
print('----------------the length of queue--------',task_queue.qsize())
while task_queue.qsize() > 0:
img_inum, img_num, img_path = task_queue.get(True, 1)
print('{:d}/{:d}: Being processed is {:}'.format(img_inum, img_num, img_path))
img,TestData = process_image(img_path,model_h,model_w)
# session = onnxruntime.InferenceSession("./onnx_models/xn.onnx", providers=['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider'])
session = onnxruntime.InferenceSession("./onnx_models/xn.onnx", providers=['CUDAExecutionProvider'])
# session = onnxruntime.InferenceSession("./onnx_models/xn.onnx",providers=['CPUExecutionProvider'])
inname = [input.name for input in session.get_inputs()][0]
outname = [output.name for output in session.get_outputs()]
print('----------------the length of end queue--------',task_queue.qsize())
print("inputs name:",inname,"outputs name:",outname)
s = time.time()
prediction = session.run(outname, {inname:TestData})
print('-------------the model output--------------\n',prediction)
print("推理照片 %s 耗时：% .2fms" % (img_path, ((time.time() - s)*1000)))
boxes = getBoxes(prediction,confidence_th,nms_th,model_h,model_w,classes,anchors)
drawBox(img_path,boxes,img,model_h,model_w)
if __name__ == '__main__':
multiprocessing = multiprocessing.get_context("spawn")
gpu_num = 1
process_num = 8
task_queue = multiprocessing.Queue()
img_paths = os.listdir('test')
img_num = len(img_paths)
img_inum = 0
for i in img_paths:
img_inum = img_inum + 1
task_queue.put((img_inum, img_num, os.path.join('test',i)))
# start multiprocessing
starttime = time.time()
worker_list = []
for gpu_id in range(0, gpu_num):
for process_id in range(0, process_num):
worker = multiprocessing.Process(
target = worker_proc,
args = (gpu_id, process_id, task_queue))
worker.start()
worker_list.append(worker)
for worker in worker_list:
worker.join()
print('------------------game is over------------------------------------------------')
多模型多进程
import onnxruntime
import numpy as np
import cv2
import os,time,logging,multiprocessing
from shutil import copyfile
def sigmoid(x):
s = 1 / (1 + np.exp(-1*x))
return s
def process_image(img_path,model_h,model_w):
img = cv2.imread(img_path)
image = cv2.resize(img, (model_h, model_w))
image = image[:,:,::-1].transpose((2,0,1))
image = image[np.newaxis,:,:,:]/255
image = np.array(image,dtype=np.float32)
#返回原图像和处理后的数组
return img,image
def getMaxClassScore(class_scores):
class_score = 0
class_index = 0
for i in range(len(class_scores)):
if class_scores[i] > class_score:
class_index = i+1
class_score = class_scores[i]
return class_score,class_index
def getBBox(feat,anchors,image_shape,confidence_threshold,classes):
box = []
total=classes+5
for i in range(len(anchors)):
for cx in range(feat.shape[0]):
for cy in range(feat.shape[1]):
tx = feat[cx][cy][0 + total * i]
ty = feat[cx][cy][1 + total * i]
tw = feat[cx][cy][2 + total * i]
th = feat[cx][cy][3 + total * i]
cf = feat[cx][cy][4 + total * i]
cp = feat[cx][cy][5 + total * i:total + total * i]
bx = (sigmoid(tx) + cx)/feat.shape[0]
by = (sigmoid(ty) + cy)/feat.shape[1]
bw = anchors[i][0]*np.exp(tw)/image_shape[0]
bh = anchors[i][1]*np.exp(th)/image_shape[1]
b_confidence = sigmoid(cf)
b_class_prob = sigmoid(cp)
b_scores = b_confidence*b_class_prob
b_class_score,b_class_index = getMaxClassScore(b_scores)
if b_class_score > confidence_threshold:
box.append([bx,by,bw,bh,b_class_score,b_class_index])
return box
def donms(boxes,nms_threshold):
b_x = boxes[:, 0]
b_y = boxes[:, 1]
b_w = boxes[:, 2]
b_h = boxes[:, 3]
scores = boxes[:,4]
areas = (b_w)*(b_h)
order = scores.argsort()[::-1]
keep = [] # 保留的结果框集合
while order.size > 0:
i = order[0]
keep.append(i) # 保留该类剩余box中得分最高的一个
# 得到相交区域,左上及右下
xx1 = np.maximum(b_x[i], b_x[order[1:]])
yy1 = np.maximum(b_y[i], b_y[order[1:]])
xx2 = np.minimum(b_x[i] + b_w[i], b_x[order[1:]] + b_w[order[1:]])
yy2 = np.minimum(b_y[i] + b_h[i], b_y[order[1:]] + b_h[order[1:]])
#相交面积,不重叠时面积为0
w = np.maximum(0.0, xx2 - xx1)
h = np.maximum(0.0, yy2 - yy1)
inter = w * h
#相并面积,面积1+面积2-相交面积
union = areas[i] + areas[order[1:]] - inter
# 计算IoU：交 /（面积1+面积2-交）
IoU = inter / union
# 保留IoU小于阈值的box
inds = np.where(IoU <= nms_threshold)[0]
order = order[inds + 1] # 因为IoU数组的长度比order数组少一个,所以这里要将所有下标后移一位
final_boxes = [boxes[i] for i in keep]
return final_boxes
def output_xml(result_list, src_img_name,img_shape, xml_path):
"""output result to xml
Args:
result_list (list): [ (code,score,xmin,ymin,xmax,ymax), ]
src_img_name (str): source image name.
img_shape(list): img_shape
xml_path: xml path
Returns:
None
"""
if 0 == len(result_list):
return None
xml_fmt = '''
VOC2007
{:}
{:d}
{:d}
{:d}
{:d}
{:} '''
object_fmt = ''''''
h,w,c = img_shape
objects_str = ""
for inx, res in enumerate(result_list):
item_str = object_fmt.format(
(inx + 1),
res[0],
int(res[2]),
int(res[3]),
int(res[4]),
int(res[5])
)
objects_str = objects_str + item_str
xml_str = xml_fmt.format(
src_img_name,
w,
h,
c,
len(result_list),
objects_str
)
with open(xml_path, "w") as outfile:
outfile.write(xml_str)
def drawBox(obj_path,item,labels,boxes,img,model_h,model_w):
image_name = os.path.basename(obj_path)
img_size = img.shape
img_h=img_size[0]
img_w=img_size[1]
m_scale_w = (float)(img_w) /(float)(model_w)
m_scale_h = (float)(img_h) /(float)(model_h)
result_list = []
for box in boxes:
x1 = int((box[0]-box[2]/2)*(float)(model_w)*m_scale_w)
y1 = int((box[1]-box[3]/2)*(float)(model_h)*m_scale_h)
x2 = int((box[0]+box[2]/2)*(float)(model_w)*m_scale_w)
y2 = int((box[1]+box[3]/2)*(float)(model_h)*m_scale_h)
confidence = float(round(box[4],3))
code = labels[int(box[5])]
result = [code,confidence,x1,y1,x2,y2]
result_list.append(result)
print('---------------result of the %s---------------'%image_name, result_list)
# 生成xml
xml_output_folder = "./suzhou1/xml_{}/".format(item)
if not os.path.exists(xml_output_folder):
os.mkdir(xml_output_folder)
if (xml_output_folder is not None) and (os.path.exists(xml_output_folder)):
copyfile(obj_path, os.path.join(xml_output_folder,image_name))
output_xml(result_list, image_name, img.shape, os.path.join(xml_output_folder, os.path.splitext(image_name)[0] + ".xml"))
record = 'suzhou/record.txt'
f = open(record,'a')
f.write(image_name+str(result_list)+'\n')
def getBoxes(prediction,confidence_threshold,model_h,model_w,classes,anchors):
boxes = []
for i in range(len(prediction)):
feature_map = prediction[i][0].transpose((2, 1, 0))
box = getBBox(feature_map, anchors[i], [model_h, model_w], confidence_threshold,classes)
boxes.extend(box)
return boxes
def parse_txt_list(params):
anchors = params[1].split('=')[1]
alist = [int(float(anchor)) for anchor in anchors[1:-2].split(',')]
anchors = [[(alist[0],alist[1]), (alist[2],alist[3]), (alist[4],alist[5])], [(alist[6],alist[7]), (alist[8],alist[9]), (alist[10],alist[11])], [(alist[12],alist[13]), (alist[14],alist[15]), (alist[16],alist[17])]]
model_w = int(params[3].split('=')[1].strip())
model_h = int(params[4].split('=')[1].strip())
classes = int(params[5].split('=')[1].strip())
confidence_th = float(params[6].split('=')[1].strip())
confidence_th = 0.8
nms_th = float(params[7].split('=')[1].strip())
labels = params[8].split('=')[1].strip()
labels = [i.strip()[1:-1] for i in labels[1:-1].split(',')]
labels.insert(0,'background')
return anchors,model_w,model_h,classes,confidence_th,nms_th,labels
def work_proc(gpu_id,process_id,model_to_scan,path_to_scan,path_list):
# params_list anchors,model_w,model_h,classes,confidence_th,nms_th,labels
device_str = 'export CUDA_VISIBLE_DEVICES={:d}'.format(gpu_id)
os.system(device_str)
for order in range(0,len(model_to_scan)):
item = model_to_scan[order]
onnxruntime.set_default_logger_severity(3)
session = onnxruntime.InferenceSession("./onnx_models/{}.onnx".format(item),providers=['CUDAExecutionProvider'])
with open("./onnx_models/{}.txt".format(item)) as f:
params = f.readlines()
params_list = parse_txt_list(params)
anchors,model_w,model_h,classes,confidence_th,nms_th,labels = params_list
image_num = len(path_list)
image_inum = 0
for image in path_list:
if image[-4:] == ".jpg":
obj_path = os.path.join(path_to_scan, image)
logging.critical('----------being detected images is %s---%d/%d----------%s'%(item,image_inum, image_num, image))
img,TestData = process_image(obj_path,model_h,model_w)
inname = [input.name for input in session.get_inputs()][0]
outname = [output.name for output in session.get_outputs()]
prediction = session.run(outname, {inname:TestData})
boxes = getBoxes(prediction,confidence_th,model_h,model_w,classes,anchors)
image_inum = image_inum + 1
print('---------boxes of %s--------'%(obj_path),boxes)
# while task_queue.qsize() > 0:
# img_inum, img_num, obj_path = task_queue.get(True,1)
# logging.critical('----------being detected images is %s---%d/%d----------'%(item,img_inum,img_num))
# img,TestData = process_image(obj_path,model_h,model_w)
# inname = [input.name for input in session.get_inputs()][0]
# outname = [output.name for output in session.get_outputs()]
# prediction = session.run(outname, {inname:TestData})
# boxes = getBoxes(prediction,confidence_th,model_h,model_w,classes,anchors)
# print('---------boxes of %s--------'%(obj_path),boxes)
if len(boxes)==0:
continue
else:
boxes = donms(np.array(boxes),nms_th)
drawBox(obj_path,item,labels,boxes,img,model_h,model_w)
if __name__ == '__main__':
data_dir = "sz"
gpu_num = 1
process_num = 8
multiprocessing = multiprocessing.get_context('spawn')
for root,dirs,files in os.walk(data_dir):
for name in dirs:
path_to_scan = os.path.join(root, name)
model_to_scan = [ "num"]
# model_to_scan = [ "oil", "cn", "xn", "yb", "light","tyfh", "110"]
path_list = os.listdir(path_to_scan)
# img_num = len(path_list)
# img_inum = 0
# task_queue = multiprocessing.Queue()
# for jpg in path_list:
# if jpg[-4:] == '.jpg':
# obj_path = os.path.join(path_to_scan, jpg)
# img_inum = img_inum + 1
# task_queue.put((img_inum, img_num, obj_path))
worker_list = []
for i in range(0,gpu_num):
for j in range(0,process_num):
worker = multiprocessing.Process(target=work_proc,args=(i,j,model_to_scan,path_to_scan,path_list))
worker.start()
worker_list.append(worker)
for worker in worker_list:
worker.join()