generate_layout_thumbnail.py

import uuid

import numpy as np
import yaml

from cooperative_cuisine import ROOT_DIR
from cooperative_cuisine.counter_factory import convert_words_to_chars
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()

    grid = []
    counter = []
    player = []
    x = 0
    for line in layout.split("\n"):
        row = []
        if line.startswith(";") or not line:
            continue
        for y, char in enumerate(line.strip()):
            t = character_map[char]
            row.append(0)
            match t:
                case "Agent":
                    player.append([float(y), float(x)])
                case "Free" | "Water" | "Ice" | "Lava":
                    pass
                case _:
                    counter.append(([float(y), float(x)], t))
        grid.append(row)
        x += 1

    image = np.array(grid).T

    def map_t(t):
        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 f"{t}Dispenser"
        return "Counter"

    def map_t_occupied(t):
        if t in counter_set:
            return None
        if t in item_lookup:
            i = item_lookup[t]
            if i.type == ItemType.Equipment and i.equipment:
                return {"category": "ItemCookingEquipment", "type": t, "content_list": [], "content_ready": None,
                        "active_effects": []}
            if i.type == ItemType.Ingredient:
                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,
         "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]},
    }

    grid_size = int(max_size // max(state["kitchen"]["width"], state["kitchen"]["height"]))
    vis.set_grid_size(grid_size)
    return vis.draw_gamescreen(state, [])