scan4llm/tests/test_post_crop.py

from __future__ import annotations

import tempfile
from pathlib import Path

import cv2

from src.split.slicer import crop_to_content, find_content_bounds, process_image
from tests.helpers import create_test_image, image_to_path


class TestFindContentBounds:
    def test_uniform_image(self) -> None:
        image = create_test_image(200, 150)
        image[:, :] = (255, 255, 255)
        x1, y1, x2, y2 = find_content_bounds(image)
        assert (x1, y1, x2, y2) == (0, 0, 200, 150)

    def test_single_dot_center(self) -> None:
        image = create_test_image(200, 150)
        image[:, :] = (255, 255, 255)
        image[75, 100] = (0, 0, 0)
        x1, y1, x2, y2 = find_content_bounds(image)
        assert abs(x1 - 100) <= 1 and abs(x2 - 100) <= 1
        assert abs(y1 - 75) <= 1 and abs(y2 - 75) <= 1

    def test_rectangle_content(self) -> None:
        image = create_test_image(200, 150)
        image[:, :] = (255, 255, 255)
        image[30:100, 50:140] = (0, 0, 0)
        x1, y1, x2, y2 = find_content_bounds(image)
        assert x1 == 50
        assert y1 == 30
        assert x2 == 139
        assert y2 == 99

    def test_content_at_edges(self) -> None:
        image = create_test_image(200, 150)
        image[:, :] = (255, 255, 255)
        image[0:10, :] = (0, 0, 0)
        image[:, 0:10] = (0, 0, 0)
        x1, y1, _x2, _y2 = find_content_bounds(image)
        assert y1 == 0
        assert x1 == 0

    def test_grayscale_content(self) -> None:
        image = create_test_image(100, 100)
        image[:, :] = (255, 255, 255)
        image[20:30, 20:30] = (128, 128, 128)
        x1, y1, x2, y2 = find_content_bounds(image)
        assert x1 <= 29 and x2 >= 20
        assert y1 <= 29 and y2 >= 20


class TestCropToContent:
    def test_crop_centered_content(self) -> None:
        image = create_test_image(200, 150)
        image[:, :] = (255, 255, 255)
        image[30:50, 40:60] = (0, 0, 0)
        result = crop_to_content(image, border_px=10)
        assert result.shape[0] >= 20
        assert result.shape[1] >= 20

    def test_crop_uniform_image_no_change(self) -> None:
        image = create_test_image(100, 100)
        image[:, :] = (255, 255, 255)
        result = crop_to_content(image, border_px=10)
        assert result.shape == (100, 100, 3)

    def test_border_capped_by_image_edge(self) -> None:
        image = create_test_image(100, 100)
        image[:, :] = (255, 255, 255)
        image[10:30, 10:30] = (0, 0, 0)
        result = crop_to_content(image, border_px=200)
        assert result.shape == (100, 100, 3)

    def test_border_equals_distance_to_edge(self) -> None:
        image = create_test_image(100, 100)
        image[:, :] = (255, 255, 255)
        image[40:60, 30:70] = (0, 0, 0)
        result = crop_to_content(image, border_px=50)
        assert result.shape[1] <= 100
        assert result.shape[0] <= 100

    def test_dark_bg_no_crop_uniform(self) -> None:
        image = create_test_image(100, 100)
        image[:, :] = (0, 0, 0)
        result = crop_to_content(image, border_px=0)
        assert result.shape == (100, 100, 3)

    def test_zero_border_exact_crop(self) -> None:
        image = create_test_image(200, 150)
        image[:, :] = (255, 255, 255)
        image[20:100, 30:170] = (0, 0, 0)
        result = crop_to_content(image, border_px=0)
        assert result.shape == (80, 140, 3)


class TestProcessImagePostCrop:
    def test_post_crop_reduces_size(self) -> None:
        image = create_test_image(400, 400)
        image[:, :] = (255, 255, 255)
        image[100:200, 100:200] = (0, 0, 0)
        path = image_to_path(image, ".png")
        with tempfile.TemporaryDirectory() as tmp:
            result = process_image(Path(path), 1, 1, Path(tmp), border_px=30, post_crop=True)
            assert len(result) == 1
            loaded = cv2.imread(str(result[0]))
            assert loaded.shape[0] < 400 or loaded.shape[1] < 400

    def test_post_crop_off_no_crop(self) -> None:
        image = create_test_image(400, 300)
        path = image_to_path(image, ".png")
        with tempfile.TemporaryDirectory() as tmp:
            result = process_image(Path(path), 1, 1, Path(tmp), post_crop=False)
            assert len(result) == 1
            loaded = cv2.imread(str(result[0]))
            assert loaded.shape == (400, 500, 3)

    def test_post_crop_e2e_2x2(self) -> None:
        image = create_test_image(400, 400)
        image[:, :] = (255, 255, 255)
        image[50:150, 50:150] = (0, 0, 0)
        image[50:150, 250:350] = (0, 0, 0)
        image[250:350, 50:150] = (0, 0, 0)
        image[250:350, 250:350] = (0, 0, 0)
        path = image_to_path(image, ".png")
        with tempfile.TemporaryDirectory() as tmp:
            result_crop = process_image(Path(path), 2, 2, Path(tmp) / "crop", border_px=30, post_crop=True)
            result_no = process_image(Path(path), 2, 2, Path(tmp) / "no", border_px=30, post_crop=False)
            assert len(result_crop) == 4
            for cp, np in zip(result_crop, result_no):
                c = cv2.imread(str(cp))
                n = cv2.imread(str(np))
                assert c.shape[0] < n.shape[0]
                assert c.shape[1] < n.shape[1]