dl----darknet模型转onnx推理部署（python）

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 = '''<annotation>
<folder>VOC2007</folder>
<filename>{:}</filename>
<size>
<width>{:d}</width>
<height>{:d}</height>
<depth>{:d}</depth>
</size>
<objectsum>{:d}</objectsum>
{:}</annotation>'''

    object_fmt = '''<object>
<Serial>{:d}</Serial>
<name>{:}</name>
<bndbox>
<xmin>{:d}</xmin>
<ymin>{:d}</ymin>
<xmax>{:d}</xmax>
<ymax>{:d}</ymax>
</bndbox>
</object>'''
    
    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()