scan4llm/tests/test_detect_grid.py

from __future__ import annotations

from src.split.slicer import detect_grid
from tests.helpers import create_test_image


class TestDetectGrid:
    def test_2x2(self) -> None:
        image = create_test_image(400, 300)
        image[:, :] = (255, 255, 255)
        dark = (0, 0, 0)
        gap = 20
        bw = (400 - 3 * gap) // 2
        bh = (300 - 3 * gap) // 2
        for r in range(2):
            for c in range(2):
                y = gap + r * (bh + gap)
                x = gap + c * (bw + gap)
                image[y : y + bh, x : x + bw] = dark
        cols, rows = detect_grid(image)
        assert cols == 2
        assert rows == 2

    def test_2x4(self) -> None:
        image = create_test_image(800, 400)
        image[:, :] = (255, 255, 255)
        gap = 15
        bw = (800 - 5 * gap) // 4
        bh = (400 - 3 * gap) // 2
        dark = (0, 0, 0)
        for r in range(2):
            for c in range(4):
                y = gap + r * (bh + gap)
                x = gap + c * (bw + gap)
                image[y : y + bh, x : x + bw] = dark
        cols, rows = detect_grid(image)
        assert cols == 4
        assert rows == 2

    def test_1x1(self) -> None:
        image = create_test_image(200, 200)
        image[:, :] = (255, 255, 255)
        image[20:180, 20:180] = (0, 0, 0)
        cols, rows = detect_grid(image)
        assert cols == 1
        assert rows == 1

    def test_3x3(self) -> None:
        image = create_test_image(600, 450)
        image[:, :] = (255, 255, 255)
        dark = (0, 0, 0)
        gap = 10
        bw = (600 - 4 * gap) // 3
        bh = (450 - 4 * gap) // 3
        for r in range(3):
            for c in range(3):
                y = gap + r * (bh + gap)
                x = gap + c * (bw + gap)
                image[y : y + bh, x : x + bw] = dark
        cols, rows = detect_grid(image)
        assert cols == 3
        assert rows == 3

    def test_uniform_image(self) -> None:
        image = create_test_image()
        cols, rows = detect_grid(image)
        assert cols == 1
        assert rows == 1