generate_layout_thumbnail.py

import hashlib
import uuid

import numpy as np
import yaml
import pygame

from cooperative_cuisine import ROOT_DIR
from cooperative_cuisine.counter_factory import convert_words_to_chars
from cooperative_cuisine.items import ItemType, ItemInfo

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 generate_item_lookup():
    with open(ROOT_DIR / "configs" / "item_info.yaml") as file:
        item_info = file.read()
    item_lookup = yaml.safe_load(item_info)
    for item_name in item_lookup:
        item_lookup[item_name] = ItemInfo(name=item_name, **item_lookup[item_name])
    for item_name, item_info in item_lookup.items():
        if item_info.equipment:
            item_info.equipment = item_lookup[item_info.equipment]
    return item_lookup


def layout_thumbnail(layout_path, vis, max_size, item_lookup, layout_hashes, layout_string: str | None = None,
                     as_np=False):
    if layout_string is None:
        with open(layout_path) as file:
            layout = file.read()
    else:
        layout = layout_string
    
    thumbnail_path = ROOT_DIR / "generated" / "layout_thumbnails"
    thumbnail_path.mkdir(parents=True, exist_ok=True)
    
    l_hash = hashlib.sha1(layout.encode("utf-8")).hexdigest()
    if l_hash in layout_hashes:
        return pygame.image.load(
            thumbnail_path / f"layout-{l_hash}.png"
        ).convert_alpha()
    
    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)
    image = vis.draw_gamescreen(state, [])
    pygame.image.save(image, thumbnail_path / f"layout-{l_hash}.png")
    
    if as_np:
        return pygame.surfarray.pixels3d(image).transpose((1, 0, 2))
    else:
        return image