40 lines
935 B
Python
40 lines
935 B
Python
from ultralytics import YOLO
|
|
import cv2
|
|
from orbbec_camera.ColorViewer import ColorViewer
|
|
|
|
|
|
class ObjectClassify:
|
|
def __init__(self):
|
|
self.model = YOLO('yolov8n.pt')
|
|
|
|
def image_classify(self, path=None):
|
|
res = self.model(path)
|
|
|
|
predicted_img = res[0].plot()
|
|
|
|
cv2.imshow("分类结果", predicted_img)
|
|
cv2.waitKey(0)
|
|
cv2.destroyAllWindows()
|
|
|
|
def video_classify(self):
|
|
cw = ColorViewer()
|
|
frame = cw.get_current_frame()
|
|
self.model(source=frame, show=True)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
oc = ObjectClassify()
|
|
# 实例化相机操作类
|
|
cw = ColorViewer()
|
|
# 启动相机
|
|
cw.start()
|
|
# 处理循环
|
|
while True:
|
|
frame = cw.get_current_frame()
|
|
if frame is not None:
|
|
oc.model.predict(source=frame, show=True)
|
|
key = cv2.waitKey(1)
|
|
if key == ord('q'):
|
|
cw.stop()
|
|
break
|