You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
58 lines
2.1 KiB
Python
58 lines
2.1 KiB
Python
import math
|
|
from collections import namedtuple
|
|
|
|
import pygame
|
|
|
|
RGB_COLOR_CODES = [
|
|
(0, 0, 0), # black
|
|
(70, 70, 70), # grey
|
|
(0, 0, 255), # blue
|
|
(0, 255, 0), # green
|
|
(255, 0, 0), # red
|
|
(255, 255, 255) # white
|
|
]
|
|
|
|
Colors = namedtuple("Colors", "black,grey,blue,green,red,white", defaults=RGB_COLOR_CODES)
|
|
|
|
|
|
class Environment:
|
|
|
|
def __init__(self, map_path: str, map_height: int, map_width: int, window_name: str = "RRT path planning",
|
|
colors: Colors = None):
|
|
self.point_cloud = []
|
|
self.external_map = pygame.image.load(map_path)
|
|
self.map_height = map_height
|
|
self.map_width = map_width
|
|
self.window_name = window_name
|
|
self.colors = colors if colors else Colors()
|
|
self._initialize_pygame()
|
|
|
|
def _initialize_pygame(self):
|
|
pygame.init()
|
|
pygame.display.set_caption(self.window_name)
|
|
self.map = pygame.display.set_mode((self.map_width, self.map_height))
|
|
# overlay external map
|
|
self.map.blit(self.external_map, (0, 0))
|
|
|
|
def calc_2d_position(self, distance, angle, robot_position):
|
|
"""Calculates a 2D point from a robots position and a Lidar sensor reading."""
|
|
x = distance * math.cos(angle) + robot_position[0]
|
|
y = -distance * math.sin(angle) + robot_position[1]
|
|
return (int(x), int(y))
|
|
|
|
# TODO add type hints
|
|
def store_points(self, sensor_sweep_data):
|
|
"""Stores sensor reading in the environment's point cloud."""
|
|
for element in sensor_sweep_data:
|
|
# TODO this could be refactored into a LidarReading & RobotPosition named tuples
|
|
point = self.calc_2d_position(element[0], element[1], element[2])
|
|
# check for dupes
|
|
# TODO ok an optimization could probably be added here too
|
|
if point not in self.point_cloud:
|
|
self.point_cloud.append(point)
|
|
|
|
def show_sensor_data(self):
|
|
# Draw data on new map
|
|
self.info_map = self.map.copy()
|
|
for point in self.point_cloud:
|
|
self.info_map.set_at((int(point[0]), int(point[1])), self.colors.red) |