Refactoring

cooperative_cuisine/pygame_2d_vis/grid.py → cooperative_cuisine/pygame_2d_vis/generate_layout_thumbnail.py

 import uuid
-from pprint import pprint
 
-import matplotlib.pyplot as plt
 import numpy as np
 import yaml
-from collections import defaultdict
 
 from cooperative_cuisine import ROOT_DIR
 from cooperative_cuisine.counter_factory import convert_words_to_chars
-from cooperative_cuisine.items import ItemInfo, ItemType
-from cooperative_cuisine.pygame_2d_vis.drawing import Visualizer
+from cooperative_cuisine.items import ItemType
 
 dummy_env_config_path = ROOT_DIR / "configs" / "dummy_environment_config.yaml"
 with open(dummy_env_config_path) as c:
     dummy_env_config = yaml.load(c, Loader=yaml.FullLoader)
 character_map = dummy_env_config["layout_chars"]
 character_map = convert_words_to_chars(character_map)
+counter_set = {"Counter", "PlateDispenser", "CuttingBoard", "Trashcan", "ServingWindow", "Sink", "SinkAddon"}
 
 
 def layout_thumbnail(layout_path, vis, max_size, item_lookup):
     with open(layout_path) as file:
         layout = file.read()
 
-    colors = defaultdict(lambda: 0)
-
-    v = list(set(character_map.values()))
-    idx_free = v.index("Free")
-    zero = v[0]
-    v[0] = "Free"
-    v[idx_free] = zero
-
-    colors.update({c: i for i, c in enumerate(v)})
-
     grid = []
     counter = []
     player = []
@@ -42,7 +29,7 @@ def layout_thumbnail(layout_path, vis, max_size, item_lookup):
             continue
         for y, char in enumerate(line.strip()):
             t = character_map[char]
-            row.append(colors[t])
+            row.append(0)
             match t:
                 case "Agent":
                     player.append([float(y), float(x)])
@@ -56,18 +43,18 @@ def layout_thumbnail(layout_path, vis, max_size, item_lookup):
     image = np.array(grid).T
 
     def map_t(t):
-        if t in {"Counter", "PlateDispenser", "CuttingBoard", "Trashcan", "ServingWindow", "Sink", "SinkAddon"}:
+        if t in counter_set:
             return t
         if t in item_lookup:
             i = item_lookup[t]
             if i.type == ItemType.Equipment and i.equipment:
                 return i.equipment.name
             if i.type == ItemType.Ingredient:
-                return t + "Dispenser"
+                return f"{t}Dispenser"
         return "Counter"
 
     def map_t_occupied(t):
-        if t in {"Counter", "PlateDispenser", "CuttingBoard", "Trashcan", "ServingWindow", "Sink", "SinkAddon"}:
+        if t in counter_set:
             return None
         if t in item_lookup:
             i = item_lookup[t]
@@ -75,23 +62,26 @@ def layout_thumbnail(layout_path, vis, max_size, item_lookup):
                 return {"category": "ItemCookingEquipment", "type": t, "content_list": [], "content_ready": None,
                         "active_effects": []}
             if i.type == ItemType.Ingredient:
-                return {"category": "Item", "type": t, "content_list": [], "content_ready": None, "active_effects": []}
+                return {"category": "Item", "type": t, "active_effects": []}
         return None
 
     state = {"players": [
-        {"id": str(i), "pos": p, "facing_direction": [0, 1], "holding": None, "current_nearest_counter_pos": None,
+        {"id": str(i),
+         "pos": p,
+         "facing_direction": [0, 1],
+         "holding": None,
+         "current_nearest_counter_pos": None,
          "current_nearest_counter_id": None} for i, p in enumerate(player)],
-        "counters": [{"id": uuid.uuid4().hex, "pos": pos, "orientation": [0, 1], "type": map_t(t),
-                      "occupied_by": map_t_occupied(t), "active_effects": []} for pos, t in counter],
-        "kitchen": {"width": image.shape[0], "height": image.shape[1]},
+        "counters": [{"id": uuid.uuid4().hex,
+                      "pos": pos,
+                      "orientation": [0, 1],
+                      "type": map_t(t),
+                      "occupied_by": map_t_occupied(t),
+                      "active_effects": []} for pos, t in counter],
+        "kitchen": {"width": image.shape[0],
+                    "height": image.shape[1]},
     }
 
-    # im = viz.get_state_image(grid_size=32, state=state).transpose((1, 0, 2))
-
-    grid_size = max_size / max(state["kitchen"]["width"], state["kitchen"]["height"])
+    grid_size = int(max_size // max(state["kitchen"]["width"], state["kitchen"]["height"]))
     vis.set_grid_size(grid_size)
-    surf = vis.draw_gamescreen(state, [])
-    return surf
-
-    # plt.imshow(im)
-    # plt.show()
+    return vis.draw_gamescreen(state, [])