(资料图片)

烧录完成以后插到Jetson Orin Nano开发板的风扇下方的卡槽中，图示如下：

说明一切准备工作就绪了。

说明一下，安装过程中要求输入提示的都输入 y

这个时间大概在五分钟左右，需要等一下才可以转换好。

1importtensorrtastrt  2fromtorchvisionimporttransforms  3importtorchast  4fromcollectionsimportOrderedDict,namedtuple  5importcv2ascv  6importtime  7importnumpyasnp  8  9img_transform=transforms.Compose([transforms.ToTensor(), 10transforms.Resize((640,640)) 11]) 12 13defload_classes(): 14withopen("classes.txt","r")asf: 15class_list=[cname.strip()forcnameinf.readlines()] 16returnclass_list 17 18 19defformat_yolov8(frame): 20row,col,_=frame.shape 21_max=max(col,row) 22result=np.zeros((_max,_max,3),np.uint8) 23result[0:row,0:col]=frame 24result=cv.cvtColor(result,cv.COLOR_BGR2RGB) 25returnresult 26 27defwrap_detection(input_image,output_data): 28class_ids=[] 29confidences=[] 30boxes=[] 31out_data=output_data.T 32rows=out_data.shape[0] 33 34image_width,image_height,_=input_image.shape 35 36x_factor=image_width/640.0 37y_factor=image_height/640.0 38 39forrinrange(rows): 40row=out_data[r] 41classes_scores=row[4:] 42class_id=np.argmax(classes_scores) 43if(classes_scores[class_id]>.25): 44class_ids.append(class_id) 45confidences.append(classes_scores[class_id]) 46x,y,w,h=row[0].item(),row[1].item(),row[2].item(),row[3].item() 47left=int((x-0.5*w)*x_factor) 48top=int((y-0.5*h)*y_factor) 49width=int(w*x_factor) 50height=int(h*y_factor) 51box=np.array([left,top,width,height]) 52boxes.append(box) 53 54indexes=cv.dnn.NMSBoxes(boxes,confidences,0.25,0.25) 55 56result_class_ids=[] 57result_confidences=[] 58result_boxes=[] 59 60foriinindexes: 61result_confidences.append(confidences[i]) 62result_class_ids.append(class_ids[i]) 63result_boxes.append(boxes[i]) 64 65returnresult_class_ids,result_confidences,result_boxes 66defgpu_trt_demo(): 67class_list=load_classes() 68device=t.device("cuda:0") 69Binding=namedtuple("Binding",("name","dtype","shape","data","ptr")) 70logger=trt.Logger(trt.Logger.INFO) 71withopen("yolov8n.engine","rb")asf,trt.Runtime(logger)asruntime: 72model=runtime.deserialize_cuda_engine(f.read()) 73bindings=OrderedDict() 74forindexinrange(model.num_bindings): 75name=model.get_binding_name(index) 76dtype=trt.nptype(model.get_binding_dtype(index)) 77shape=model.get_binding_shape(index) 78data=t.from_numpy(np.empty(shape,dtype=np.dtype(dtype))).to(device) 79bindings[name]=Binding(name,dtype,shape,data,int(data.data_ptr())) 80binding_addrs=OrderedDict((n,d.ptr)forn,dinbindings.items()) 81context=model.create_execution_context() 82 83capture=cv.VideoCapture("test.mp4") 84colors=[(255,255,0),(0,255,0),(0,255,255),(255,0,0)] 85whileTrue: 86_,frame=capture.read() 87ifframeisNone: 88print("Endofstream") 89break 90fh,fw,fc=frame.shape 91start=time.time() 92image=format_yolov8(frame) 93x_input=img_transform(image).view(1,3,640,640).to(device) 94binding_addrs["images"]=int(x_input.data_ptr()) 95context.execute_v2(list(binding_addrs.values())) 96out_prob=bindings["output0"].data.cpu().numpy() 97end=time.time() 98 99class_ids,confidences,boxes=wrap_detection(image,np.squeeze(out_prob,0))100for(classid,confidence,box)inzip(class_ids,confidences,boxes):101ifbox[2]>fw*0.67:102continue103color=colors[int(classid)%len(colors)]104cv.rectangle(frame,box,color,2)105cv.rectangle(frame,(box[0],box[1]-20),(box[0]+box[2],box[1]),color,-1)106cv.putText(frame,class_list[classid]+""+("%.2f"%confidence),(box[0],box[1]-10),cv.FONT_HERSHEY_SIMPLEX,.5,(0,0,0))107108inf_end=end-start109fps=1/inf_end110fps_label="FPS:%.2f"%fps111cv.putText(frame,fps_label,(10,25),cv.FONT_HERSHEY_SIMPLEX,1,(0,0,255),2)112cv.imshow("YOLOv8+TensorRT8.5.xObjectDetection",frame)113cc=cv.waitKey(1)114ifcc==27:115break116cv.waitKey(0)117cv.destroyAllWindows()118119120if__name__=="__main__":121gpu_trt_demo()

扫码查看OpenCV+OpenVIO+Pytorch系统化学习路线图

标签：