painter.py

1import numpy as np2from typing import Dict, List, Optional, TYPE_CHECKING3from mrlypy.core.colors import alpha, black, blue, green, red, white4from mrlypy.core.enums import Mode5from mrlypy.core.state import state67if TYPE_CHECKING:8    from mrlypy.core.colors import Color9    from .models import Cell3d1011def get_mapping():12    return {0: [white], 1: [black], 2: [alpha], 3: [red], 4: [green], 5: [blue]}1314def paint_3d(cell: "Cell3d", mapping: Optional[Dict[int, List["Color"]]] = None, mode: Optional[Mode] = None) -> "Cell3d":15    mapping = mapping or get_mapping()16    mode = mode or Mode.TYPE17    for key, colors in mapping.items():18        mask = cell.types == key19        if not np.any(mask):20            continue21        match mode:22            case Mode.TYPE:23                cell.colors[mask] = colors[0].to_rgba()24            case Mode.RANDOM:25                palette = np.array([c.to_rgba() for c in colors], dtype=np.uint8)26                indices = state.rng.randint(0, len(palette), size=np.sum(mask))27                cell.colors[mask] = palette[indices]28            case Mode.ENUMERATE:29                palette = np.array([c.to_rgba() for c in colors], dtype=np.uint8)30                indices = np.arange(np.sum(mask)) % len(palette)31                cell.colors[mask] = palette[indices]32            case _:33                source_array = {34                    Mode.INDEX: np.arange(cell.width * cell.height * cell.depth).reshape(cell.depth, cell.height, cell.width),35                    Mode.TAG: cell.tags,36                    Mode.DEPTH: np.indices((cell.depth, cell.height, cell.width))[0],37                    Mode.ROW: np.indices((cell.depth, cell.height, cell.width))[1],38                    Mode.COLUMN: np.indices((cell.depth, cell.height, cell.width))[2],39                }.get(mode)40                if source_array is not None:41                    palette = np.array([c.to_rgba() for c in colors], dtype=np.uint8)42                    masked_values = source_array[mask]43                    indices = masked_values % len(palette)44                    cell.colors[mask] = palette[indices]45    return cell