scan4llm/src/image_prepare.py

from __future__ import annotations

import logging
import os
from pathlib import Path
from typing import Annotated

import typer

from src.cli_utils import run_main
from src.split.slicer import (
    DEFAULT_BORDER_PX,
    VALID_ROTATIONS,
    denoise_image,
    detect_grid,
    load_image,
    parse_slice,
    process_image,
    rotate_image,
    unwarp_image,
)

logging.basicConfig(
    level=logging.INFO,
    format="%(levelname)-8s %(message)s",
)
app = typer.Typer(add_completion=False, no_args_is_help=True)


def _validate_input(input_path: Path | None) -> Path:
    if input_path is None:
        raise typer.BadParameter(
            "Укажите путь к входному изображению", param_hint="--input"
        )
    if not input_path.exists():
        raise typer.BadParameter(
            f"Файл не найден: {input_path}", param_hint="--input"
        )
    if not input_path.is_file():
        raise typer.BadParameter(
            f"Не является файлом: {input_path}", param_hint="--input"
        )
    if not os.access(input_path, os.R_OK):
        raise typer.BadParameter(
            f"Нет доступа на чтение: {input_path}", param_hint="--input"
        )
    return input_path.resolve()


def _validate_rotation(degrees: int | None) -> None:
    if degrees is not None and degrees not in VALID_ROTATIONS:
        raise typer.BadParameter(
            f"Недопустимый угол: {degrees}. "
            f"Допустимые: {sorted(VALID_ROTATIONS)}"
        )


def _resolve_grid(
    slice_str: str | None,
    slice_auto: bool,
    input_path: Path,
    pre_rotate: int | None,
    unwarp: bool,
    denoise: bool,
) -> tuple[int, int]:
    if slice_str is not None and slice_auto:
        raise typer.BadParameter(
            "Нельзя указывать одновременно --slice и --slice-auto"
        )
    if slice_str is None and not slice_auto:
        raise typer.BadParameter("Укажите --slice или --slice-auto")

    if slice_auto:
        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)
        cols, rows = detect_grid(image)
        typer.echo(f"Определена сетка: {cols}×{rows}")
        return cols, rows

    if slice_str is None:
        raise typer.BadParameter("Укажите --slice или --slice-auto")
    return parse_slice(slice_str)


@app.command()
def slice_pages(
    input: Annotated[
        Path | None,
        typer.Option(
            "--input",
            "-i",
            file_okay=True,
            dir_okay=False,
            help="Путь к входному изображению (JPEG, PNG, TIFF)",
        ),
    ] = None,
    slice: Annotated[
        str | None,
        typer.Option(
            "--slice",
            "-s",
            help="Размер сетки в формате <колонки>:<строки>, например 3:2",
        ),
    ] = None,
    slice_auto: Annotated[
        bool,
        typer.Option(
            "--slice-auto",
            "-a",
            help="Автоматически определить размер сетки по содержимому изображения",
        ),
    ] = False,
    output_dir: Annotated[
        Path,
        typer.Option(
            "--output-dir",
            "-o",
            file_okay=False,
            dir_okay=True,
            writable=True,
            help="Каталог для сохранения результатов (по умолчанию — текущий)",
        ),
    ] = Path(),
    pre_rotate: Annotated[
        int | None,
        typer.Option(
            "--pre-rotate",
            "-r",
            help="Поворот изображения перед разрезанием: 90, 180 или 270",
        ),
    ] = None,
    border: Annotated[
        int,
        typer.Option(
            "--border",
            "-b",
            min=0,
            help="Отступ в пикселях (белая рамка / отступ при обрезке содержимого, по умолчанию 5)",
        ),
    ] = DEFAULT_BORDER_PX,
    unwarp: Annotated[
        bool,
        typer.Option(
            "--unwarp",
            "-u",
            help="Выпрямить перспективные искажения перед разрезанием",
        ),
    ] = False,
    denoise: Annotated[
        bool,
        typer.Option(
            "--denoise",
            help="Убрать шум сканера (Non-Local Means) перед разрезанием",
        ),
    ] = False,
    post_crop: Annotated[
        bool,
        typer.Option(
            "--post-crop / --no-post-crop",
            "-c",
            help="Обрезать каждую страницу по границам полезного содержимого (включено по умолчанию)",
        ),
    ] = True,
) -> None:
    """Подготовить отсканированное изображение: разрезать по сетке, обрезать, добавить рамку."""
    input_path = _validate_input(input)
    _validate_rotation(pre_rotate)
    cols, rows = _resolve_grid(slice, slice_auto, input_path, pre_rotate, unwarp, denoise)

    output_paths = process_image(
        input_path=input_path,
        rows=rows,
        cols=cols,
        output_dir=output_dir.resolve(),
        pre_rotate=pre_rotate,
        border_px=border,
        post_crop=post_crop,
        unwarp=unwarp,
        denoise=denoise,
    )

    typer.echo(f"Сохранено {len(output_paths)} страниц в {output_dir.resolve()}")


def main() -> None:
    run_main(app)


if __name__ == "__main__":
    main()