learn_slam/slam/main.py

import math

import pygame

from slam.environment import Environment
from slam.sensors import Lidar2DSensor

PYGAME_RUNNING = True
SENSOR_ON = False  # Mouse cursor


def main():
    global PYGAME_RUNNING
    global SENSOR_ON
    environment = Environment(map_path="../data/slam_demo_floor_plan.png", map_height=600, map_width=1200)
    environment.original_map = environment.map.copy()
    laser = Lidar2DSensor(200, environment.original_map, uncertainty=(0.5, 0.01))
    environment.map.fill(environment.colors.black)
    environment.info_map = environment.map.copy()

    while PYGAME_RUNNING:

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                PYGAME_RUNNING = False
            if pygame.mouse.get_focused():
                SENSOR_ON = True
            else:
                SENSOR_ON = False

        if SENSOR_ON:
            sim_robot_position = pygame.mouse.get_pos()
            laser.position = sim_robot_position
            sensor_data = laser.sense_obstacles()
            # Draw data to map
            environment.store_points(sensor_data)
            environment.show_sensor_data()
        # update the map
        environment.map.blit(environment.info_map, (0, 0))
        # display the map
        pygame.display.update()

    exit(0)


if __name__ == "__main__":
    main()