quat_visualization_matplot.py

import argparse
import dataclasses
import json
import math
import threading
import time
from typing import List

import matplotlib
import matplotlib.animation as animation
import matplotlib.pyplot as plt
import numpy as np
import scipy.spatial.transform as transform
from pymavlink import mavutil
from scipy.spatial.transform import Rotation as R

from nav_messages.message_service import MessageService


class QuatVisualizer:

    def __init__(self):
        self.rot_vins = transform.Rotation.from_quat([0, 0, 0, 1]).as_matrix()
        self.rot_mav = transform.Rotation.from_quat([0, 0, 0, 1]).as_matrix()
        self.yaw_align_deg = 180

    def visualize(self):
        def init():
            ax.clear()
            ax.set_xlim(-1.3, 1.3)
            ax.set_ylim(-1.3, 1.3)
            ax.set_zlim(-1.3, 1.3)
            ax.set_xlabel('X')
            ax.set_ylabel('Y')
            ax.set_zlabel('Z')
            ax.set_title('3D Rotating Arrow')
            return []

        def animate(i):
            init()
            origin = np.array([0, 0, 0])

            x = np.array([1, 0, 0])
            y = np.array([0, 1, 0])
            z = np.array([0, 0, 1])

            align_yaw_matrix = R.from_euler('z', self.yaw_align_deg, degrees=True).as_matrix()

            ax.quiver(*origin, *(self.rot_mav @ x), color=(1, 0, 0, 0.5))
            ax.quiver(*origin, *(self.rot_mav @ y), color=(0, 1, 0, 0.5))
            ax.quiver(*origin, *(self.rot_mav @ z), color=(0, 0, 1, 0.5))

            ax.quiver(*origin, *(align_yaw_matrix @ self.rot_vins @ x * 0.5), color='r')
            ax.quiver(*origin, *(align_yaw_matrix @ self.rot_vins @ y * 0.5), color='g')
            ax.quiver(*origin, *(align_yaw_matrix @ self.rot_vins @ z * 0.5), color='b')
            return []

        fig = plt.figure()
        ax = fig.add_subplot(1, 1, 1, projection='3d')

        ani = animation.FuncAnimation(fig, animate, init_func=init, frames=99999999, interval=20, blit=True)

        plt.show()

rot_mav_in_mav_coords = transform.Rotation.from_euler('xyz', [roll, pitch, yaw], degrees=True).as_matrix()
align_yaw_matrix = R.from_euler('z', q_vis.yaw_align_deg, degrees=True).as_matrix()
q_vis = QuatVisualizer()
q_vis.rot_vins = rot_mav_in_mav_coords
if visualize:
    matplotlib.use('TkAgg')
    threading.Thread(target=lambda: q_vis.visualize(), daemon=True).start()