Algorithm-Repo/SourceCode/MachineLearningUtils/ObjectDetection/YOLO_V4/demo_darknet2onnx.py

import sys
import onnx
import os
import argparse
import numpy as np
import cv2
import onnxruntime
import multiprocessing

from tool.utils import *
from tool.darknet2onnx import *


def main():


    onnx_path = 'C:\\Users\\VINNO\\Desktop\\新建文件夹 (2)\\pytorch-YOLOv4-master\\20210824_yolov4_1_320_320_static11.onnx'

    sess_options = onnxruntime.SessionOptions()
    # sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
    #
    #控制用于运行模型的线程数 controls the number of threads to use to run the model
    sess_options.intra_op_num_threads = 1
    #
    # #When sess_options.execution_mode = rt.ExecutionMode.ORT_PARALLEL,
    # # you can set sess_options.inter_op_num_threads to control the number of threads used to parallelize the execution of the graph (across nodes).
    # sess_options.execution_mode = onnxruntime.ExecutionMode.ORT_PARALLEL
    # sess_options.inter_op_num_threads = 1

    session = onnxruntime.InferenceSession(onnx_path, sess_options)
    # session = onnxruntime.InferenceSession(onnx_path)
    # print("The model expects input shape: ", session.get_inputs()[0].shape)
    IN_IMAGE_H = session.get_inputs()[0].shape[2]
    IN_IMAGE_W = session.get_inputs()[0].shape[3]
    input_name = session.get_inputs()[0].name

    print(IN_IMAGE_H)

    cap = cv2.VideoCapture("C:\\Users\\VINNO\\Desktop\\2.mp4")
    t1 = time.time()
    t = 0
    forwardtime = 0
    alltime = 0
    forward_progresstime = 0
    while (cap.isOpened()):
        ret, frame = cap.read()
        if ret == True:

            ttt_alltime = time.time()
            # Input
            resized = cv2.resize(frame, (IN_IMAGE_W, IN_IMAGE_H), interpolation=cv2.INTER_LINEAR)

            # img_in = cv2.cvtColor(resized, cv2.COLOR_BGR2GRAY)
            # img_in = np.expand_dims(img_in, axis=2)
            img_in = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)

            img_in = np.transpose(img_in, (2, 0, 1)).astype(np.float32)
            img_in = np.expand_dims(img_in, axis=0)
            img_in /= 255.0

            ttt_forwardtime = time.time()
            outputs = session.run(None, {input_name: img_in})
            temp_forwardtime = (time.time() - ttt_forwardtime) * 1000
            forwardtime += temp_forwardtime

            boxes = post_processing(img_in, 0.4, 0.6, outputs)

            temp_forward_progresstime = (time.time() - ttt_forwardtime) * 1000
            forward_progresstime += temp_forward_progresstime

            namesfile = 'breast.names'
            class_names = load_class_names(namesfile)
            dst_img = plot_boxes_cv2(frame, boxes[0], class_names=class_names)

            temp_alltime = (time.time() - ttt_alltime) * 1000
            alltime += temp_alltime
            print(temp_alltime)

            t += 1
            cv2.imshow('111', frame)

            if cv2.waitKey(1) & 0xff == ord('q'):
                break
        else:
            break
    tttt = time.time() - t1
    print("时间:{}".format(tttt))
    print(t)
    print("alltime平均时间：{}".format(alltime/t))
    print("forwardtime平均时间：{}".format(forwardtime/t))
    print("forward_progresstime平均时间：{}".format(forward_progresstime/t))


    print('{:.3f}'.format(tttt/t*1000))
    cap.release()
    cv2.destroyAllWindows()



if __name__ == '__main__':
    print("Converting to onnx and running demo ...")

    # cfg_file = 'C:\\Users\\VINNO\\Desktop\\新建文件夹 (2)\\pytorch-YOLOv4-master\\yolov4-tiny-breast-256-anchors---20210820.cfg'
    # weight_file = 'C:\\Users\\VINNO\\Desktop\\新建文件夹 (2)\\pytorch-YOLOv4-master\\yolov4-tiny-breast-256-anchors---20210820_last.weights'
    # batch_size = 1
    #
    # if batch_size <= 0:
    #     onnx_path_demo = transform_to_onnx(cfg_file, weight_file, batch_size)
    # else:
    #     # Transform to onnx as specified batch size
    #     transform_to_onnx(cfg_file, weight_file, batch_size)
    #     # Transform to onnx as demo
    #     onnx_path_demo = transform_to_onnx(cfg_file, weight_file, 1)
    # session = onnxruntime.InferenceSession(onnx_path_demo)
    #
    #
    # cpuCount = os.cpu_count()
    # print("Number of CPUs in the system:", cpuCount)

    main()