Visualization.py

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Ellipse

def draw_ellipse(ax, mu, sigma, color='r'):
    cov = sigma[:2, :2]
    vals, vecs = np.linalg.eigh(cov)
    order = vals.argsort()[::-1]
    vals, vecs = vals[order], vecs[:, order]
    theta = np.degrees(np.atan2(*vecs[:, 0][::-1]))
    w, h = 2 * np.sqrt(np.maximum(vals, 0))
    ell = Ellipse(xy=mu, width=w, height=h, angle=theta, 
                  edgecolor=color, fc='none', lw=1, alpha=0.6)
    ax.add_patch(ell)

def plot_state(gt, est, landmarks, ekf_mu, ekf_Sigma, map_landmarks, 
               associations=None, measurements=None, pause_time=0.01):
    ax = plt.gca()
    plt.cla()
    
    plt.scatter(landmarks[:, 0], landmarks[:, 1], c='k', marker='x', label='True Landmarks')
    plt.plot(gt[:, 0], gt[:, 1], 'g-', label="Ground Truth")
    plt.plot(est[:, 0], est[:, 1], 'r--', label="EKF Estimate")

    draw_ellipse(ax, ekf_mu[:2], ekf_Sigma[:2, :2], color='r')

    for l_id, idx in map_landmarks.items():
        l_mu = ekf_mu[idx:idx+2]
        l_sigma = ekf_Sigma[idx:idx+2, idx:idx+2]
        plt.scatter(l_mu[0], l_mu[1], c='b', marker='o', s=10)
        draw_ellipse(ax, l_mu, l_sigma, color='b')

    if associations is not None:
        robot_est = ekf_mu[:2]
        for assoc in associations:
            landmark_id = assoc[0]
            is_candidate = assoc[1]
            
            if not is_candidate and landmark_id in map_landmarks:
                idx = map_landmarks[landmark_id]
                l_est = ekf_mu[idx:idx+2]
                plt.plot([robot_est[0], l_est[0]], [robot_est[1], l_est[1]], 
                         'y-', lw=1.5, alpha=0.8, label='Association' if 'Association' not in plt.gca().get_legend_handles_labels()[1] else "")

    if measurements is not None:
        rx, ry, r_theta = gt[-1] 
        for z in measurements:
            if z[0] is not None:
                angle_world = r_theta + z[1]
                lx = rx + z[0] * np.cos(angle_world)
                ly = ry + z[0] * np.sin(angle_world)
                plt.plot([rx, lx], [ry, ly], 'b-', alpha=0.1)

    plt.title(f"EKF SLAM: {len(map_landmarks)} Landmarks | Step: {len(est)}")
    plt.legend(loc='upper right', fontsize='x-small')
    plt.axis("equal")
    plt.grid(True, linestyle=':', alpha=0.6)
    plt.pause(pause_time)