100 lines
3.3 KiB
Python
100 lines
3.3 KiB
Python
import asyncio
|
||
import os
|
||
import threading
|
||
|
||
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.FRRobot import FRRobot
|
||
|
||
# 全局变量
|
||
model = YOLO('yolov8n.pt')
|
||
image_width = 640
|
||
# 矫正系数,用于映射像素尺寸和实际尺寸
|
||
# 开发时系数为25/32 = 0.78125
|
||
correction_factor = 1
|
||
# 创建摄像头实例
|
||
camera = OrbbecCamera('HW', True, image_width=image_width)
|
||
center_point = 320
|
||
# 创建FRRobot实例
|
||
robot_controller = FRRobot()
|
||
|
||
|
||
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, )
|
||
# print("中心点的深度值: ", camera.get_center_distance(), "mm")
|
||
x_pos = 0
|
||
y_pos = 0
|
||
high = 0
|
||
width = 0
|
||
for r in results:
|
||
if len(r.boxes.xywh) > 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) * correction_factor
|
||
y_distance = (y_pos - center_point) * correction_factor
|
||
print("X距离:", x_distance)
|
||
print("Y距离:", y_distance)
|
||
move_commands = []
|
||
if abs(x_distance) > 15:
|
||
x_time = abs(x_distance) / 21
|
||
x_dir = 1 if x_distance > 0 else 0
|
||
move_commands.append(("x", x_distance))
|
||
|
||
if abs(y_distance) > 15:
|
||
y_time = abs(y_distance) / 21
|
||
y_dir = 1 if y_distance < 0 else 0
|
||
move_commands.append(("y", y_distance))
|
||
|
||
if len(move_commands) > 0:
|
||
if not robot_controller.is_action:
|
||
# 添加移动指令
|
||
robot_controller.add_move_command(move_commands)
|
||
|
||
key = cv2.waitKey(1)
|
||
if key == ord('q'):
|
||
camera.stop()
|
||
robot_controller.add_move_command("stop")
|
||
break
|