graduation-design/img_process/ObjectClassify.py
GitLab dae25cdf47 1. 调试了模型预测的输出结果
2. 成功控制机械臂移动
2024-03-31 22:26:56 +08:00

107 lines
3.3 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import asyncio
import os
from time import sleep
import Robot
import torch
from torch import tensor
from ultralytics import YOLO
import cv2
from orbbec_camera.OrbbecCamera import OrbbecCamera
from robot_control import FRRobot
# 全局变量
model = YOLO('yolov8n.pt')
image_width = 640
camera = OrbbecCamera('HW', True, image_width=image_width)
center_point = 320
robot = Robot.RPC('192.168.3.102')
def init():
global image_width
# 检测GPU是否存在
if not torch.cuda.is_available():
print("你忘了打开独显,大聪明")
exit(1)
# 输入的视频的宽度
image_width = 640
camera.run()
if __name__ == '__main__':
# 初始化
init()
# 处理循环
while True:
color_image = camera.get_color_image()
if color_image is not None:
# conf: 置信度,小于该值的框将被过滤
# imgsz: 图像的尺寸
# half: 是否使用FP16,可有效提升速度
# vid_stride: 帧预测间隔可降低算力消耗间隔x帧进行一次预测
results = model.predict(source=color_image,
show=True,
conf=0.5,
half=True,
imgsz=image_width,
verbose=False,
vid_stride=5)
print("中心点的深度值: ", camera.get_center_distance(), "mm")
x_pos = 0
y_pos = 0
high = 0
width = 0
for r in results:
if len(r.boxes.xyxy) > 0:
positions = r.boxes.xywh.tolist()
x_pos = positions[0][0]
y_pos = positions[0][1]
high = positions[0][2]
width = positions[0][3]
# print("X中心点:", x_pos)
# print("Y中心点:", y_pos)
# print("宽度", high)
# print("高度", width)
# 计算与中心点的距离
x_distance = x_pos - center_point
y_distance = y_pos - center_point
# print("X距离:", x_distance)
# print("Y距离:", y_distance)
# 移动方向
# dir0 - 负方向1 - 正方向;
x_dir = 0
y_dir = 0
# 移动时间
x_time = 0
y_time = 0
# 确定移动方向
if abs(x_distance) > 25:
x_time = abs(x_distance) / 21
if x_distance > 0:
x_dir = 1
else:
x_dir = 0
robot.StartJOG(ref=4, nb=1, dir=x_dir, max_dis=30, vel=20.0, acc=100.0)
sleep(0.5)
robot.ImmStopJOG()
if abs(y_distance) > 25:
y_time = abs(y_distance) / 21
if y_distance < 0:
y_dir = 1
else:
y_dir = 0
robot.StartJOG(ref=4, nb=2, dir=y_dir, max_dis=30, vel=20.0, acc=100.0)
sleep(0.5)
robot.ImmStopJOG()
key = cv2.waitKey(1)
if key == ord('q'):
camera.stop()
# cw.stop()
break