utils/split_image.py

import numpy as np
from typing import Literal

def split_image(image: np.ndarray, width_parts: int, height_parts: int, result: Literal['np.ndarray', 'bboxes'] = 'np.ndarray'):
  tile_width = image.shape[0] // width_parts
  tile_height = image.shape[1] // height_parts
  for height in range(height_parts):
    for width in range(width_parts):
      width_start = width * tile_width
      width_end = tile_width * (width + 1) if (width + 1) < width_parts else image.shape[0]
      height_start = height * tile_height
      height_end = tile_height * (height + 1) if (height + 1) < height_parts else image.shape[1]
      if result == 'np.ndarray':
        # np.ndarray(height, width, channels)
        yield image[height_start:height_end, width_start:width_end]
      else:
        # Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and ``0 <= y1 < y2``.
        yield (width_start, height_start, width_end, height_end)