50 lines
1.4 KiB
Python
50 lines
1.4 KiB
Python
from ultralytics import YOLO
|
|
import cv2
|
|
from abandon_code.ColorViewer import ColorViewer
|
|
from orbbec_camera.OrbbecCamera import OrbbecCamera
|
|
|
|
|
|
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()
|
|
ob = OrbbecCamera('HW', True)
|
|
ob.run()
|
|
# 处理循环
|
|
while True:
|
|
color_image = ob.get_color_image()
|
|
if color_image is None:
|
|
continue
|
|
cv2.imshow("彩色图像", color_image)
|
|
if ob.get_depth_image() is None:
|
|
continue
|
|
cv2.imshow("深度图像", ob.get_depth_image())
|
|
if ob.get_combined_image() is None:
|
|
continue
|
|
cv2.imshow("合成图像", ob.get_combined_image())
|
|
print(f"中心点的深度值: {ob.get_center_distance()}mm")
|
|
if color_image is not None:
|
|
oc.model.predict(source=color_image, show=True)
|
|
key = cv2.waitKey(1)
|
|
if key == ord('q'):
|
|
ob.stop()
|
|
# cw.stop()
|
|
break
|