graduation-design/img_process/ObjectClassify.py

51 lines
1.2 KiB
Python

from ultralytics import YOLO
import cv2
from orbbec_camera.ColorViewer import ColorViewer
from orbbec_camera.DepthViewer import DepthViewer
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()
# 实例化深度图像操作类
dw = DepthViewer()
# 启动深度相机
dw.start()
# 处理循环
while True:
# frame = cw.get_current_frame()
frame = dw.get_current_image()
if frame is not None:
cv2.imshow("深度图像", frame)
print(dw.get_depth())
# if frame is not None:
# oc.model.predict(source=frame, show=True)
key = cv2.waitKey(1)
if key == ord('q'):
dw.stop()
# cw.stop()
break