preview_preprocessing.py

import argparse
import random
from pathlib import Path

import cv2
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image

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


def parse_args():
    parser = argparse.ArgumentParser(
        description="Preview Original vs CLAHE vs Sobel vs CLAHE+Sobel on sampled images."
    )
    parser.add_argument("--data-root", type=str, default="data_frames", help="Root folder containing images.")
    parser.add_argument(
        "--num-samples",
        type=int,
        default=6,
        help="Number of images to preview.",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed for sampling images.",
    )
    parser.add_argument(
        "--clahe-clip-limit",
        type=float,
        default=4,
        help="CLAHE clip limit.",
    )
    parser.add_argument(
        "--clahe-tile-grid-size",
        type=int,
        default=8,
        help="CLAHE tile grid size (single integer for square grid).",
    )
    parser.add_argument(
        "--sobel-ksize",
        type=int,
        default=3,
        help="Sobel kernel size (odd positive integer).",
    )
    parser.add_argument(
        "--sobel-mode",
        type=str,
        default="magnitude",
        choices=["magnitude", "x", "y"],
        help="Sobel output mode.",
    )
    parser.add_argument(
        "--output",
        type=str,
        default="preprocess_preview.png",
        help="Output PNG path.",
    )
    return parser.parse_args()


def collect_images(data_root: Path):
    images = [p for p in data_root.rglob("*") if p.is_file() and p.suffix.lower() in VALID_EXTS]
    images.sort()
    return images


def apply_clahe(gray: np.ndarray, clip_limit: float, tile_grid_size: int) -> np.ndarray:
    clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=(tile_grid_size, tile_grid_size))
    return clahe.apply(gray)


def apply_sobel(gray: np.ndarray, ksize: int, mode: str) -> np.ndarray:
    sobel_x = cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=ksize)
    sobel_y = cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=ksize)

    if mode == "x":
        sobel = np.abs(sobel_x)
    elif mode == "y":
        sobel = np.abs(sobel_y)
    else:
        sobel = cv2.magnitude(sobel_x, sobel_y)

    return cv2.convertScaleAbs(sobel)


def to_rgb(gray: np.ndarray) -> np.ndarray:
    return cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)


def build_preview_variants(
    img_path: Path,
    clahe_clip_limit: float,
    clahe_tile_grid_size: int,
    sobel_ksize: int,
    sobel_mode: str,
):
    rgb = np.array(Image.open(img_path).convert("RGB"))
    gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)

    clahe = apply_clahe(gray, clip_limit=clahe_clip_limit, tile_grid_size=clahe_tile_grid_size)
    sobel = apply_sobel(gray, ksize=sobel_ksize, mode=sobel_mode)
    clahe_sobel = apply_sobel(clahe, ksize=sobel_ksize, mode=sobel_mode)

    return {
        "Original": rgb,
        "CLAHE": to_rgb(clahe),
        f"Sobel-{sobel_mode}": to_rgb(sobel),
        f"CLAHE+Sobel-{sobel_mode}": to_rgb(clahe_sobel),
    }


def main():
    args = parse_args()

    if args.num_samples <= 0:
        raise ValueError("--num-samples must be > 0")
    if args.sobel_ksize <= 0 or args.sobel_ksize % 2 == 0:
        raise ValueError("--sobel-ksize must be an odd positive integer, e.g. 1, 3, 5")
    if args.clahe_tile_grid_size <= 0:
        raise ValueError("--clahe-tile-grid-size must be > 0")

    data_root = Path(args.data_root)
    if not data_root.exists():
        raise FileNotFoundError(
            f"Data root not found: {data_root}. Use --data-root to point to your image folder."
        )

    all_images = collect_images(data_root)
    if len(all_images) == 0:
        raise RuntimeError(
            f"No images found under {data_root}. Expected extensions: {sorted(VALID_EXTS)}"
        )

    rng = random.Random(args.seed)
    sample_count = min(args.num_samples, len(all_images))
    sampled = rng.sample(all_images, k=sample_count)

    col_titles = ["Original", "CLAHE", f"Sobel-{args.sobel_mode}", f"CLAHE+Sobel-{args.sobel_mode}"]
    nrows, ncols = sample_count, len(col_titles)

    fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=(4.0 * ncols, 3.2 * nrows))
    if nrows == 1:
        axes = np.expand_dims(axes, axis=0)

    for row_idx, img_path in enumerate(sampled):
        variants = build_preview_variants(
            img_path=img_path,
            clahe_clip_limit=args.clahe_clip_limit,
            clahe_tile_grid_size=args.clahe_tile_grid_size,
            sobel_ksize=args.sobel_ksize,
            sobel_mode=args.sobel_mode,
        )

        for col_idx, title in enumerate(col_titles):
            ax = axes[row_idx][col_idx]
            ax.imshow(variants[title])
            if row_idx == 0:
                ax.set_title(title)
            if col_idx == 0:
                rel_path = img_path.relative_to(data_root)
                ax.set_ylabel(str(rel_path), fontsize=8)
            ax.set_xticks([])
            ax.set_yticks([])

    fig.suptitle(
        "Image Preprocessing Preview\n"
        f"root={data_root} | samples={sample_count} | "
        f"clahe(clip={args.clahe_clip_limit}, tile={args.clahe_tile_grid_size}) | "
        f"sobel(ksize={args.sobel_ksize}, mode={args.sobel_mode})",
        fontsize=11,
    )
    fig.tight_layout(rect=[0, 0, 1, 0.95])

    out_path = Path(args.output)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path, dpi=180)
    plt.close(fig)

    print(f"Found images: {len(all_images)}")
    print(f"Sampled images: {sample_count}")
    print(f"Saved preview: {out_path}")


if __name__ == "__main__":
    main()