primer.py

1.8 kB · python · 65 lines

1import numpy as np2from mrlypy.tile.models import Tile3from mrlypy.two import carpet_2d4from .enums import Edition, Scheme, Type5from .models import Paint6from .randomizer import random_secondary, random_shades78# HELPERS910def remap_tags(p: Paint, tile: Tile) -> int:11    match p.target:12        case Type.FILL:13            type_mask = (tile.cell.types == 0)14        case Type.VOID:15            type_mask = (tile.cell.types == 1)16    relevant_tags = tile.cell.tags[type_mask]17    if relevant_tags.size == 0:18        return 019    unique_tags = np.unique(relevant_tags)20    new_tags = np.zeros_like(tile.cell.tags)21    for i, t in enumerate(unique_tags):22        new_tags[tile.cell.tags == t] = i23    tile.cell.tags = new_tags24    return len(unique_tags)2526def apply_colors(p: Paint, max_val: int) -> Paint:27    match p.scheme:28        case Scheme.MULTICOLOR:29            p.wipe()30            p.secondary = random_secondary(max_val, p.primary)31        case Scheme.MULTITONE:32            p.wipe()33            p.secondary = random_secondary(1, None)34            p.shades = random_shades(max_val, p.primary)35    return p3637# LAYERS3839def prime_layers(p: Paint, tile: Tile, mask: Tile = None) -> Paint:40    tile.cell.layers()41    max_val = remap_tags(p, tile)42    p = apply_colors(p, max_val)43    return p4445# NEIGHBORS4647def prime_neighbors(p: Paint, tile: Tile, mask: Tile = None) -> Paint:48    tile.cell.neighbors(mask.cell.types, 1)49    max_val = remap_tags(p, tile)50    p = apply_colors(p, max_val)51    return p5253# PRIME5455PRIME_FACTORY = {56    Edition.LAYERS: prime_layers,57    Edition.NEIGHBORS: prime_neighbors,58}5960def prime(p: Paint, tile: Tile, mask: Tile = None) -> Paint:61    if p.edition not in PRIME_FACTORY:62        return p63    if p.edition == Edition.NEIGHBORS and mask is None:64        mask = Tile(cell=carpet_2d(3))65    return PRIME_FACTORY[p.edition](p, tile, mask)