Course/coursework/test.py - master - tracie/Sirius_projects

Course/coursework/test.py
160 строк6 KB
Ksenia Vasileva
done
13 окт 2025, 07:51
13 окт 2025, 07:517ae8d58
import cv2
import cv2.aruco as aruco
import socket
import egm_pb2
import time
import numpy as np
from scipy.spatial.transform import Rotation as R, Slerp
from scipy.interpolate import CubicSpline

def euler_angles_to_quaternion(roll, pitch, yaw):
    cy, sy, cp = np.cos(yaw * 0.5), np.sin(yaw * 0.5), np.cos(pitch * 0.5)
    sp, cr, sr = np.sin(pitch * 0.5), np.cos(roll * 0.5), np.sin(roll * 0.5)
    w = cr * cp * cy + sr * sp * sy
    x = sr * cp * cy - cr * sp * sy
    y = cr * sp * cy + sr * cp * sy
    z = cr * cp * sy - sr * sp * cy
    return w, x, y, z

def smooth_trajectory(points):
    if len(points) < 2:
        return points
    t = np.linspace(0, 1, len(points))
    cs = CubicSpline(t, points, bc_type='clamped')
    t_new = np.linspace(0, 1, len(points) * 2)
    smoothed_points = cs(t_new)
    return smoothed_points

def linear_interpolation(start_point, end_point):
    def distance(p1, p2):
        return np.linalg.norm(p1 - p2)

    start_point = np.array(start_point)
    end_point = np.array(end_point)

    start_coords, start_quat = start_point[:3], start_point[3:]
    end_coords, end_quat = end_point[:3], end_point[3:]

    total_distance = distance(start_coords, end_coords)
    max_distance = 0.2
    num_points = int(np.ceil(total_distance / max_distance)) + 1

    t_values = np.linspace(0, 1, num_points)
    t_values = np.sin(t_values * np.pi / 2)

    interpolated_coords = np.outer(1 - t_values, start_coords) + np.outer(t_values, end_coords)

    # SLERP для интерполяции кватернионов
    slerp = Slerp([0, 1], R.from_quat([start_quat, end_quat]))
    interpolated_quats = slerp(t_values).as_quat()

    # Объединение интерполированных координат и кватернионов
    interpolated_points = np.hstack((interpolated_coords, interpolated_quats))
    return interpolated_points

def create_sensor_message(point):
    x_p, y_p, z_p, q1, q2, q3, q4 = point

    sensor_message = egm_pb2.EgmSensor()

    # Создание и заполнение EgmCartesian
    cart = egm_pb2.EgmCartesian()
    cart.x = x_p
    cart.y = y_p
    cart.z = z_p

    # Создание и заполнение EgmQuaternion
    quat = egm_pb2.EgmQuaternion()
    quat.u0 = q1
    quat.u1 = q2
    quat.u2 = q3
    quat.u3 = q4

    pose = egm_pb2.EgmPose()
    pose.pos.CopyFrom(cart)
    pose.orient.CopyFrom(quat)

    # Создание и заполнение EgmPlanned
    planned = egm_pb2.EgmPlanned()
    planned.cartesian.CopyFrom(pose)
    # Заполнение EgmSensor
    sensor_message.planned.CopyFrom(planned)

    # Создание и заполнение EgmHeader
    header = egm_pb2.EgmHeader()
    header.mtype = egm_pb2.EgmHeader.MSGTYPE_CORRECTION
    header.seqno = 0
    header.tm = int(time.time())
    sensor_message.header.CopyFrom(header)

    return sensor_message

def motion_p2p(start, end):
    path = linear_interpolation(start[:3] + start[3:], end[:3] + end[3:])
    path = smooth_trajectory(path)
    for i in path:
        sensor_message = create_sensor_message(i)
        serialized_message = sensor_message.SerializeToString()
        udp_socket.sendto(serialized_message, ("192.168.1.10", 1025))
    print("Send")

if __name__ == "__main__":
    robtarget_start = [0, 0, 700, 0, 0, 0.707106785, -0.707106778]
    aruco_dict = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_6X6_250)
    parameters = aruco.DetectorParameters()

    mtx = np.array([[1.42711770e+03, 0.00000000e+00, 9.62248529e+02],
                    [0.00000000e+00, 1.45194135e+03, 5.55365484e+02],
                    [0.00000000e+00, 0.00000000e+00, 1.00000000e+00]])
    dist = np.array([0.02359394, 0.25472047, 0.0062465, -0.00323776, -0.53118869])
    cap = cv2.VideoCapture(0)

    ip_cr = "192.168.1.20"
    port_cr = 1025
    udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp_socket.bind((ip_cr, port_cr))

    first = True
    while True:
        ret, frame = cap.read()
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        corners, ids, rejected = aruco.detectMarkers(gray, aruco_dict, parameters=parameters)
        if len(corners) > 0:
            aruco.drawDetectedMarkers(frame, corners)
            for i in range(len(ids)):
                # Correct length marker
                rvec, tvec, _ = aruco.estimatePoseSingleMarkers(corners[i], 0.115, mtx, dist)
                x, y, z = tvec[0][0]
                x_mm, y_mm, z_mm = x * 1000, y * 1000, z * 1000
                z_mm -= 34
                quat_vec = np.array(euler_angles_to_quaternion(rvec[0][0][0], rvec[0][0][1], rvec[0][0][2]))
                point = [x_mm, y_mm, z_mm, quat_vec[0], quat_vec[1], quat_vec[2], quat_vec[3]]
                if first:
                    buf_point = robtarget_start
                    motion_p2p(buf_point, point)
                    first = False
                    buf_point = point
                    continue
                motion_p2p(buf_point, point)
                buf_point = point
                axis_points = np.float32([[0, 0, 0], [0, 0, 0.12], [-0.12, 0, 0], [0, -0.12, 0]]).reshape(-1, 3, 1)
                axis_points, _ = cv2.projectPoints(axis_points, rvec, tvec, mtx, dist)
                origin = tuple(axis_points[0].ravel().astype(int))
                z_axis_end = tuple(axis_points[1].ravel().astype(int))
                x_axis_end = tuple(axis_points[2].ravel().astype(int))
                y_axis_end = tuple(axis_points[3].ravel().astype(int))

                cv2.line(frame, origin, z_axis_end, (255, 0, 0), 2)
                cv2.line(frame, origin, x_axis_end, (0, 0, 255), 2)
                cv2.line(frame, origin, y_axis_end, (0, 255, 0), 2)

                cv2.putText(frame, f'x: {x_mm:.2f} mm, y: {y_mm:.2f} mm, z: {z_mm:.2f} mm',
                            (int(corners[i][0][0][0]), int(corners[i][0][0][1]) - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 4)
                cv2.imshow("CameraAruco", frame)
        cv2.imshow('CameraAruco', frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()