scan4llm/src/split/slicer.py

from __future__ import annotations

import logging
from pathlib import Path

import cv2
import numpy as np

logger = logging.getLogger(__name__)

SUPPORTED_EXTENSIONS: frozenset[str] = frozenset({".jpg", ".jpeg", ".png", ".tiff", ".tif"})
PAGE_NUMBER_WIDTH: int = 2
DEFAULT_BORDER_PX: int = 5
BORDER_COLOR: tuple[int, int, int] = (255, 255, 255)
VALID_ROTATIONS: frozenset[int] = frozenset({90, 180, 270})
DENSITY_RATIO: float = 0.005
AUTO_GRID_MAX: int = 4
EMPTY_DARK_THRESHOLD: int = 100
EMPTY_DARK_RATIO: float = 0.001


def load_image(path: Path) -> np.ndarray:
    if path.suffix.lower() not in SUPPORTED_EXTENSIONS:
        raise ValueError(
            f"Неподдерживаемый формат: {path.suffix}. "
            f"Поддерживаются: {', '.join(sorted(SUPPORTED_EXTENSIONS))}"
        )
    if not path.exists():
        raise FileNotFoundError(f"Файл не найден: {path}")
    image = cv2.imread(str(path), cv2.IMREAD_UNCHANGED)
    if image is None:
        raise ValueError(f"Не удалось загрузить изображение: {path}")
    return image


def rotate_image(image: np.ndarray, degrees: int) -> np.ndarray:
    if degrees == 90:
        return cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)
    if degrees == 180:
        return cv2.rotate(image, cv2.ROTATE_180)
    if degrees == 270:
        return cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)
    raise ValueError(f"Недопустимый угол поворота: {degrees}. Допустимые: {sorted(VALID_ROTATIONS)}")


def slice_grid(image: np.ndarray, rows: int, cols: int) -> list[np.ndarray]:
    height, width = image.shape[:2]
    base_h = height // rows
    base_w = width // cols
    extra_h = height % rows
    extra_w = width % cols

    pages: list[np.ndarray] = []
    y = 0
    for row in range(rows):
        cell_h = base_h + (1 if row >= rows - extra_h else 0)
        x = 0
        for col in range(cols):
            cell_w = base_w + (1 if col >= cols - extra_w else 0)
            page = image[y : y + cell_h, x : x + cell_w]
            pages.append(page)
            x += cell_w
        y += cell_h
    return pages


def add_border(image: np.ndarray, border_px: int) -> np.ndarray:
    return cv2.copyMakeBorder(
        image,
        top=border_px,
        bottom=border_px,
        left=border_px,
        right=border_px,
        borderType=cv2.BORDER_CONSTANT,
        value=BORDER_COLOR,
    )


def save_page(image: np.ndarray, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    success = cv2.imwrite(str(path), image)
    if not success:
        raise OSError(f"Не удалось сохранить изображение: {path}")


def generate_output_paths(input_path: Path, page_count: int, output_dir: Path) -> list[Path]:
    stem = input_path.stem
    suffix = input_path.suffix.lower()
    if page_count == 1:
        return [output_dir / f"{stem}{suffix}"]
    max_digits = max(PAGE_NUMBER_WIDTH, len(str(page_count)))
    return [
        output_dir / f"{stem}_{i:0{max_digits}d}{suffix}"
        for i in range(1, page_count + 1)
    ]


def parse_slice(value: str) -> tuple[int, int]:
    parts = value.split(":")
    if len(parts) != 2:
        raise ValueError(
            f"Неверный формат --slice: {value}. Ожидается <колонки>:<строки>, например 3:2"
        )
    try:
        cols = int(parts[0])
        rows = int(parts[1])
    except ValueError:
        raise ValueError(
            f"Неверный формат --slice: {value}. Колонки и строки должны быть целыми числами"
        )
    if rows < 1 or cols < 1:
        raise ValueError(
            f"Неверное значение --slice: {value}. Строки и столбцы должны быть >= 1"
        )
    return cols, rows


def is_empty(image: np.ndarray) -> bool:
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    dark_pixels = (gray < EMPTY_DARK_THRESHOLD).sum()
    threshold = int(image.shape[0] * image.shape[1] * EMPTY_DARK_RATIO)
    return dark_pixels <= threshold


def _binarize(image: np.ndarray) -> np.ndarray:
    """Бинаризация: серый → Гаусс-блюр → Otsu-порог → морф.закрытие.
    При слишком низком пороге Otsu (<50) — адаптивный порог (Gaussian, окно 31).
    Морфологическое закрытие (3×3) склеивает фрагменты букв в непрерывные регионы."""
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(gray, (3, 3), 0)
    otsu_th, binary = cv2.threshold(
        blurred, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU,
    )
    if otsu_th < 50:
        binary = cv2.adaptiveThreshold(
            blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
            cv2.THRESH_BINARY_INV, 31, 10,
        )
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    return cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)


def detect_grid(image: np.ndarray) -> tuple[int, int]:
    """Автоопределение сетки: бинаризация → проекция плотности → дилатация → подсчёт полос.
    Дилатация (ядро = max(5, длина/40)) сливает строки текста в непрерывные блоки,
    но не перекрывает межстраничные пробелы (они шире)."""
    binary = _binarize(image)

    cols = _count_bands(binary, axis=0)
    rows = _count_bands(binary, axis=1)

    return cols, rows


