107 lines
3.3 KiB
Python
107 lines
3.3 KiB
Python
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)
|
||
# 移动方向
|
||
# dir:0 - 负方向,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
|