Visualization.py

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Circle, Wedge
import matplotlib.cm as cm


def plot_state(gt, est, landmarks, particles=None, measurements=None, 
               show_particles=True, pause_time=0.2):

    plt.cla()
    
    plt.scatter(landmarks[:,0], landmarks[:,1], 
                c='black', marker='x', s=150, linewidths=2.5, 
                label='Landmarks', zorder=5)
    
    plt.plot(gt[:,0], gt[:,1], 'g-', linewidth=2, 
             label='Ground Truth', zorder=3)
    
    rx_gt, ry_gt, rtheta_gt = gt[-1]
    draw_robot(rx_gt, ry_gt, rtheta_gt, color='green', 
               size=0.4, alpha=0.8, zorder=6)
    
    if len(est) > 0:
        plt.plot(est[:,0], est[:,1], 'r--', linewidth=2, 
                 label='FastSLAM Est', zorder=4)
        
        rx_est, ry_est, rtheta_est = est[-1]
        draw_robot(rx_est, ry_est, rtheta_est, color='red', 
                   size=0.4, alpha=0.8, zorder=7)
    
    if show_particles and particles is not None:

        sorted_particles = sorted(particles, key=lambda p: p.weight)
        
        weights = np.array([p.weight for p in sorted_particles])
        
        if weights.max() > weights.min():
            weights_norm = (weights - weights.min()) / (weights.max() - weights.min())
        else:
            weights_norm = np.ones_like(weights)

        cmap = cm.get_cmap('YlOrRd')
        
        for i, particle in enumerate(sorted_particles):
            px, py, ptheta = particle.pose
            w_norm = weights_norm[i]
            
            color = cmap(w_norm)
            
            size = 0.15 + 0.15 * w_norm
            alpha = 0.3 + 0.4 * w_norm
            
            draw_robot(px, py, ptheta, color=color, 
                      size=size, alpha=alpha, zorder=2)
            
            if w_norm > 0.8 and len(particle.landmarks) > 0:
                for lm in particle.landmarks:
                    plt.scatter(lm.mu[0], lm.mu[1], 
                              c='orange', marker='o', s=30, 
                              alpha=0.2, zorder=1)
        
        n_particles = len(particles)
        weights_std = weights.std()
        
        info_text = (f"Particles: {n_particles}\n"
                    f"Weight std: {weights_std:.4f}\n"
                    f"Max weight: {weights.max():.4f}")
        
        plt.text(0.02, 0.98, info_text, 
                transform=plt.gca().transAxes,
                fontsize=10, verticalalignment='top',
                bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.7))
    
    if measurements is not None and len(measurements) > 0:
        rx, ry, rtheta = gt[-1]
        
        for z in measurements:
            if z[0] is None:
                continue
            
            dist = z[0]
            bearing = z[1]
            angle_world = rtheta + bearing
            
            lx = rx + dist * np.cos(angle_world)
            ly = ry + dist * np.sin(angle_world)
            
            plt.plot([rx, lx], [ry, ly], 'b-', alpha=0.4, linewidth=1.5, zorder=1)
            plt.scatter(lx, ly, c='blue', marker='o', s=40, alpha=0.5, zorder=2)

    plt.legend(loc='upper right', fontsize=10, framealpha=0.9)
    plt.axis('equal')
    plt.grid(True, alpha=0.3)
    plt.xlabel('X (m)', fontsize=11)
    plt.ylabel('Y (m)', fontsize=11)
    plt.title('FastSLAM 2.0 - Todas as Partículas', fontsize=13, fontweight='bold')
    
    plt.pause(pause_time)


def draw_robot(x, y, theta, color='red', size=0.3, alpha=1.0, zorder=5):

    front = size
    back = -size * 0.5
    side = size * 0.6
    
    vertices_local = np.array([
        [front, 0],      
        [back, side],    
        [back, -side]    
    ])
    
    cos_theta = np.cos(theta)
    sin_theta = np.sin(theta)
    rotation = np.array([
        [cos_theta, -sin_theta],
        [sin_theta, cos_theta]
    ])
    
    vertices_world = (rotation @ vertices_local.T).T
    vertices_world[:, 0] += x
    vertices_world[:, 1] += y
 
    triangle = plt.Polygon(vertices_world, 
                          facecolor=color, 
                          edgecolor='black' if alpha > 0.5 else color,
                          linewidth=1.5 if alpha > 0.5 else 0.5,
                          alpha=alpha, 
                          zorder=zorder)
    plt.gca().add_patch(triangle)


def plot_resample_effect(particles_before, particles_after, pause_time=0.5):

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))

    ax1.set_title('ANTES da Reamostragem', fontsize=14, fontweight='bold')
    
    weights_before = np.array([p.weight for p in particles_before])
    poses_before = np.array([p.pose[:2] for p in particles_before])
    
    scatter1 = ax1.scatter(poses_before[:,0], poses_before[:,1],
                          c=weights_before, cmap='YlOrRd',
                          s=50, alpha=0.7, edgecolors='black', linewidths=0.5)
    plt.colorbar(scatter1, ax=ax1, label='Weight')
    
    ax1.set_xlabel('X (m)')
    ax1.set_ylabel('Y (m)')
    ax1.grid(True, alpha=0.3)
    ax1.axis('equal')

    info1 = f"N={len(particles_before)}\nWeight std={weights_before.std():.4f}"
    ax1.text(0.02, 0.98, info1, transform=ax1.transAxes,
            fontsize=11, verticalalignment='top',
            bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))

    ax2.set_title('DEPOIS da Reamostragem', fontsize=14, fontweight='bold', color='green')
    
    weights_after = np.array([p.weight for p in particles_after])
    poses_after = np.array([p.pose[:2] for p in particles_after])
    
    scatter2 = ax2.scatter(poses_after[:,0], poses_after[:,1],
                          c=weights_after, cmap='YlOrRd',
                          s=50, alpha=0.7, edgecolors='black', linewidths=0.5)
    plt.colorbar(scatter2, ax=ax2, label='Weight')
    
    ax2.set_xlabel('X (m)')
    ax2.set_ylabel('Y (m)')
    ax2.grid(True, alpha=0.3)
    ax2.axis('equal')

    info2 = f"N={len(particles_after)}\nWeight std={weights_after.std():.4f}"
    ax2.text(0.02, 0.98, info2, transform=ax2.transAxes,
            fontsize=11, verticalalignment='top',
            bbox=dict(boxstyle='round', facecolor='lightgreen', alpha=0.8))
    
    plt.tight_layout()
    plt.pause(pause_time)
    plt.close()