PiLidarMapperTest.py

import pygame
import sys
import math
import time
from rplidar import RPLidar
import numpy as np
import os
import threading

class LidarMapper:
    def __init__(self, port='/dev/ttyUSB0', baudrate=115200):
        self.port = port
        self.baudrate = baudrate
        self.lidar = None
        self.is_scanning = False
        self.scan_thread = None
        
        # Pygame settings
        self.screen_size = 600
        self.screen = None
        self.center_x = self.screen_size // 2
        self.center_y = self.screen_size // 2
        self.scale = 0.3
        self.point_size = 2
        
        # Colors
        self.bg_color = (10, 10, 40)
        self.robot_color = (255, 50, 50)
        self.axis_color = (50, 255, 50)
        self.text_color = (255, 255, 100)
        self.measure_color = (255, 255, 0)
        
        # Data storage
        self.all_points = []
        self.scan_count = 0
        self.start_time = time.time()
        self.font = None
        
        # Settings
        self.show_measurements = True
        self.room_data = None
        self.color_mode = 'quality'
        
        # Error handling
        self.error_count = 0
        self.max_errors = 5
        
    def connect(self):
        """Connect to RPLidar with robust error handling"""
        try:
            print(f"🔌 Connecting to RPLidar on {self.port}...")
            self.lidar = RPLidar(self.port, baudrate=self.baudrate, timeout=3)
            time.sleep(2)
            
            # Get device info
            info = self.lidar.get_info()
            health = self.lidar.get_health()
            
            print("✅ Connected to RPLidar successfully!")
            print(f"   Model: {info['model']}")
            print(f"   Firmware: {info['firmware']}")
            print(f"   Health: {health}")
            
            # Reset error counter on successful connection
            self.error_count = 0
            return True
            
        except Exception as e:
            print(f"❌ Connection error: {e}")
            self.error_count += 1
            return False
    
    def init_pygame(self):
        """Initialize Pygame"""
        try:
            pygame.init()
            self.screen = pygame.display.set_mode((self.screen_size, self.screen_size))
            pygame.display.set_caption("RPLidar Room Mapper - Stable Version")
            
            self.font = pygame.font.SysFont('Arial', 16)
            self.large_font = pygame.font.SysFont('Arial', 18, bold=True)
            
            print("✅ Pygame initialized successfully")
            return True
        except Exception as e:
            print(f"❌ Pygame init error: {e}")
            return False
    
    def get_color_by_quality(self, quality):
        """Get color based on quality"""
        q = max(0, min(quality, 255))
        if q < 128:
            r = 255
            g = int(255 * (q / 128))
            b = 0
        else:
            r = int(255 * (1 - (q - 128) / 127))
            g = 255
            b = 0
        return (r, g, b)
    
    def get_color_by_distance(self, distance):
        """Get color based on distance"""
        if distance > 16000:
            distance = 16000
        normalized = int(255 * (distance / 16000))
        if normalized < 85:
            r = 0
            g = int(255 * (normalized / 85))
            b = 255
        elif normalized < 170:
            r = 0
            g = 255
            b = int(255 * (1 - (normalized - 85) / 85))
        else:
            r = int(255 * ((normalized - 170) / 85))
            g = 255 - int(128 * ((normalized - 170) / 85))
            b = 0
        return (r, g, b)
    
    def robust_scan_worker(self):
        """Robust scanning with comprehensive error handling"""
        print("🔄 Starting robust scan worker...")
        consecutive_errors = 0
        max_consecutive_errors = 3
        
        while self.is_scanning and consecutive_errors < max_consecutive_errors:
            try:
                # Simple scanning without extra parameters
                for scan in self.lidar.iter_scans():
                    if not self.is_scanning:
                        break
                    
                    new_points = []
                    points_processed = 0
                    
                    for quality, angle, distance in scan:
                        if not self.is_scanning:
                            break
                            
                        if distance > 0 and distance < 16000:
                            # Process only every 3rd point to reduce load
                            if points_processed % 3 == 0:
                                angle_rad = np.radians(angle)
                                x = distance * np.cos(angle_rad)
                                y = distance * np.sin(angle_rad)
                                
                                if self.color_mode == 'quality':
                                    color = self.get_color_by_quality(quality)
                                else:
                                    color = self.get_color_by_distance(distance)
                                
                                new_points.append((x, y, color))
                            
                            points_processed += 1
                    
                    if new_points and self.is_scanning:
                        self.all_points.extend(new_points)
                        
                        # Keep limited points for performance
                        if len(self.all_points) > 600:
                            self.all_points = self.all_points[-600:]
                        
                        self.scan_count += 1
                        
                        # Reset error counter on successful scan
                        consecutive_errors = 0
                        
                        # Print status occasionally
                        if self.scan_count % 25 == 0:
                            print(f"📊 Scan {self.scan_count}: {len(new_points)} points")
                
            except Exception as e:
                consecutive_errors += 1
                print(f"⚠️ Scan error #{consecutive_errors}: {e}")
                
                if consecutive_errors >= max_consecutive_errors:
                    print("🔴 Too many consecutive errors, stopping scan worker")
                    break
                
                # Wait before retry
                time.sleep(1)
        
        print("🛑 Scan worker stopped")
        self.is_scanning = False
    
    def start_scan(self):
        """Start the scanning process"""
        try:
            print("🚀 Starting scan...")
            self.is_scanning = True
            
            # Start scan thread
            self.scan_thread = threading.Thread(target=self.robust_scan_worker, daemon=True)
            self.scan_thread.start()
            
            # Main display loop
            clock = pygame.time.Clock()
            last_measurement_time = time.time()
            
            while self.is_scanning:
                current_time = time.time()
                
                # Update measurements every 2 seconds
                if current_time - last_measurement_time > 2.0:
                    self.room_data = self.measure_room()
                    last_measurement_time = current_time
                
                self.draw_frame()
                
                if self.handle_events():
                    break
                
                # Conservative frame rate
                clock.tick(15)
                
        except Exception as e:
            print(f"❌ Main loop error: {e}")
        finally:
            self.stop_scan()
    
    def measure_room(self):
        """Calculate room dimensions"""
        if len(self.all_points) > 30:
            try:
                xs = [p[0] for p in self.all_points]
                ys = [p[1] for p in self.all_points]
                
                if not xs or not ys:
                    return None
                    
                width = max(xs) - min(xs)
                height = max(ys) - min(ys)
                
                room_area = width * height / 1000000
                perimeter = 2 * (width + height)
                center_x = (max(xs) + min(xs)) / 2
                center_y = (max(ys) + min(ys)) / 2
                
                return {
                    'width': width,
                    'height': height, 
                    'area': room_area,
                    'perimeter': perimeter,
                    'center': (center_x, center_y),
                    'bounds': (min(xs), min(ys), max(xs), max(ys)),
                    'point_count': len(self.all_points)
                }
            except Exception as e:
                print(f"⚠️ Measurement error: {e}")
                return None
        return None
    
    def draw_frame(self):
        """Draw one frame"""
        self.screen.fill(self.bg_color)
        
        # Draw axes
        pygame.draw.line(self.screen, self.axis_color, 
                        (self.center_x, 0), (self.center_x, self.screen_size), 1)
        pygame.draw.line(self.screen, self.axis_color, 
                        (0, self.center_y), (self.screen_size, self.center_y), 1)
        
        # Draw points
        if self.all_points:
            for x, y, color in self.all_points:
                screen_x = int(x * self.scale) + self.center_x
                screen_y = int(y * self.scale) + self.center_y
                
                if 0 <= screen_x < self.screen_size and 0 <= screen_y < self.screen_size:
                    pygame.draw.circle(self.screen, color, (screen_x, screen_y), self.point_size)
        
        # Draw robot
        pygame.draw.circle(self.screen, self.robot_color, 
                         (self.center_x, self.center_y), 8)
        
        # Draw measurements and stats
        self.draw_room_measurements()
        self.draw_stats()
        
        pygame.display.flip()
    
    def draw_room_measurements(self):
        """Draw room measurements"""
        if not self.room_data or not self.show_measurements:
            return
            
        try:
            min_x, min_y, max_x, max_y = self.room_data['bounds']
            
            screen_min_x = int(min_x * self.scale) + self.center_x
            screen_min_y = int(min_y * self.scale) + self.center_y
            screen_max_x = int(max_x * self.scale) + self.center_x
            screen_max_y = int(max_y * self.scale) + self.center_y
            
            rect = pygame.Rect(screen_min_x, screen_min_y, 
                              screen_max_x - screen_min_x, 
                              screen_max_y - screen_min_y)
            pygame.draw.rect(self.screen, self.measure_color, rect, 2)
            
            width_m = self.room_data['width'] / 1000.0
            height_m = self.room_data['height'] / 1000.0
            area_m2 = self.room_data['area']
            
            measurements = [
                "ROOM MEASUREMENTS:",
                f"Width: {width_m:.2f} m",
                f"Height: {height_m:.2f} m", 
                f"Area: {area_m2:.2f} m²",
                f"Points: {self.room_data['point_count']}",
                f"Scans: {self.scan_count}",
                f"Color: {self.color_mode.upper()}"
            ]
            
            measure_bg = pygame.Rect(self.screen_size - 220, 80, 210, 170)
            pygame.draw.rect(self.screen, (0, 0, 0, 180), measure_bg)
            pygame.draw.rect(self.screen, self.measure_color, measure_bg, 1)
            
            for i, text in enumerate(measurements):
                color = self.measure_color if i == 0 else self.text_color
                font = self.large_font if i == 0 else self.font
                text_surface = font.render(text, True, color)
                self.screen.blit(text_surface, (self.screen_size - 210, 90 + i * 25))
        except Exception as e:
            print(f"⚠️ Drawing measurements error: {e}")
    
    def draw_stats(self):
        """Draw statistics"""
        elapsed_time = time.time() - self.start_time
        
        stats = [
            f"Time: {elapsed_time:.1f}s",
            f"Zoom: {self.scale:.2f}",
            f"Points: {len(self.all_points)}",
            "ESC: Exit  M: Measurements",
            "+/-: Zoom  R: Reset",
            "C: Color Mode"
        ]
        
        for i, text in enumerate(stats):
            text_surface = self.font.render(text, True, self.text_color)
            self.screen.blit(text_surface, (10, 10 + i * 20))
    
    def handle_events(self):
        """Handle Pygame events"""
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return True
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    return True
                elif event.key == pygame.K_PLUS or event.key == pygame.K_EQUALS:
                    self.scale = min(self.scale * 1.2, 2.0)
                elif event.key == pygame.K_MINUS:
                    self.scale = max(self.scale * 0.8, 0.1)
                elif event.key == pygame.K_r:
                    self.all_points = []
                    self.scan_count = 0
                    self.start_time = time.time()
                    self.room_data = None
                    print("🔄 View reset")
                elif event.key == pygame.K_m:
                    self.show_measurements = not self.show_measurements
                elif event.key == pygame.K_s:
                    self.save_scan_data()
                elif event.key == pygame.K_c:
                    self.color_mode = 'distance' if self.color_mode == 'quality' else 'quality'
                    print(f"🎨 Color mode: {self.color_mode}")
        
        return False
    
    def save_scan_data(self):
        """Save scan data to file"""
        try:
            filename = f"lidar_scan_{time.strftime('%Y%m%d_%H%M%S')}.txt"
            with open(filename, 'w') as f:
                f.write(f"Lidar Scan Data - {time.ctime()}\n")
                f.write("=" * 50 + "\n")
                
                if self.room_data:
                    f.write(f"Room Width: {self.room_data['width']/1000:.3f} m\n")
                    f.write(f"Room Height: {self.room_data['height']/1000:.3f} m\n")
                    f.write(f"Room Area: {self.room_data['area']:.3f} m²\n")
                    f.write(f"Total Points: {self.room_data['point_count']}\n\n")
                
                f.write("Point Data (x, y in mm):\n")
                for point in self.all_points:
                    f.write(f"{point[0]:.1f}, {point[1]:.1f}\n")
            
            print(f"💾 Scan data saved to: {filename}")
            return True
            
        except Exception as e:
            print(f"❌ Error saving data: {e}")
            return False
    
    def stop_scan(self):
        """Stop scanning safely"""
        print("🛑 Stopping scan...")
        self.is_scanning = False
        
        if self.lidar:
            try:
                self.lidar.stop()
                time.sleep(1)
                self.lidar.disconnect()
                print("✅ Scan stopped successfully")
            except Exception as e:
                print(f"⚠️ Stop warning: {e}")
    
    def disconnect(self):
        """Clean up resources"""
        self.stop_scan()
        if pygame.get_init():
            pygame.quit()
        print("🔌 Disconnected successfully")

def main():
    print("🎯 RPLidar Room Mapper")
    print("=" * 50)
    
    port = '/dev/ttyUSB0'
    
    mapper = LidarMapper(port)
    
    try:
        if not mapper.init_pygame():
            return
        
        if mapper.connect():
            print("\n🎮 Controls:")
            print("  + / - : Zoom in/out")
            print("  M : Toggle measurements")
            print("  R : Reset view")
            print("  S : Save scan data")
            print("  C : Toggle color mode")
            print("  ESC : Exit")
            print("\n🚀 Starting room mapping...")
            
            mapper.start_scan()
        else:
            print("❌ Failed to connect to RPLidar")
            
    except KeyboardInterrupt:
        print("\n⏹️ Program interrupted by user")
    except Exception as e:
        print(f"❌ Unexpected error: {e}")
    finally:
        mapper.disconnect()

if __name__ == "__main__":
    main()