Project_ZHUAN/PreProcessing_MatrixZhuan/scripts/run_preprocessing.py

from __future__ import annotations

import argparse
from pathlib import Path

import cv2
import numpy as np

from process_characters import (
    IMAGE_SUFFIXES,
    binarize_foreground,
    clean_mask,
    matrix_map,
    mask_to_display,
    read_image,
    small_matrix_to_svg,
    write_image,
)


def matrix_to_text(matrix: np.ndarray) -> str:
    return "\n".join(
        "".join("1" if value > 0 else "0" for value in row) for row in matrix
    ) + "\n"


def iter_image_paths(input_dir: Path) -> list[Path]:
    return sorted(
        path
        for path in input_dir.iterdir()
        if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES
    )


def process_image(
    image_path: Path,
    cleaned_dir: Path,
    matrix_dir: Path,
    svg_dir: Path,
    grid_size: int = 17,
    cell_samples: int = 32,
    open_kernel_size: int = 1,
    close_kernel_size: int = 0,
    median_size: int = 1,
    min_component_area: int = 8,
    triangle_threshold: float = 0.18,
) -> None:
    original = read_image(image_path)
    gray = cv2.cvtColor(original, cv2.COLOR_BGR2GRAY)
    binary = binarize_foreground(gray)
    cleaned = clean_mask(
        binary,
        open_kernel_size=open_kernel_size,
        close_kernel_size=close_kernel_size,
        median_size=median_size,
        min_component_area=min_component_area,
        grid_size=grid_size,
    )
    mapping = matrix_map(
        cleaned,
        grid_size=grid_size,
        triangle_threshold=triangle_threshold,
        cell_samples=cell_samples,
    )

    stem = image_path.stem
    write_image(cleaned_dir / f"{stem}_cleaned.png", mask_to_display(cleaned))
    (matrix_dir / f"{stem}_matrix.txt").write_text(
        matrix_to_text(mapping.matrix),
        encoding="utf-8",
    )
    height, width = gray.shape
    svg = small_matrix_to_svg(
        mapping.matrix,
        width,
        height,
        stem,
        mapping.x_edges,
        mapping.y_edges,
        mapping.triangles,
    )
    (svg_dir / f"{stem}.svg").write_text(svg, encoding="utf-8")


def process_directory(
    input_dir: Path,
    cleaned_dir: Path,
    matrix_dir: Path,
    svg_dir: Path,
    grid_size: int = 17,
    cell_samples: int = 32,
    open_kernel_size: int = 1,
    close_kernel_size: int = 0,
    median_size: int = 1,
    min_component_area: int = 8,
    triangle_threshold: float = 0.18,
) -> int:
    cleaned_dir.mkdir(parents=True, exist_ok=True)
    matrix_dir.mkdir(parents=True, exist_ok=True)
    svg_dir.mkdir(parents=True, exist_ok=True)

    image_paths = iter_image_paths(input_dir)
    if not image_paths:
        raise SystemExit(f"No images found in {input_dir}")

    for image_path in image_paths:
        process_image(
            image_path=image_path,
            cleaned_dir=cleaned_dir,
            matrix_dir=matrix_dir,
            svg_dir=svg_dir,
            grid_size=grid_size,
            cell_samples=cell_samples,
            open_kernel_size=open_kernel_size,
            close_kernel_size=close_kernel_size,
            median_size=median_size,
            min_component_area=min_component_area,
            triangle_threshold=triangle_threshold,
        )
        print(f"processed: {image_path.name}")

    return len(image_paths)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Clean seal-script images and export 17x17 matrix text files."
    )
    parser.add_argument("--input-dir", type=Path, required=True)
    parser.add_argument("--cleaned-dir", type=Path, required=True)
    parser.add_argument("--matrix-dir", type=Path, required=True)
    parser.add_argument("--svg-dir", type=Path, required=True)
    parser.add_argument("--grid-size", type=int, default=17)
    parser.add_argument("--cell-samples", type=int, default=32)
    parser.add_argument("--open-kernel-size", type=int, default=1)
    parser.add_argument("--close-kernel-size", type=int, default=0)
    parser.add_argument("--median-size", type=int, default=1)
    parser.add_argument("--min-component-area", type=int, default=8)
    parser.add_argument("--triangle-threshold", type=float, default=0.18)
    return parser


def main() -> None:
    args = build_parser().parse_args()
    process_directory(
        input_dir=args.input_dir,
        cleaned_dir=args.cleaned_dir,
        matrix_dir=args.matrix_dir,
        svg_dir=args.svg_dir,
        grid_size=args.grid_size,
        cell_samples=args.cell_samples,
        open_kernel_size=args.open_kernel_size,
        close_kernel_size=args.close_kernel_size,
        median_size=args.median_size,
        min_component_area=args.min_component_area,
        triangle_threshold=args.triangle_threshold,
    )


if __name__ == "__main__":
    main()