人物跟随.ipynb

import cv2
import imutils
import math
from picamera2 import Picamera2
from IPython.display import display, Image
import ipywidgets as widgets
import threading
import mediapipe as mp
from base_ctrl import BaseController
import time
import queue
import random

# 创建一个“停止”按钮，用户可以通过点击它来停止视频流
stopButton = widgets.ToggleButton(
    value=False,
    description='Stop',
    disabled=False,
    button_style='danger',
    tooltip='Description',
    icon='square'
)

# 初始化 MediaPipe 的绘图工具和姿态检测模型
mpDraw = mp.solutions.drawing_utils
mp_pose = mp.solutions.pose
pose = mp_pose.Pose(static_image_mode=False, 
                    model_complexity=1, 
                    smooth_landmarks=True, 
                    min_detection_confidence=0.5, 
                    min_tracking_confidence=0.5)

# 用于检测树莓派的函数
def is_raspberry_pi5():
    with open('/proc/cpuinfo', 'r') as file:
        for line in file:
            if 'Model' in line:
                if 'Raspberry Pi 5' in line:
                    return True
                else:
                    return False

# 根据树莓派的型号来确定 GPIO 串口设备名称
if is_raspberry_pi5():
    base = BaseController('/dev/ttyAMA0', 115200)
else:
    base = BaseController('/dev/serial0', 115200)

# 控制机器人运动的函数
def move_robot(direction, rotation=0):
    if direction == "forward":
        base.send_command({"T": 1, "L": 0.2, "R": 0.2})  # 前进
    elif direction == "backward":
        base.send_command({"T": 1, "L": -0.2, "R": -0.2})  # 后退
    elif direction == "stop":
        base.send_command({"T": 1, "L": 0, "R": 0})  # 停止
    elif rotation > 0:
        base.send_command({"T": 1, "L": 0.2, "R": -0.2})  # 右转
    elif rotation < 0:
        base.send_command({"T": 1, "L": -0.2, "R": 0.2})  # 左转

# 计算人体中心位置的函数
def get_body_center(landmarks):
    left_shoulder = landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER]
    right_shoulder = landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER]
    left_hip = landmarks[mp_pose.PoseLandmark.LEFT_HIP]
    right_hip = landmarks[mp_pose.PoseLandmark.RIGHT_HIP]
    
    x_center = (left_shoulder.x + right_shoulder.x + left_hip.x + right_hip.x) / 4
    y_center = (left_shoulder.y + right_shoulder.y + left_hip.y + right_hip.y) / 4
    
    return x_center, y_center

# 计算两点之间的距离（像素距离）
def calculate_distance(center_x, center_y, target_x, target_y):
    return math.sqrt((center_x - target_x)**2 + (center_y - target_y)**2)

# 云台回到原点的函数
def reset_gimbal():
    base.gimbal_ctrl(0, 0, 0, 0)  # 假设原点位置是 (0, 0)

# 定义显示函数，用于处理视频帧并进行姿态检测
def view(button):
    # 使用 Picamera2 来获取树莓派摄像头的视频流
    picam2 = Picamera2()  
    picam2.configure(picam2.create_preview_configuration(main={"size": (320, 240)}))
    picam2.start()

    display_handle = display(None, display_id=True)  # 创建显示句柄用于更新显示的图像
    
    previous_center = None  # 用于跟踪前一帧人体的中心位置
    target_distance_min = 0.5  # 目标最小距离
    target_distance_max = 0.7  # 目标最大距离
    buffer_zone = 0.05  # 设置误差范围，避免过于频繁的移动

    # 云台调整的平滑过渡量
    smoothing_factor = 0.1  # 调整每一帧对角度变化的响应程度，0.1 是适中的速度

    # 设置偏离中心的阈值，只有当人物的偏离超过这些阈值时，才调整云台
    x_threshold = 0.2  # 当人物偏离图像中心 20% 时，开始调整
    y_threshold = 0.2  # 同样对于 y 坐标，偏离 20% 时开始调整

    # 定义一个最大查找时间，防止无限循环
    max_find_time = 10  # 最大查找时间为10秒
    find_time = 0
    
    # 用于缓存帧的队列
    frame_queue = queue.Queue(maxsize=2)

    # 捕获视频的线程
    def capture_video():
        nonlocal find_time
        while True:
            frame = picam2.capture_array()  # 获取摄像头捕获的帧
            if not frame_queue.full():
                frame_queue.put(frame)
            else:
                frame_queue.get()
                frame_queue.put(frame)
            time.sleep(0.05)  # 降低捕获帧频率，避免过快

    # 启动视频捕获线程
    capture_thread = threading.Thread(target=capture_video)
    capture_thread.daemon = True
    capture_thread.start()

    # 处理视频帧
    while True:
        if not frame_queue.empty():
            frame = frame_queue.get()

            results = pose.process(frame)  # 使用 MediaPipe 处理图像，获取姿态检测结果

            # 如果检测到姿态关键点
            if results.pose_landmarks:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  # 确保色彩空间转换为 RGB
                mpDraw.draw_landmarks(frame, results.pose_landmarks, mp_pose.POSE_CONNECTIONS)
                
                # 获取当前人体的中心位置
                center_x, center_y = get_body_center(results.pose_landmarks.landmark)

                # 计算当前与目标的距离
                distance = calculate_distance(center_x, center_y, 0.5, 0.5)
                print(f"Center (x, y): ({center_x}, {center_y})")
                print(f"Distance: {distance} Target Distance: {target_distance_min}-{target_distance_max}")

                # 控制机器人保持在目标距离范围内
                if distance < target_distance_min - buffer_zone:
                    print("Moving forward to maintain distance")
                    move_robot("forward")
                elif distance > target_distance_max + buffer_zone:
                    print("Moving backward to maintain distance")
                    move_robot("backward")
                else:
                    print("Stopping to maintain distance")
                    move_robot("stop")
                    # 停顿 5 到 7 秒
                    time.sleep(random.uniform(5, 7))

                # 云台控制：如果人物偏离中心，进行调整
                input_x, input_y = 0, 0  # 默认不进行云台调整

                if abs(center_x - 0.5) > x_threshold:
                    input_x = (center_x - 0.5) * 90
                if abs(center_y - 0.5) > y_threshold:
                    input_y = (center_y - 0.5) * 90

                # 限制云台转动的角度范围
                input_x = max(min(input_x, 45), -45)
                input_y = max(min(input_y, 45), -45)

                # 平滑过渡
                if previous_center is not None:
                    input_x = previous_center[0] + smoothing_factor * (input_x - previous_center[0])
                    input_y = previous_center[1] + smoothing_factor * (input_y - previous_center[1])

                base.gimbal_ctrl(input_x, input_y, 0, 0)  # 发送云台控制指令

                previous_center = (input_x, input_y)

                # 如果检测到人体，重置查找时间
                find_time = 0

            else:
                # 如果没检测到人体，开始寻找
                find_time += 1
                print(f"Searching... Time: {find_time}s")
                if find_time >= max_find_time:
                    print("Max search time reached, stopping.")
                    break

# 显示“停止”按钮并启动显示函数的线程
display(stopButton)
thread = threading.Thread(target=view, args=(stopButton,))
thread.start()  # 启动线程