visualize_wireframe.py

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

def play_wireframe_matplotlib(npy_path, target_fps=12.0):
    landmarks = np.load(npy_path)
    num_frames = landmarks.shape[0]
    
    # Set up the plotting window
    fig, ax = plt.subplots(figsize=(6, 6))
    ax.set_xlim(-0.1, 1.1)
    ax.set_ylim(1.1, -0.1)  # Inverted Y-axis to match traditional image coordinates
    ax.set_aspect('equal')
    ax.axis('off')
    
    # Connection lines for the 28-point anime-face topology (matching
    # 0.4_test_inference.py and vts_playback_driver.py).
    #  0-4: face outline,  5-7: right brow,  8-10: left brow
    # 11-16: left eye, 17-22: right eye,  23: nose tip
    # 24-27: mouth
    connections = [
        # Face outline
        (0, 1), (1, 2), (2, 3), (3, 4),
        # Right brow (subject's right)
        (5, 6), (6, 7),
        # Left brow
        (8, 9), (9, 10),
        # Left eye contour (subject's left)
        (11, 12), (12, 13), (13, 16), (16, 15), (15, 14), (14, 11),
        # Right eye contour (subject's right)
        (17, 18), (18, 19), (19, 22), (22, 21), (21, 20), (20, 17),
        # Mouth loop
        (24, 25), (25, 26), (26, 27), (27, 24),
    ]
    
    # Initialize scatter plot and line structures
    scatter = ax.scatter([], [], c='#33cc33', s=15, zorder=3)
    lines = [ax.plot([], [], color='#4a90e2', lw=1.5, zorder=2)[0] for _ in connections]
    title_text = ax.set_title("Frame 0", fontsize=12)
    # One persistent text label per landmark index (0-27), moved each frame.
    n_pts = 28
    labels = [ax.text(0, 0, str(i), fontsize=7, color='red', zorder=4) for i in range(n_pts)]
    
    def update(frame_idx):
        lm = landmarks[frame_idx]
        
        if np.isnan(lm).any():
            # If tracking was lost, clear the face details
            scatter.set_offsets(np.empty((0, 2)))
            for line in lines:
                line.set_data([], [])
            for lbl in labels:
                lbl.set_position((-1, -1))
            title_text.set_text(f"Frame {frame_idx}/{num_frames} [No Face Tracked]")
        else:
            scatter.set_offsets(lm)
            for i, (start, end) in enumerate(connections):
                xs = [lm[start, 0], lm[end, 0]]
                ys = [lm[start, 1], lm[end, 1]]
                lines[i].set_data(xs, ys)
            for i in range(n_pts):
                labels[i].set_position((lm[i, 0], lm[i, 1]))
            title_text.set_text(f"Frame {frame_idx}/{num_frames}")
            
        return [scatter] + lines + labels

    # Playback interval matches 12 FPS (~83.3ms per frame)
    ani = FuncAnimation(
        fig, 
        update, 
        frames=num_frames, 
        interval=int(1000 / target_fps),
        blit=True,
        repeat=False
    )
    plt.show()

if __name__ == "__main__":
    play_wireframe_matplotlib("dataset\\landmarks\\XTRlxbOA8Pg_12fps_001153_to_end_norm_landmarks.npy")