def _count_bands(binary: np.ndarray, axis: int) -> int:
    other_dim = binary.shape[axis]
    length = binary.shape[1 - axis]
    projection = (binary > 0).sum(axis=axis).astype(np.float64)
    density = projection / other_dim

    dilate_width = max(5, length // 80)
    kernel_1d = np.ones(dilate_width, dtype=np.uint8)
    density_u8 = (density * 255).clip(0, 255).astype(np.uint8)
    density_2d = density_u8.reshape(1, -1)
    dilated = cv2.dilate(density_2d, kernel_1d)
    is_content = (dilated[0] > 1).tolist()

    bands = 0
    in_band = False
    for val in is_content:
        if val and not in_band:
            bands += 1
            in_band = True
        elif not val:
            in_band = False
    return max(1, bands)


def find_content_bounds(image: np.ndarray) -> tuple[int, int, int, int]:
    binary = _binarize(image)

    h, w = binary.shape
    min_density = int(w * DENSITY_RATIO)

    row_density = (binary > 0).sum(axis=1)
    col_density = (binary > 0).sum(axis=0)

    y1 = 0
    while y1 < h and row_density[y1] < min_density:
        y1 += 1
    y2 = h - 1
    while y2 >= 0 and row_density[y2] < min_density:
        y2 -= 1
    x1 = 0
    while x1 < w and col_density[x1] < min_density:
        x1 += 1
    x2 = w - 1
    while x2 >= 0 and col_density[x2] < min_density:
        x2 -= 1

    if y1 > y2 or x1 > x2:
        return 0, 0, w, h
    return x1, y1, x2, y2


def crop_to_content(image: np.ndarray, border_px: int = DEFAULT_BORDER_PX) -> np.ndarray:
    h, w = image.shape[:2]
    x1, y1, x2, y2 = find_content_bounds(image)

    left = min(border_px, x1)
    top = min(border_px, y1)
    right = min(border_px, w - 1 - x2)
    bottom = min(border_px, h - 1 - y2)

    x1_crop = x1 - left
    y1_crop = y1 - top
    x2_crop = x2 + right + 1
    y2_crop = y2 + bottom + 1

    return image[y1_crop:y2_crop, x1_crop:x2_crop]


def _order_points(pts: np.ndarray) -> np.ndarray:
    rect = np.zeros((4, 2), dtype=np.float32)
    s = pts.sum(axis=1)
    rect[0] = pts[np.argmin(s)]
    rect[2] = pts[np.argmax(s)]
    diff = np.diff(pts, axis=1)
    rect[1] = pts[np.argmin(diff)]
    rect[3] = pts[np.argmax(diff)]
    return rect


def _apply_perspective(image: np.ndarray, rect: np.ndarray) -> np.ndarray:
    (tl, tr, br, bl) = rect
    max_w = int(max(np.linalg.norm(br - bl), np.linalg.norm(tr - tl)))
    max_h = int(max(np.linalg.norm(tr - br), np.linalg.norm(tl - bl)))
    dst = np.array([[0, 0], [max_w - 1, 0], [max_w - 1, max_h - 1], [0, max_h - 1]], dtype=np.float32)
    mtx = cv2.getPerspectiveTransform(rect, dst)
    return cv2.warpPerspective(image, mtx, (max_w, max_h), flags=cv2.INTER_CUBIC)


def unwarp_image(image: np.ndarray) -> np.ndarray:
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    edges = cv2.Canny(blurred, 50, 150)
    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return image

    h, w = image.shape[:2]
    min_area = w * h * 0.3

    contours = sorted(contours, key=cv2.contourArea, reverse=True)
    for cnt in contours[:10]:
        area = cv2.contourArea(cnt)
        if area < min_area:
            continue
        peri = cv2.arcLength(cnt, True)
        approx = cv2.approxPolyDP(cnt, 0.02 * peri, True)
        if len(approx) == 4:
            rect = _order_points(approx.reshape(4, 2))
            return _apply_perspective(image, rect)
    return image


def denoise_image(image: np.ndarray) -> np.ndarray:
    """Non-Local Means denoising. Убирает шум, сохраняя границы символов."""
    return cv2.fastNlMeansDenoisingColored(image, None, 10, 10, 7, 21)


def process_image(
    input_path: Path,
    rows: int | None = None,
    cols: int | None = None,
    output_dir: Path = Path(),
    pre_rotate: int | None = None,
    border_px: int = DEFAULT_BORDER_PX,
    post_crop: bool = False,
    unwarp: bool = False,
    denoise: bool = False,
) -> list[Path]:
    image = load_image(input_path)

    if pre_rotate is not None:
        image = rotate_image(image, pre_rotate)

    if unwarp:
        image = unwarp_image(image)

    if denoise:
        image = denoise_image(image)

    if rows is None or cols is None:
        cols, rows = detect_grid(image)

    pages = slice_grid(image, rows, cols)

    if post_crop:
        pages = [crop_to_content(page, border_px) for page in pages]

    pages_with_border = [add_border(page, border_px) for page in pages]

    non_empty: list[np.ndarray] = []
    skipped: int = 0
    for i, page in enumerate(pages_with_border):
        if is_empty(page):
            logger.info("Пропущена пустая страница %d (размер %dx%d)", i + 1, *page.shape[:2])
            skipped += 1
        else:
            non_empty.append(page)

    output_paths = generate_output_paths(input_path, len(non_empty), output_dir)
    for page, path in zip(non_empty, output_paths):
        save_page(page, path)

    return output_paths