數(shù)據(jù)集制作

我用手機(jī)拍了一張圖像

然后自己寫個代碼，每旋轉(zhuǎn)一度保存一張圖像，這樣就成功生成了360張圖像及其注釋文件，分為訓(xùn)練集與驗(yàn)證集。訓(xùn)練文件夾包含 320張帶有注釋的圖像。測試和驗(yàn)證文件夾都包含 40 張帶有注釋的圖像。數(shù)據(jù)集部分圖像顯示如下：

模型訓(xùn)練

準(zhǔn)備好數(shù)據(jù)集以后，直接按下面的命令行運(yùn)行即可：

yolo obb train data=pen_dataset.yaml model=yolov8s-obb.pt epochs=25 imgsz=640

導(dǎo)出與測試

模型導(dǎo)出與測試

# export model
yolo export model=yolov8s-obb.pt format=onnx
# inference model
yoloobbpredictmodel=pen_best.ptsource=pen_rotate_test.png

部署推理

轉(zhuǎn)成ONNX格式文件以后，基于OpenVINO-Python部署推理，相關(guān)代碼如下

class_list=["pen"]
colors=[(255,255,0),(0,255,0),(0,255,255),(255,0,0)]

ie=Core()
fordeviceinie.available_devices:
print(device)

#ReadIR
model=ie.read_model(model="pen_best.onnx")
compiled_model=ie.compile_model(model=model,device_name="CPU")
output_layer=compiled_model.output(0)

##xywhr
frame=cv.imread("D:/python/my_yolov8_train_demo/four_pen.jpg")
bgr=format_yolov8(frame)
img_h,img_w,img_c=bgr.shape

start=time.time()
image=cv.dnn.blobFromImage(bgr,1/255.0,(640,640),swapRB=True,crop=False)

res=compiled_model([image])[output_layer]#1x25x8400
rows=np.squeeze(res,0).T
boxes,confidences,angles,class_ids=post_process(rows)

indexes=cv.dnn.NMSBoxes(boxes,confidences,0.25,0.45)
M=np.zeros((2,3),dtype=np.float32)
forindexinindexes:
box=boxes[index]
d1=-angles[index]
color=colors[int(class_ids[index])%len(colors)]
pts=[(box[0],box[1]),(box[0]+box[2],box[1]),(box[0]+box[2],box[1]+box[3]),(box[0],box[1]+box[3])]
rrt_pts=get_rotate_point(pts,M,d1,box)
cv.drawContours(frame,[np.asarray(rrt_pts).astype(np.int32)],0,(255,0,255),2)
cv.putText(frame,class_list[class_ids[index]],(int(box[0]+box[2]/2),int(box[1]+box[3]/2)),cv.FONT_HERSHEY_SIMPLEX,1.0,(0,0,255),2)

end=time.time()
inf_end=end-start
fps=1/inf_end
fps_label="FPS:%.2f"%fps
cv.putText(frame,fps_label,(20,45),cv.FONT_HERSHEY_SIMPLEX,1,(0,0,255),2)

cv.imshow("YOLOv8-OBBRotateObjectDetection",frame)
cv.imwrite("D:/pen_result.jpg",frame)
cv.waitKey(0)
cv.destroyAllWindows()

審核編輯：劉清