heatmap.py

1import os2import sys34# MRLYPROD56HERE = os.path.dirname(os.path.abspath(__file__))7MRLYPROD = os.path.normpath(os.path.join(HERE, ".."))89if not os.path.isdir(os.path.join(MRLYPROD, "mrlypy", "six")):10    sys.exit(f"missing mrlypy: expected it at {MRLYPROD}")1112sys.path.insert(0, MRLYPROD)1314import mrlypy.two as m215import numpy as np16from config import DATA_DIR, IMAGE_SIZE17from helpers import hex_key18from PIL import Image19from typing import Iterable2021INVERT = False22LEVEL = 123LIMIT = 5024MODE = "FIT"25PAINT = "GRAYSCALE"26ODDS = [i for i in range(1, LIMIT, 2)]27NUMBERS = ODDS28FUNCS = [29    m2.carpet_2d,30    m2.net_2d,31    m2.tree_2d,32    m2.void_2d,33]3435def pad_grid(grid: np.ndarray, max_height: int, max_width: int) -> np.ndarray:36    pad_height = max_height - grid.shape[0]37    pad_width = max_width - grid.shape[1]38    pad_top = pad_height // 239    pad_bottom = pad_height - pad_top40    pad_left = pad_width // 241    pad_right = pad_width - pad_left42    return np.pad(grid, ((pad_top, pad_bottom), (pad_left, pad_right)), 'constant')4344def resize_grid(grid: np.ndarray) -> np.ndarray:45    if grid.shape == IMAGE_SIZE:46        return grid47    old_height, old_width = grid.shape48    new_y, new_x = np.indices(IMAGE_SIZE)49    old_x = (new_x * old_width / IMAGE_SIZE[1]).astype(int)50    old_y = (new_y * old_height / IMAGE_SIZE[0]).astype(int)51    return grid[old_y, old_x]5253def create_color_map(max_value: int) -> np.ndarray:54    if max_value == 0:55        return np.array([[0, 0, 0]], dtype=np.uint8)56    grayscale_map = np.linspace(0, 255, max_value + 1, dtype=np.uint8)57    color_maps = {58        "GRAYSCALE": np.stack([grayscale_map] * 3, axis=1),59        "RANDOM": np.random.randint(0, 256, size=(max_value + 1, 3), dtype=np.uint8),60    }61    color_map = color_maps[PAINT]62    color_map[0] = [0, 0, 0]63    return color_map6465def draw_step(grid: np.ndarray, name: str):66    fp = f"{DATA_DIR}/{name}.png"67    max_val = np.max(grid)68    color_map = create_color_map(max_val)69    image_array = color_map[grid]70    image = Image.fromarray(image_array, 'RGB')71    image = image.resize(IMAGE_SIZE, Image.Resampling.NEAREST)72    image.save(fp, "PNG")73    print(f"Saved: {fp}")7475def generate_steps(title: str, grids: Iterable[np.ndarray], mode: str):76    heatmap = None77    if mode == "FIT":78        heatmap = np.zeros(IMAGE_SIZE, dtype=np.int16)79        for grid in grids:80            heatmap += resize_grid(grid)81    elif mode == "GROW":82        for grid in grids:83            if heatmap is None:84                heatmap = grid.astype(np.int16)85            else:86                h_h, h_w = heatmap.shape87                g_h, g_w = grid.shape88                new_h = max(h_h, g_h)89                new_w = max(h_w, g_w)90                if new_h > h_h or new_w > h_w:91                    heatmap = pad_grid(heatmap, new_h, new_w)92                padded_grid = pad_grid(grid, new_h, new_w)93                heatmap += padded_grid94    if heatmap is not None:95        draw_step(heatmap, f"{title}_{hex_key(4)}")9697def main():98    for func in FUNCS:99        if INVERT:100            grids = (func(number, LEVEL).invert().to_array() for number in NUMBERS)101        else:102            grids = (func(number, LEVEL).to_array() for number in NUMBERS)103        generate_steps(func.__name__, grids, MODE)104105if __name__ == "__main__":106    main()