|
本帖最后由 zx896825328 于 2022-11-15 20:32 编辑
请求大佬给个例子感谢!
from ctypes import *
import cv2
import numpy as np
import numpy.ctypeslib as npct
class Detector():
def __init__(self,model_path,dll_path):
self.yolov5 = CDLL(dll_path)
self.yolov5.Detect.argtypes = [c_void_p,c_int,c_int,POINTER(c_ubyte),npct.ndpointer(dtype = np.float32, ndim = 2, shape = (50, 6), flags="C_CONTIGUOUS")]
self.yolov5.Init.restype = c_void_p
self.yolov5.Init.argtypes = [c_void_p]
self.yolov5.cuda_free.argtypes = [c_void_p]
self.c_point = self.yolov5.Init(model_path)
def predict(self,img):
rows, cols = img.shape[0], img.shape[1]
res_arr = np.zeros((50,6),dtype=np.float32)
self.yolov5.Detect(self.c_point,c_int(rows), c_int(cols), img.ctypes.data_as(POINTER(c_ubyte)),res_arr)
self.bbox_array = res_arr[~(res_arr==0).all(1)]
return self.bbox_array
def free(self):
self.yolov5.cuda_free(self.c_point)
def visualize(img,bbox_array):
for temp in bbox_array:
bbox = [temp[0],temp[1],temp[2],temp[3]] #xywh
clas = int(temp[4])
score = temp[5]
cv2.rectangle(img,(int(temp[0]),int(temp[1])),(int(temp[0]+temp[2]),int(temp[1]+temp[3])), (105, 237, 249), 2)
img = cv2.putText(img, "class:"+str(clas)+" "+str(round(score,2)), (int(temp[0]),int(temp[1])-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (105, 237, 249), 1)
return img
det = Detector(model_path=b"./yolov5s.engine",dll_path="./yolov5.dll") # b'' is needed
img = cv2.imread("./pictures/zidane.jpg")
result = det.predict(img)
img = visualize(img,result)
cv2.imshow("img",img)
cv2.waitKey(0)
det.free()
cv2.destroyAllWindows()
|
|