Project_ZHUAN/PreProcessing_Common/scripts/make_clean_comparisons.py

from __future__ import annotations

import argparse
from pathlib import Path

import cv2
import numpy as np


IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp", ".webp"}


def read_gray(path: Path) -> np.ndarray:
    raw = np.fromfile(str(path), dtype=np.uint8)
    image = cv2.imdecode(raw, cv2.IMREAD_GRAYSCALE)
    if image is None:
        raise ValueError(f"Cannot read image: {path}")
    return image


def write_image(path: Path, image: np.ndarray) -> None:
    ok, encoded = cv2.imencode(path.suffix, image)
    if not ok:
        raise ValueError(f"Cannot encode image: {path}")
    encoded.tofile(str(path))


def fit_on_canvas(image: np.ndarray, width: int, height: int) -> np.ndarray:
    scale = min(width / image.shape[1], height / image.shape[0])
    resized_width = max(1, int(round(image.shape[1] * scale)))
    resized_height = max(1, int(round(image.shape[0] * scale)))
    resized = cv2.resize(
        image,
        (resized_width, resized_height),
        interpolation=cv2.INTER_AREA,
    )
    canvas = np.full((height, width), 255, dtype=np.uint8)
    x = (width - resized_width) // 2
    y = (height - resized_height) // 2
    canvas[y : y + resized_height, x : x + resized_width] = resized
    return canvas


def make_comparison(
    original_path: Path,
    cleaned_path: Path,
    output_path: Path,
    panel_size: int,
    label_height: int,
) -> None:
    original = read_gray(original_path)
    cleaned = read_gray(cleaned_path)
    panel_width = panel_size
    panel_height = panel_size

    original_panel = fit_on_canvas(original, panel_width, panel_height)
    cleaned_panel = fit_on_canvas(cleaned, panel_width, panel_height)
    divider = np.full((panel_height + label_height, 2), 210, dtype=np.uint8)
    comparison = np.full(
        (panel_height + label_height, panel_width * 2 + divider.shape[1]),
        255,
        dtype=np.uint8,
    )
    comparison[:panel_height, :panel_width] = original_panel
    comparison[:panel_height, panel_width + divider.shape[1] :] = cleaned_panel
    comparison[:, panel_width : panel_width + divider.shape[1]] = divider

    baseline = panel_height + 24
    cv2.putText(
        comparison,
        "original",
        (12, baseline),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.7,
        0,
        2,
        cv2.LINE_AA,
    )
    cv2.putText(
        comparison,
        "cleaned",
        (panel_width + divider.shape[1] + 12, baseline),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.7,
        0,
        2,
        cv2.LINE_AA,
    )
    cv2.putText(
        comparison,
        original_path.stem,
        (12, panel_height + label_height - 12),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.45,
        80,
        1,
        cv2.LINE_AA,
    )
    write_image(output_path, comparison)


def iter_originals(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 build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Create side-by-side original/cleaned comparison images."
    )
    parser.add_argument(
        "--original-dir",
        type=Path,
        default=Path("PreProcessing_Common/original_characters"),
    )
    parser.add_argument(
        "--cleaned-dir",
        type=Path,
        default=Path("PreProcessing_Common/cleaned_characters"),
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("PreProcessing_Common/comparisons"),
    )
    parser.add_argument("--panel-size", type=int, default=520)
    parser.add_argument("--label-height", type=int, default=72)
    return parser


def main() -> None:
    args = build_parser().parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)

    originals = iter_originals(args.original_dir)
    if not originals:
        raise SystemExit(f"No images found in {args.original_dir}")

    count = 0
    for original_path in originals:
        cleaned_path = args.cleaned_dir / f"{original_path.stem}_cleaned.png"
        if not cleaned_path.exists():
            print(f"missing cleaned image: {cleaned_path}")
            continue

        make_comparison(
            original_path=original_path,
            cleaned_path=cleaned_path,
            output_path=args.output_dir / f"{original_path.stem}_comparison.png",
            panel_size=args.panel_size,
            label_height=args.label_height,
        )
        print(f"comparison: {original_path.name}")
        count += 1

    if count == 0:
        raise SystemExit("No comparisons were written.")


if __name__ == "__main__":
    main()