preprocess.py

import cv2
import numpy as np
import subprocess
import argparse
import os
import tempfile
from pathlib import Path


def enhance_imagemagick(image_path: str, output_path: str) -> bool:
    cmd = [
        'convert', image_path,
        '-resize', '50%',
        '-colorspace', 'gray',
        '-blur', '0x0.5',
        '-normalize',
        '-lat', '15x15-8%',
        '-threshold', '45%',
        '-morphology', 'Open', 'Disk:0.4',
        '-morphology', 'Dilate', 'Disk:0.4',
        output_path
    ]
    try:
        result = subprocess.run(cmd, capture_output=True, timeout=120)
        return result.returncode == 0
    except Exception as e:
        print(f"  ImageMagick error: {e}")
        return False


def filter_connected_components(image: np.ndarray,
                                 min_area: int = 30,
                                 max_area_ratio: float = 0.01,
                                 dilate: bool = True) -> np.ndarray:
    if len(image.shape) == 3:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    else:
        gray = image.copy()

    _, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY_INV)

    h, w = gray.shape
    max_area = int(h * w * max_area_ratio)

    num, labels, stats, _ = cv2.connectedComponentsWithStats(binary)

    clean = np.ones_like(gray) * 255
    kept = 0
    for i in range(1, num):
        area = stats[i, cv2.CC_STAT_AREA]
        if min_area < area < max_area:
            clean[labels == i] = 0
            kept += 1

    if dilate:
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2))
        clean = cv2.erode(clean, kernel, iterations=1)

    print(f"  CC filter: {kept}/{num-1} components kept "
          f"(min={min_area}, max={max_area})")

    return cv2.cvtColor(clean, cv2.COLOR_GRAY2BGR)


def analyze_image(image_path: str) -> dict:
    img = cv2.imread(image_path)
    if img is None:
        return {'needs_enhance': False}

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    mean = float(gray.mean())
    std = float(gray.std())
    noise = float(cv2.Laplacian(gray, cv2.CV_64F).var())

    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    saturation = float(hsv[:,:,1].mean())

    is_colored_bg = saturation > 30 and mean < 220

    needs_enhance = (
        mean < 190 or
        std < 30 or
        is_colored_bg
    )

    return {
        'needs_enhance': needs_enhance,
        'mean': round(mean, 1),
        'std': round(std, 1),
        'noise': round(noise, 1),
        'is_colored_bg': is_colored_bg,
        'saturation': round(saturation, 1),
    }


def preprocess(image_path: str, output_path: str,
               min_area: int = 30,
               max_area_ratio: float = 0.01,
               dilate: bool = True,
               auto: bool = True,
               keep_temp: bool = False) -> bool:
    print(f"\nProcessing: {image_path}")

    os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)

    if auto:
        info = analyze_image(image_path)
        print(f"  [ANALYZE] mean={info['mean']}, std={info['std']}, "
              f"noise={info['noise']:.0f}, colored_bg={info['is_colored_bg']}")

        if not info.get('needs_enhance', True):
            print(f"  [SKIP] Image is clean, no enhancement needed")
            import shutil
            shutil.copy2(image_path, output_path)
            return True
        else:
            print(f"  [ENHANCE] Image needs enhancement")

    os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)

    tmp = tempfile.NamedTemporaryFile(suffix='.jpg', delete=False)
    tmp_path = tmp.name
    tmp.close()

    print(f"  [1] ImageMagick LAT...")
    ok = enhance_imagemagick(image_path, tmp_path)
    if not ok:
        print("  ImageMagick failed! Check if installed: brew install imagemagick")
        return False

    print(f"  [2] Connected Components filtering...")
    img = cv2.imread(tmp_path)
    if img is None:
        print("  Cannot read enhanced image!")
        return False

    result = filter_connected_components(img, min_area, max_area_ratio, dilate)

    cv2.imwrite(output_path, result)
    print(f"  Saved: {output_path}")

    if not keep_temp:
        os.remove(tmp_path)

    return True


def preprocess_dir(input_dir: str, output_dir: str,
                   extensions: set = None, **kwargs) -> int:
    if extensions is None:
        extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}

    paths = sorted(
        p for p in Path(input_dir).iterdir()
        if p.suffix.lower() in extensions
    )

    print(f"Found {len(paths)} images in {input_dir}")
    success = 0

    for i, p in enumerate(paths):
        out_path = os.path.join(output_dir, p.stem + '_processed.jpg')
        print(f"\n[{i+1}/{len(paths)}]")
        if preprocess(str(p), out_path, **kwargs):
            success += 1

    print(f"\nDone! {success}/{len(paths)} images processed.")
    return success


def compare(original_path: str, processed_path: str,
            output_path: str = None) -> np.ndarray:
    orig = cv2.imread(original_path)
    proc = cv2.imread(processed_path)

    if orig is None or proc is None:
        return None

    target_h = 800
    orig_h, orig_w = orig.shape[:2]
    proc_h, proc_w = proc.shape[:2]

    orig_r = cv2.resize(orig, (int(orig_w * target_h / orig_h), target_h))
    proc_r = cv2.resize(proc, (int(proc_w * target_h / proc_h), target_h))

    label_h = 40
    canvas_w = orig_r.shape[1] + proc_r.shape[1] + 10
    canvas = np.ones((target_h + label_h, canvas_w, 3), dtype=np.uint8) * 240

    canvas[label_h:label_h+target_h, :orig_r.shape[1]] = orig_r
    cv2.putText(canvas, 'ORIGINAL', (10, 30),
                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (50,50,50), 2)

    x_off = orig_r.shape[1] + 10
    canvas[label_h:label_h+target_h, x_off:x_off+proc_r.shape[1]] = proc_r
    cv2.putText(canvas, 'PROCESSED', (x_off+10, 30),
                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,100,0), 2)

    if output_path:
        cv2.imwrite(output_path, canvas)
        print(f"  Compare saved: {output_path}")

    return canvas


def main():
    parser = argparse.ArgumentParser(
        description="Image Preprocessor - Tiền xử lý ảnh scan sách cũ"
    )
    parser.add_argument("--image", required=True,
                        help="Đường dẫn ảnh hoặc thư mục")
    parser.add_argument("--output", default="./output/enhanced",
                        help="Thư mục lưu ảnh đã xử lý")
    parser.add_argument("--min_area", type=int, default=30,
                        help="Diện tích tối thiểu của CC (mặc định: 30)")
    parser.add_argument("--max_area_ratio", type=float, default=0.01,
                        help="Tỷ lệ diện tích tối đa (mặc định: 0.01)")
    parser.add_argument("--no_dilate", action="store_true")
    parser.add_argument("--no_auto", action="store_true",
                        help="Tắt auto-detect, luôn enhance")
    parser.add_argument("--compare", action="store_true",
                        help="Lưu ảnh so sánh before/after")
    args = parser.parse_args()

    kwargs = {
        'min_area': args.min_area,
        'max_area_ratio': args.max_area_ratio,
        'dilate': not args.no_dilate,
        'auto': not args.no_auto,
    }

    if os.path.isdir(args.image):
        preprocess_dir(args.image, args.output, **kwargs)
    else:
        stem = Path(args.image).stem
        out_path = os.path.join(args.output, f"{stem}_processed.jpg")

        ok = preprocess(args.image, out_path, **kwargs)

        if ok and args.compare:
            compare_path = os.path.join(args.output, f"{stem}_compare.jpg")
            compare(args.image, out_path, compare_path)


if __name__ == "__main__":
    main()