graduation-design/img_process/ObjectClassify.py

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