Add visualization.py

fsd_model/visualization.py

@@ -0,0 +1,163 @@
+"""
+Visualization utilities for FSD model outputs.
+Generates visual representations of:
+- Sensor placement on vehicle
+- BEV perception outputs
+- Planned trajectory
+- Control commands
+"""
+
+import torch
+import numpy as np
+from typing import Dict, Optional, Tuple, List
+import math
+
+
+def visualize_sensor_layout_ascii(vehicle_config) -> str:
+    """Generate ASCII art of sensor placement on vehicle."""
+    sc = vehicle_config.sensor_config
+    half_l = vehicle_config.length / 2
+    half_w = vehicle_config.width / 2
+    
+    # Create grid
+    grid_h, grid_w = 30, 50
+    grid = [[' '] * grid_w for _ in range(grid_h)]
+    
+    # Scale factors
+    scale_x = (grid_w - 10) / (vehicle_config.length + 2)
+    scale_y = (grid_h - 6) / (vehicle_config.width + 2)
+    
+    cx, cy = grid_w // 2, grid_h // 2
+    
+    # Draw vehicle outline
+    vw = int(vehicle_config.length * scale_x / 2)
+    vh = int(vehicle_config.width * scale_y / 2)
+    
+    for x in range(cx - vw, cx + vw + 1):
+        if 0 <= x < grid_w:
+            if 0 <= cy - vh < grid_h:
+                grid[cy - vh][x] = '─'
+            if 0 <= cy + vh < grid_h:
+                grid[cy + vh][x] = '─'
+    
+    for y in range(cy - vh, cy + vh + 1):
+        if 0 <= y < grid_h:
+            if 0 <= cx - vw < grid_w:
+                grid[y][cx - vw] = '│'
+            if 0 <= cx + vw < grid_w:
+                grid[y][cx + vw] = '│'
+    
+    # Corners
+    for (dy, dx), ch in [
+        ((-vh, -vw), '┌'), ((-vh, vw), '┐'),
+        ((vh, -vw), '└'), ((vh, vw), '┘')
+    ]:
+        gy, gx = cy + dy, cx + dx
+        if 0 <= gy < grid_h and 0 <= gx < grid_w:
+            grid[gy][gx] = ch
+    
+    # Direction arrow
+    if 0 <= cy < grid_h and cx + vw + 1 < grid_w:
+        grid[cy][cx + vw + 1] = '►'
+    grid[cy][cx] = '+'  # center
+    
+    # Place cameras
+    for i, cam in enumerate(sc.cameras):
+        gx = cx + int(cam.placement.x * scale_x)
+        gy = cy - int(cam.placement.y * scale_y)
+        if 0 <= gy < grid_h and 0 <= gx < grid_w:
+            grid[gy][gx] = 'C'
+    
+    # Place ultrasonics
+    for i, us in enumerate(sc.ultrasonics):
+        gx = cx + int(us.placement.x * scale_x)
+        gy = cy - int(us.placement.y * scale_y)
+        if 0 <= gy < grid_h and 0 <= gx < grid_w:
+            if grid[gy][gx] == ' ' or grid[gy][gx] in '─│':
+                grid[gy][gx] = 'U'
+    
+    # Add labels
+    result = "Vehicle Sensor Layout (Top View)\n"
+    result += "C = Camera, U = Ultrasonic, + = Center, ► = Forward\n"
+    result += "=" * grid_w + "\n"
+    for row in grid:
+        result += ''.join(row) + "\n"
+    result += "=" * grid_w + "\n"
+    
+    return result
+
+
+def format_model_output(output: Dict[str, torch.Tensor]) -> str:
+    """Format model output as human-readable text."""
+    lines = []
+    lines.append("╔══════════════════════════════════════════╗")
+    lines.append("║       FSD Model Output Summary           ║")
+    lines.append("╠══════════════════════════════════════════╣")
+    
+    # Control outputs
+    if "control/steering_deg" in output:
+        steer = output["control/steering_deg"].mean().item()
+        throttle = output["control/throttle"].mean().item()
+        brake = output["control/brake"].mean().item()
+        lines.append(f"║ Steering:  {steer:+.2f}°")
+        lines.append(f"║ Throttle:  {throttle:.3f}")
+        lines.append(f"║ Brake:     {brake:.3f}")
+    
+    # Safety
+    if "planning/collision_risk" in output:
+        risk = output["planning/collision_risk"].mean().item()
+        emergency = output["planning/emergency_brake"].mean().item()
+        lines.append(f"║ Collision Risk: {risk:.3f}")
+        lines.append(f"║ Emergency Brake: {emergency:.3f}")
+    
+    # Behavior
+    if "planning/behavior_logits" in output:
+        behaviors = [
+            "keep_lane", "turn_left", "turn_right", "change_left",
+            "change_right", "stop", "yield", "park", "reverse", "emergency"
+        ]
+        probs = torch.softmax(output["planning/behavior_logits"], dim=-1)
+        top_prob, top_idx = probs.mean(0).topk(3)
+        lines.append("║ Top behaviors:")
+        for p, i in zip(top_prob, top_idx):
+            lines.append(f"║   {behaviors[i.item()]}: {p.item():.3f}")
+    
+    # Waypoints
+    if "planning/safe_waypoints" in output:
+        wp = output["planning/safe_waypoints"]
+        lines.append(f"║ Planned waypoints: {wp.shape[1]}")
+        lines.append(f"║   First: ({wp[0,0,0]:.2f}, {wp[0,0,1]:.2f})")
+        lines.append(f"║   Last:  ({wp[0,-1,0]:.2f}, {wp[0,-1,1]:.2f})")
+    
+    # Controller weights
+    if "control/controller_weights" in output:
+        w = output["control/controller_weights"].mean(0)
+        lines.append(f"║ Controller mix: Neural={w[0]:.2f} Stanley={w[1]:.2f} PID={w[2]:.2f}")
+    
+    lines.append("╚══════════════════════════════════════════╝")
+    
+    return "\n".join(lines)
+
+
+def format_parameter_count(counts: Dict[str, int]) -> str:
+    """Format parameter counts as a nice table."""
+    lines = []
+    lines.append("┌─────────────────┬──────────────┬───────────┐")
+    lines.append("│ Module          │ Parameters   │ Size (MB) │")
+    lines.append("├─────────────────┼──────────────┼───────────┤")
+    
+    for name, count in counts.items():
+        if name in ["total", "total_trainable"]:
+            continue
+        size_mb = count * 4 / (1024 * 1024)  # float32
+        lines.append(f"│ {name:<15} │ {count:>12,} │ {size_mb:>8.2f}  │")
+    
+    lines.append("├─────────────────┼──────────────┼───────────┤")
+    total = counts.get("total", 0)
+    trainable = counts.get("total_trainable", 0)
+    total_mb = total * 4 / (1024 * 1024)
+    lines.append(f"│ {'TOTAL':<15} │ {total:>12,} │ {total_mb:>8.2f}  │")
+    lines.append(f"│ {'Trainable':<15} │ {trainable:>12,} │          │")
+    lines.append("└─────────────────┴──────────────┴───────────┘")
+    
+    return "\n".join(lines)