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()