/-- 1: install onnx on Ubuntu 20.04 pip3 install onnx==1.8.1 #或临时换国内源 pip3 install onnx==1.8.1 -i https://pypi.tuna.tsinghua.edu.cn/simple 2: install onnxruntime pip3 install onnxruntime -i https://pypi.tuna.tsinghua.edu.cn/simple 3: 应用 python3 run_onnx.py --onnx ./model/pix2pixHD_20230918.onnx --input ./images/demo.png ---/ import os import cv2 import onnx import onnxruntime import numpy as np import time import argparse \# import torch VIDEO_SUFFIX = [".avi", ".mp4"] IMAGE_SUFFIX = [".jpg", ".png"] WINDOW_NAME = "Lane Detection" def find_lanes(channel, image_vis, color): _, mask = cv2.threshold(channel, 100, 255, cv2.THRESH_BINARY) contours, _ = cv2.findContours(mask, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) contours = [contour for contour in contours if cv2.arcLength(contour, True) > 1000] for contour in contours: length = cv2.arcLength(contour, False) x_coors = contour[:, :, 0][:, 0] y_coors = contour[:, :, 1][:, 0] coeff = np.polyfit(y_coors, x_coors, 5) start = np.min(y_coors) end = np.max(y_coors) poly_y_coors = np.linspace(start, end, int((end - start))) poly_x_coors = np.polyval(coeff, poly_y_coors) poly_points = [[x, y] for [x, y] in zip(poly_x_coors, poly_y_coors)] cv2.polylines(image_vis, np.int32([poly_points]), isClosed=False, color=color, thickness=3) def infer_image(image_file, onnx_session, need_read_image=True): if need_read_image: image = cv2.imread(image_file) if image is None: print("Can't open image: {}".format(image_file)) return else: image = image_file src_img_size = image.shape src_width = src_img_size[1] src_height = src_img_size[0] print('src image size:', src_width, 'x', src_height) image_vis = image.copy() _, _, h, w = onnx_session.get_inputs()[0].shape image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) image = cv2.resize(image, (w, h)) #image.tofile("dataRGBPACKAGE.bin") # np.save("data.bin", image) #save rgb input data image = np.transpose(image, (2, 0, 1)) # whc->chw #image.tofile("dataRGBPlanner.bin") image =image.astype(np.float32) # image = image / 127.5 - 1.0 image = image/255 input_tensor = np.expand_dims(image, axis=0) print(input_tensor.shape) ## model inference start = time.time() output_names = [out.name for out in onnx_session.get_outputs()] inputs = {onnx_session.get_inputs()[0].name: input_tensor} outputs = onnx_session.run(output_names, inputs) end = time.time() print("Model inference cost time: %.4f ms." % ((end - start) * 1e3)) out_idx = 0 for outdata in outputs: print("out layer ---- >>> name shape") # print(outdata.name) print(output_names[out_idx]) print(outdata.shape) outdataname="./output_%d.bin"%(out_idx) outdata.tofile(outdataname) out_idx = out_idx+1 ## procee model output generated = np.squeeze(outputs[0]) #print(generated.shape) generated = np.transpose(generated, (1, 2, 0)) #print(generated.shape) generated = np.clip(generated, 0, 255).astype(np.uint8) #print(generated.shape, type(generated)) ret, generated = cv2.threshold(generated, 1, 255, cv2.THRESH_BINARY) cv2.imwrite('lanes.png', generated) ch = np.dsplit(generated, 3) ch1 = np.array(ch[0]) # 640x352 ch2 = np.array(ch[1]) ch3 = np.array(ch[2]) cv2.imwrite('ch1.png', ch1) cv2.imwrite('ch2.png', ch2) cv2.imwrite('ch3.png', ch3) CH1 = cv2.resize(ch1, (src_width, src_height)) # 1920x1080 CH2 = cv2.resize(ch2, (src_width, src_height)) # 1920x1080 CH3 = cv2.resize(ch3, (src_width, src_height)) # 1920x1080 cv2.imwrite('CH1.png', CH1) cv2.imwrite('CH2.png', CH2) cv2.imwrite('CH3.png', CH3) generated = cv2.cvtColor(generated, cv2.COLOR_BGR2RGB) h, w, c = image_vis.shape generated = cv2.resize(generated, (w, h)) mask = generated.copy() ### 远端的轮廓会粘连到一起，直接把这部分的轮廓去掉 mask[0 : h // 3, :, :] = 0 ## find contours and fit lines find_lanes(mask[:, :, 1], image_vis, color=(0, 255, 0)) find_lanes(mask[:, :, 2], image_vis, color=(0, 0, 255)) image_show = np.zeros([h // 2, w, c], dtype=image_vis.dtype) generated = cv2.resize(generated, (w // 2, h // 2)) image_vis = cv2.resize(image_vis, (w // 2, h // 2)) image_show[:, 0 : w // 2, :] = image_vis image_show[:, w // 2 : w, :] = generated cv2.imshow(WINDOW_NAME, image_show) if need_read_image: key = cv2.waitKey(0) & 0xFF else: key = cv2.waitKey(1) & 0xFF return key def infer_video(video_file, onnx_session): cap = cv2.VideoCapture(video_file) if not cap.isOpened(): print("Can't open video: {}".format(video_file)) return while True: ret, frame = cap.read() if not ret: break key = infer_image(frame, onnx_session, False) if key == 27: exit() elif key == 113: #'q' break def run_on_file(f, onnx_session): suffix = os.path.splitext(f)[-1] if suffix in VIDEO_SUFFIX: infer_video(f, onnx_session) elif suffix in IMAGE_SUFFIX: infer_image(f, onnx_session) else: print("This file is not a image(.png,.jpg) or a video(.mp4,.avi) file: {}".format(f)) def run_onnx(args): print("Loadding onnx model {}".format(args.onnx)) onnx_model=onnx.load(args.onnx) onnx_model_graph=onnx_model.graph nodes=onnx_model_graph.node for inode in nodes: print(inode.name) # print(onnx_model) # onnx_model.graph.output.extend([onnx.ValueInfoProto(name="128")]) print(onnx_model.graph.output) # onnx_model_graph session = onnxruntime.InferenceSession(onnx_model.SerializeToString(), providers=["CUDAExecutionProvider", "CPUExecutionProvider"]) print(session.get_inputs()[0].name) print(session.get_inputs()[0].shape) if os.path.isfile(args.input): run_on_file(args.input, session) elif os.path.isdir(args.input): for root, _, files in os.walk(args.input): for f in files: f = os.path.join(root, f) run_on_file(f, session) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( "--onnx", type=str, default="pix2pixHD.onnx", help="the onnx model", ) parser.add_argument( "--input", type=str, default="test.jpg", help="the path of image(s) or video(s), file or directory", ) args = parser.parse_args() cv2.namedWindow(WINDOW_NAME, cv2.WINDOW_NORMAL) cv2.resizeWindow(WINDOW_NAME, 1920, 1080) run_onnx(args)

标签: none

评论已关闭