optitrack_stream.py

import argparse
import time
import sys
import signal
from NatNetClient import NatNetClient
from collections import defaultdict
import math

class SwarmTracker:
    def __init__(self, server_ip="127.0.0.1", client_ip="127.0.0.1", display_mode="summary", display_rate=2.0):
        self.running = True
        self.server_ip = server_ip
        self.client_ip = client_ip
        self.last_positions = {}
        self.robot_velocities = {}
        self.last_update_times = {}
        
        # Configuration
        self.show_rigid_bodies = True
        self.show_markers = False  # Turn off individual markers
        self.show_velocities = True
        self.update_threshold = 0.001  # Minimum position change to report (meters)
        self.velocity_smoothing = 0.3  # Exponential smoothing factor
        
        # Display rate limiting
        self.display_rate = display_rate  # Hz - how often to update display (2 = twice per second)
        self.last_display_time = 0
        self.display_mode = display_mode  # 'summary', 'detailed', or 'minimal'
        self.clear_screen = True  # Clear screen between updates
        
        # Data buffering
        self.latest_robot_data = {}
        
        # Stats tracking
        self.frame_count = 0
        self.last_stats_time = time.time()
        self.robot_stats = defaultdict(lambda: {'updates': 0, 'last_seen': 0})
        
        signal.signal(signal.SIGINT, self.signal_handler)
        
    def signal_handler(self, sig, frame):
        print("\n🛑 Shutting down gracefully...")
        self.running = False
        
    def calculate_velocity(self, robot_id, new_pos, timestamp):
        """Calculate velocity using exponential smoothing"""
        if robot_id in self.last_positions and robot_id in self.last_update_times:
            dt = timestamp - self.last_update_times[robot_id]
            if dt > 0:
                # Calculate instantaneous velocity
                old_pos = self.last_positions[robot_id]
                instant_vel = tuple((new_pos[i] - old_pos[i]) / dt for i in range(3))
                
                # Apply exponential smoothing
                if robot_id in self.robot_velocities:
                    smooth_vel = tuple(
                        self.velocity_smoothing * instant_vel[i] + 
                        (1 - self.velocity_smoothing) * self.robot_velocities[robot_id][i]
                        for i in range(3)
                    )
                    self.robot_velocities[robot_id] = smooth_vel
                else:
                    self.robot_velocities[robot_id] = instant_vel
                    
        self.last_update_times[robot_id] = timestamp
        
    def clear_terminal(self):
        """Clear terminal screen"""
        if sys.platform == 'win32':
            sys.stdout.write('\033[2J\033[H')  # ANSI escape codes
            sys.stdout.flush()
        else:
            sys.stdout.write('\033c')
            sys.stdout.flush()
    
    def should_display(self):
        """Check if enough time has passed to update display"""
        current_time = time.time()
        if current_time - self.last_display_time >= (1.0 / self.display_rate):
            self.last_display_time = current_time
            return True
        return False
    
    def display_summary(self):
        """Display a clean summary of all robots"""
        if self.clear_screen:
            self.clear_terminal()
        
        print(f"╔{'═'*78}╗")
        print(f"║ OptiTrack Swarm Tracker - {time.strftime('%H:%M:%S')} - Frame: {self.frame_count:<10} ║")
        print(f"╠{'═'*78}╣")
        print(f"║ {'Robot ID':<12} {'Position (x, y, z)':<30} {'Speed':<10} {'Status':<15} ║")
        print(f"╠{'═'*78}╣")
        
        current_time = time.time()
        
        # Sort robots by ID for consistent display
        for robot_id in sorted(self.latest_robot_data.keys()):
            data = self.latest_robot_data[robot_id]
            pos = data['position']
            
            # Calculate time since last update
            time_since = current_time - data['timestamp']
           
            status = "✅ Active"
            
            
            # Get speed if available
            speed_str = "---"
            if robot_id in self.robot_velocities:
                vel = self.robot_velocities[robot_id]
                speed = math.sqrt(sum(v**2 for v in vel))
                speed_str = f"{speed:.3f} m/s"
            
            # Format position
            pos_str = f"{pos[0]:7.3f}, {pos[1]:7.3f}, {pos[2]:7.3f}"
            
            print(f"║ {robot_id:<12} {pos_str:<30} {speed_str:<10} {status:<15} ║")
        
        print(f"╚{'═'*78}╝")
        print(f"\nDisplay Rate: {self.display_rate} Hz | Mode: {self.display_mode} | "
              f"Press Ctrl+C to exit")
    
    def display_detailed(self):
        """Display detailed information for each robot"""
        if self.clear_screen:
            self.clear_terminal()
        
        print(f"\n{'='*80}")
        print(f"Frame {self.frame_count} | Time: {time.strftime('%H:%M:%S')}")
        print(f"{'='*80}\n")
        
        for robot_id in sorted(self.latest_robot_data.keys()):
            data = self.latest_robot_data[robot_id]
            print(f"🤖 {robot_id}:")
            print(f"   Position: {self.format_vector(data['position'])}")
            print(f"   Rotation: {self.format_vector(data['rotation'], 4)}")
            
            if robot_id in self.robot_velocities:
                vel = self.robot_velocities[robot_id]
                speed = math.sqrt(sum(v**2 for v in vel))
                print(f"   Velocity: {self.format_vector(vel)} | Speed: {speed:.3f} m/s")
            
            if 'error' in data:
                print(f"   Tracking Error: {data['error']:.4f}")
            print()
    
    def display_minimal(self):
        """Display minimal one-line updates"""
        for robot_id in sorted(self.latest_robot_data.keys()):
            data = self.latest_robot_data[robot_id]
            pos = data['position']
            print(f"{robot_id}: [{pos[0]:6.3f}, {pos[1]:6.3f}, {pos[2]:6.3f}]", end="  ")
        print()  # New line at end
        
    def has_position_changed(self, robot_id, new_pos):
        """Check if position has changed significantly"""
        if robot_id not in self.last_positions:
            return True
            
        old_pos = self.last_positions[robot_id]
        distance = math.sqrt(sum((new_pos[i] - old_pos[i])**2 for i in range(3)))
        return distance > self.update_threshold
        
    def receive_frame_data(self, data_dict):
        """Main callback for processing frame data"""
        self.frame_count += 1
        timestamp = data_dict.get("timestamp", time.time())
        
        # Process MoCap data if available
        mocap_data = data_dict.get("mocap_data")
        if not mocap_data:
            return
        
        # Process rigid bodies (your swarm robots)
        if self.show_rigid_bodies:
            rigid_body_data = mocap_data.rigid_body_data
            if rigid_body_data:
                for rb in rigid_body_data.rigid_body_list:
                    if not rb.tracking_valid:
                        continue
                    
                    robot_id = f"Robot_{rb.id_num}"
                    pos = (rb.pos[0], rb.pos[1], rb.pos[2])
                    
                    # Update velocity calculations
                    if self.has_position_changed(robot_id, pos):
                        self.calculate_velocity(robot_id, pos, timestamp)
                        self.last_positions[robot_id] = pos
                    
                    # Store latest data for display
                    self.latest_robot_data[robot_id] = {
                        'position': pos,
                        'rotation': (rb.rot[0], rb.rot[1], rb.rot[2], rb.rot[3]),
                        'timestamp': timestamp,
                        'error': getattr(rb, 'error', 0.0)
                    }
                    
                    # Update stats
                    self.robot_stats[robot_id]['updates'] += 1
                    self.robot_stats[robot_id]['last_seen'] = timestamp
        
        # Only update display at specified rate
        if self.should_display():
            if self.display_mode == 'summary':
                self.display_summary()
            elif self.display_mode == 'detailed':
                self.display_detailed()
            elif self.display_mode == 'minimal':
                self.display_minimal()
    
    def print_statistics(self):
        """Print tracking statistics"""
        print("\n" + "="*60)
        print("📊 TRACKING STATISTICS")
        print("="*60)
        
        current_time = time.time()
        for robot_id, stats in sorted(self.robot_stats.items()):
            time_since_seen = current_time - stats['last_seen']
            status = "✅ Active"
            print(f"{robot_id}: {status} | Updates: {stats['updates']} | "
                  f"Last seen: {time_since_seen:.1f}s ago")
        
        print(f"\nTotal frames processed: {self.frame_count}")
        print("="*60 + "\n")
    
    def toggle_markers(self):
        """Toggle individual marker display"""
        self.show_markers = not self.show_markers
        print(f"\n{'📍 Showing' if self.show_markers else '🚫 Hiding'} individual markers")
    
    def run(self):
        """Main execution function"""
        # Create NatNet client
        self.client = NatNetClient()
        self.client.set_print_level(0)  # Reduce debug output
        
        # Configure connection (adjust these for your setup)
        self.client.set_client_address(self.client_ip)
        self.client.set_server_address(self.server_ip)
        self.client.set_use_multicast(False)
        
        # Set callback
        self.client.new_frame_with_data_listener = self.receive_frame_data
        
        print("🎯 Starting SwarmTracker for OptiTrack")
        print("\nDisplay Settings:")
        print(f"  - Update Rate: {self.display_rate} Hz (updates {self.display_rate} times per second)")
        print(f"  - Display Mode: {self.display_mode}")
        print(f"  - Clear Screen: {'Yes' if self.clear_screen else 'No'}")
        print("\nControls:")
        print("  - Ctrl+C: Exit")
        print("\nUse --rate and --mode to change the display (see --help)\n")
        
        # Start client
        if not self.client.run('d') or not self.client.connected():
            print("❌ Failed to connect to Motive")
            print("Check that:")
            print("  1. Motive is running")
            print("  2. Streaming is enabled (View -> Data Streaming)")
            print("  3. IP addresses are correct")
            return False
        
        print("✅ Connected successfully!")
        print("Waiting for robot data...\n")
        time.sleep(1)  # Give user time to read before clearing screen
        
        try:
            while self.running:
                time.sleep(0.05)  # Small sleep to prevent CPU spinning
                
        except KeyboardInterrupt:
            pass
        finally:
            print("\n🔻 Cleaning up...")
            self.client.shutdown()
            print("✅ Shutdown complete")
            
            # Final statistics
            if self.robot_stats:
                print("\nFinal Statistics:")
                for robot_id, stats in sorted(self.robot_stats.items()):
                    print(f"{robot_id}: Total updates: {stats['updates']}")
        
        return True


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Stream OptiTrack rigid bodies from Motive and show live robot positions and speeds.")
    parser.add_argument("--server", default="127.0.0.1", help="IP of the machine running Motive")
    parser.add_argument("--client", default="127.0.0.1", help="IP of this machine on the same network")
    parser.add_argument("--mode", choices=["summary", "detailed", "minimal"], default="summary", help="display layout")
    parser.add_argument("--rate", type=float, default=2.0, help="display refreshes per second")
    args = parser.parse_args()
    tracker = SwarmTracker(args.server, args.client, args.mode, args.rate)
    tracker.run()