[bugfix] - color overflow in get_island_colored_map

src_python/habitat_sim/utils/viz_utils.py

@@ -12,6 +12,7 @@
 if "google.colab" in sys.modules:
     os.environ["IMAGEIO_FFMPEG_EXE"] = "/usr/bin/ffmpeg"
 
+import random
 from typing import Any, Dict, List, Optional, Tuple
 
 import imageio
@@ -300,11 +301,22 @@ def get_island_colored_map(island_top_down_map_data: np.ndarray):
     white = int("0xffffff", base=16)
     island_map = Image.new("RGB", island_top_down_map_data.shape, color=white)
     pixels = island_map.load()
+    extra_colors: List[int] = []
+    r = lambda: random.randint(0, 255)
     for x in range(island_top_down_map_data.shape[0]):
         for y in range(island_top_down_map_data.shape[1]):
             if island_top_down_map_data[x, y] >= 0:
-                pixels[x, y] = int(
-                    d3_40_colors_hex[island_top_down_map_data[x, y]], base=16
-                )
-
+                color_index = island_top_down_map_data[x, y]
+                if color_index < len(d3_40_colors_hex):
+                    # fixed colors from a selected list
+                    # NOTE: PIL Image origin is top left, so invert y
+                    pixels[x, -y] = int(d3_40_colors_hex[color_index], base=16)
+                else:
+                    random_color_index = color_index - len(d3_40_colors_hex)
+                    # pick random colors once fixed colors are overflowed
+                    while random_color_index >= len(extra_colors):
+                        new_color = int(("0x%02X%02X%02X" % (r(), r(), r())), base=16)
+                        if new_color not in extra_colors:
+                            extra_colors.append(new_color)
+                    pixels[x, -y] = extra_colors[random_color_index]
     return island_map