heatmap.py

1from mrlypy.core.colors import gradient2from mrlypy.two import Cell2d3import numpy as np4import os5from colors import create_heatmap_colors6from config import FORMAT7from models import Saga89def create_saga_heatmap(saga: Saga, output_dir: str) -> Saga:10    os.makedirs(output_dir, exist_ok=True)11    shape = saga.grids[0].types.shape12    cursor = 013    for seg, seg_len in zip(saga.segments, saga.segment_lengths):14        seg_grids = saga.grids[cursor:cursor + seg_len]15        final = np.zeros(shape, dtype=np.int32)16        for g in seg_grids:17            final += g.types18        max_count = max(1, int(np.max(final)))19        primary, secondary = create_heatmap_colors(seg)20        gradient_colors = gradient(secondary, max_count)21        gradient_colors.insert(0, primary)22        gradient_rgba = [c.to_rgba() for c in gradient_colors]23        gradient_array = np.array(gradient_rgba, dtype=np.uint8)24        cumulative = np.zeros(shape, dtype=np.int32)25        for i, grid in enumerate(seg_grids):26            cumulative += grid.types27            colored = gradient_array[cumulative]28            cell = Cell2d(colors=colored)29            image = cell.to_image(1).convert("RGB")30            image.save(f"{output_dir}/{saga.key}_{cursor + i + 1:03d}.{FORMAT}", format=FORMAT)31        cursor += seg_len32    print(f"heatmap {len(saga.grids)} -> {output_dir}")33    return saga