crop.py

1import math2from mrlypy.two import Cell2d3import numpy as np4from typing import List5from .helpers import logger67def crop_grids(grids: List[Cell2d]) -> List[Cell2d]:8    first_grid = grids[0]9    original_h, original_w = first_grid.types.shape10    heatmap = np.zeros_like(first_grid.types, dtype=np.int32)11    for grid in grids:12        heatmap += grid.types13    if not np.any(heatmap):14        logger.debug("Found nothing to crop (empty heatmap). Returning...")15        return grids16    rows = np.any(heatmap, axis=1)17    cols = np.any(heatmap, axis=0)18    rmin, rmax = np.where(rows)[0][[0, -1]]19    cmin, cmax = np.where(cols)[0][[0, -1]]20    rows_cropped = rmin + (original_h - 1 - rmax)21    cols_cropped = cmin + (original_w - 1 - cmax)22    if rows_cropped == 0 and cols_cropped == 0:23        logger.debug("Found nothing to crop. Returning...")24        return grids25    else:26        logger.debug(f"Cropping {rows_cropped} rows / {cols_cropped} columns")27    height = rmax - rmin + 128    width = cmax - cmin + 129    final_size = max(height, width)30    pad_h = final_size - height31    pad_w = final_size - width32    pad_top = pad_h // 233    pad_bottom = pad_h - pad_top34    pad_left = pad_w // 235    pad_right = pad_w - pad_left36    trimmed_grids = []37    for grid in grids:38        cropped_grid = grid.types[rmin:rmax + 1, cmin:cmax + 1]39        padded_cropped_grid = np.pad(40            cropped_grid,41            ((pad_top, pad_bottom), (pad_left, pad_right)),42            mode="constant",43            constant_values=0,44        )45        trimmed_grids.append(Cell2d(types=padded_cropped_grid))46    return trimmed_grids4748def tessellate_grids(grids: List[Cell2d], min_canvas: int) -> List[Cell2d]:49    if not grids or min_canvas <= 0:50        return grids51    size = max(grids[0].types.shape)52    if size >= min_canvas:53        return grids54    n = math.ceil(min_canvas / size)55    logger.debug(f"Tessellating {size} -> {size * n} (x{n})")56    return [grid.copy().tile(n, n) for grid in grids]