surfaces

BerkeleyAutomation · Sep 13, 2018 · a5238f6 · a5238f6
1 parent c43cb40
commit a5238f6
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 5 deletions.
diff --git a/fluids/assets/__init__.py b/fluids/assets/__init__.py
@@ -18,7 +18,9 @@
             PedCrossing,
             Pedestrian,
             Car,
-            TrafficLight,
+            "TrafficLight-Red",
+            "TrafficLight-Green",
+            "TrafficLight-Yellow",
             CrossWalkLight,
             CrossWalk,
             Waypoint]

diff --git a/fluids/obs/grid.py b/fluids/obs/grid.py
@@ -52,7 +52,9 @@ def __init__(self, car, obs_dim=500, shape=(500,500)):
         undrivable_window = pygame.Surface((self.grid_dim, self.grid_dim))
         car_window        = pygame.Surface((self.grid_dim, self.grid_dim))
         ped_window        = pygame.Surface((self.grid_dim, self.grid_dim))
-        light_window      = pygame.Surface((self.grid_dim, self.grid_dim))
+        light_window_red  = pygame.Surface((self.grid_dim, self.grid_dim))
+        light_window_green= pygame.Surface((self.grid_dim, self.grid_dim))
+        light_window_yellow=pygame.Surface((self.grid_dim, self.grid_dim))
         direction_window  = pygame.Surface((self.grid_dim, self.grid_dim))
         direction_pixel_window \
                           = pygame.Surface((self.grid_dim, self.grid_dim))
@@ -88,9 +90,18 @@ def __init__(self, car, obs_dim=500, shape=(500,500)):
         for obj in collideable_map[Pedestrian]:
             rel_obj = obj.get_relative(rel)
             rel_obj.render(ped_window, border=None)
-        for obj in collideable_map[TrafficLight]:
+
+        for obj in collideable_map["TrafficLight-Red"]:
+            rel_obj = obj.get_relative(rel)
+            rel_obj.render(light_window_red, border=None)
+
+        for obj in collideable_map["TrafficLight-Green"]:
+            rel_obj = obj.get_relative(rel)
+            rel_obj.render(light_window_green, border=None)
+
+        for obj in collideable_map["TrafficLight-Yellow"]:
             rel_obj = obj.get_relative(rel)
-            rel_obj.render(light_window, border=None)
+            rel_obj.render(light_window_green, border=None)
 
         point = (int(gd/6), int(gd/2))
         edge_point = None
@@ -128,7 +139,9 @@ def is_on_screen(point, gd):
                                                                               undrivable_window,
                                                                               car_window,
                                                                               ped_window,
-                                                                              light_window,
+                                                                              light_window_red,
+                                                                              light_window_green,
+                                                                              light_window_yellow,
                                                                               direction_window,
                                                                               direction_pixel_window,
                                                                               direction_edge_window