Untitled

import cv2
import numpy as np
import sys
import os
import apriltag
import math
from math import cos, sin

robot_driving_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../lidar_processing/lidar_raw/autonomous_driving"))
sys.path.append(robot_driving_path)

from drive import RobotDriving

# Init robot
robot = RobotDriving()

# Initialize the detector (AprilTag 3)
detector = apriltag.Detector(None, searchpath=apriltag._get_dll_path())

# Open video capture (0 for default webcam, change if using another USB camera)
capture = cv2.VideoCapture(0)

if not capture.isOpened():
    print("Error: Could not open webcam.")
    sys.exit()

# Get camera parameters
fx, fy, cx, cy = 600, 600, 320, 240

# Bézier curve functions
def bezier_cubic(t, P0, P1, P2, P3):
    x = (1 - t)**3 * P0[0] + 3 * (1 - t)**2 * t * P1[0] + 3 * (1 - t) * t**2 * P2[0] + t**3 * P3[0]
    y = (1 - t)**3 * P0[1] + 3 * (1 - t)**2 * t * P1[1] + 3 * (1 - t) * t**2 * P2[1] + t**3 * P3[1]
    return x, y

def bezier_quad(t, P0, P1, P2):
    x = (1 - t)**2 * P0[0] + 2 * (1 - t) * t * P1[0] + t**2 * P2[0]
    y = (1 - t)**2 * P0[1] + 2 * (1 - t) * t * P1[1] + t**2 * P2[1]
    return x, y

# Function to calculate distance
def distance(p1, p2):
    return math.sqrt((p2[0] - p1[0]) ** 2 + (p2[1] - p1[1]) ** 2)

# Function to go to a target position following the cubic Bézier curve
def goto_position(start_position, target_position):
    # Compute trajectory along Bézier curve
    t = 0
    current_position = start_position
    while t <= 1.0:
        # Control points for Bézier curve
        P0 = current_position
        P2 = target_position

        if (robot.direction) != np.pi/2:
            d = (P2[0] - P0[0]) / cos(robot.direction)
            P1 = [P2[0], P0[1] + d * sin(robot.direction)]
        else:
            P1 = [P2[0], (P0[1] + P2[1]) /2]

        bx, by = bezier_quad(t, P0, P1, P2)

        # Calculate distance and forward speed
        d = distance(current_position, [bx, by])
        forward = min(d * 0.1, 0.1)  # Limit max speed to 0.5
        # Calculate angle to next point
        angle_to_checkpoint = math.atan2(by - current_position[1], bx - current_position[0])
        left = angle_to_checkpoint - robot.direction
        left = (left + np.pi) % (2 * np.pi) - np.pi

        # Move robot
        robot.move(forward, left)
        sleep(0.1)
        robot.update_position()
        current_position = robot.position
        t += 0.1  # Increment time parameter to move along the Bézier curve

    return current_position  # Return current position when the target is reached

# Main loop
detection_on = True
checkpoint_reached = False

while True:
    ret, frame = capture.read()
    if not ret:
        print("Error: Failed to capture image.")
        break

    if detection_on:
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        # Detect AprilTags
        results, overlay = apriltag.detect_tags(gray,
                                                detector,
                                                camera_params=(fx, fy, cx, cy),
                                                tag_size=0.08,
                                                vizualization=3,
                                                verbose=0,
                                                annotation=True)


        print(results)
        if results:
            for result in results:
                tag_id = result[0].tag_id

                if tag_id == 0 and not checkpoint_reached:
                    print("Tag 0 detected!")

                    # Extract pose matrix from the result
                    pose_matrix = result[1]
                    tag_position_cam = pose_matrix[:3, 3]  # Position of the tag in camera frame
                    tag_position_global = robot.position + tag_position_cam
                    x0, y0 = tag_position_global[:2]  # Get the x and y coordinates

                    # Compute checkpoint position (1m in front of the tag)
                    checkpoint_position = [x0, y0 - 1.0]  # Target 1 meter in front of the tag

                    print(f"Going to checkpoint at: {checkpoint_position}")

                    # Disable further detection and move to checkpoint
                    detection_on = False
                    current_position = [0, 0]
                    current_position = goto_position(current_position, checkpoint_position)

                    # Once checkpoint is reached, enable detection again
                    checkpoint_reached = True
                    detection_on = True
                    print("Checkpoint reached!")

                elif tag_id == 0 and checkpoint_reached:
                    print("Re-detecting tag and adjusting position.")

                    # Extract pose matrix for re-adjusting the position
                    pose_matrix = result[1]
                    tag_position = pose_matrix[:3, 3]
                    x0, y0 = tag_position[:2]

                    # Compute second target position (50cm in front of the tag)
                    target_position = [x0, y0 - 0.5]  # Target 50 cm in front of the tag

                    print(f"Going to target position: {target_position}")

                    # Move to the adjusted target position
                    current_position = goto_position(current_position, target_position)
                    print("Target position reached!")
                    detection_on = False  # Finish the loop, no further movement

                    break

    # Display the video feed
    cv2.imshow("AprilTag", overlay)

    if cv2.waitKey(1) == ord('q'):
        break

capture.release()
cv2.destroyAllWindows()