import cv2 import numpy as np import argparse import os from pathlib import Path class ImagePreprocessor: def __init__(self, to_grayscale: bool = True, normalize_bg: bool = True, denoise: bool = True, denoise_method: str = "gaussian", # gaussian | nlm clahe: bool = True, clahe_clip: float = 2.0, clahe_grid: int = 8, binarize: bool = True, binarize_method: str = "otsu", # otsu | adaptive | sauvola deskew: bool = True, sharpen: bool = True, morph_clean: bool = True, target_height: int = 48, padding: int = 4): self.to_grayscale = to_grayscale self.normalize_bg = normalize_bg self.denoise = denoise self.denoise_method = denoise_method self.clahe = clahe self.clahe_clip = clahe_clip self.clahe_grid = clahe_grid self.binarize = binarize self.binarize_method = binarize_method self.deskew = deskew self.sharpen = sharpen self.morph_clean = morph_clean self.target_height = target_height self.padding = padding def process(self, image: np.ndarray, verbose: bool = False) -> np.ndarray: if image is None or image.size == 0: return image steps = {} img = image.copy() steps["original"] = img.copy() # 1. Grayscale if len(img.shape) == 3: gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) else: gray = img.copy() steps["grayscale"] = gray.copy() if self.normalize_bg: gray = self._normalize_background(gray) steps["bg_norm"] = gray.copy() if self.denoise: if self.denoise_method == "nlm": gray = cv2.fastNlMeansDenoising(gray, h=7, templateWindowSize=7, searchWindowSize=21) else: gray = cv2.GaussianBlur(gray, (3, 3), 0) steps["denoise"] = gray.copy() if self.clahe: clahe_obj = cv2.createCLAHE( clipLimit=self.clahe_clip, tileGridSize=(self.clahe_grid, self.clahe_grid) ) gray = clahe_obj.apply(gray) steps["clahe"] = gray.copy() if self.sharpen: gray = self._sharpen(gray) steps["sharpen"] = gray.copy() if self.binarize: binary = self._binarize(gray) steps["binarize"] = binary.copy() else: binary = gray.copy() if self.deskew: binary = self._deskew(binary) steps["deskew"] = binary.copy() if self.morph_clean: binary = self._morph_clean(binary) steps["morph_clean"] = binary.copy() ch, cw = binary.shape[:2] if ch < self.target_height: scale = self.target_height / ch new_w = max(1, int(cw * scale)) binary = cv2.resize( binary, (new_w, self.target_height), interpolation=cv2.INTER_CUBIC ) steps["upscale"] = binary.copy() if self.padding > 0: binary = cv2.copyMakeBorder( binary, self.padding, self.padding, self.padding, self.padding, cv2.BORDER_CONSTANT, value=255 ) result = cv2.cvtColor(binary, cv2.COLOR_GRAY2BGR) if verbose: print(f" Original: {image.shape} → Final: {result.shape}") self._steps = steps return result def _normalize_background(self, gray: np.ndarray) -> np.ndarray: h, w = gray.shape ksize = max(h // 2, w // 8, 15) ksize = min(ksize, 60, h - 2, w - 2) # không vượt quá kích thước ảnh ksize = max(ksize, 3) ksize = ksize if ksize % 2 == 1 else ksize + 1 kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (ksize, ksize)) bg = cv2.morphologyEx(gray, cv2.MORPH_DILATE, kernel) gray_f = gray.astype(np.float32) bg_f = bg.astype(np.float32) normalized = (gray_f / (bg_f + 1e-6)) * 255.0 normalized = np.clip(normalized, 0, 255).astype(np.uint8) return normalized def _sharpen(self, gray: np.ndarray) -> np.ndarray: blurred = cv2.GaussianBlur(gray, (0, 0), 3) sharpened = cv2.addWeighted(gray, 1.5, blurred, -0.5, 0) return np.clip(sharpened, 0, 255).astype(np.uint8) def _binarize(self, gray: np.ndarray) -> np.ndarray: if self.binarize_method == "otsu": _, binary = cv2.threshold( gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU ) elif self.binarize_method == "adaptive": binary = cv2.adaptiveThreshold( gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, blockSize=15, C=8 ) elif self.binarize_method == "sauvola": # Sauvola: tốt nhất cho ảnh scan cũ binary = self._sauvola_threshold(gray) else: _, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY) # Đảm bảo chữ đen nền trắng binary = self._ensure_dark_text(binary) return binary def _sauvola_threshold(self, gray: np.ndarray, window_size: int = None, k: float = 0.2) -> np.ndarray: h, w = gray.shape if window_size is None: ws = max(11, min(h // 2, 31)) window_size = ws if ws % 2 == 1 else ws + 1 gray_f = gray.astype(np.float64) R = 128.0 mean = cv2.boxFilter(gray_f, -1, (window_size, window_size)) mean_sq = cv2.boxFilter(gray_f**2, -1, (window_size, window_size)) std = np.sqrt(np.maximum(mean_sq - mean**2, 0)) threshold = mean * (1.0 + k * (std / R - 1.0)) binary = np.where(gray_f <= threshold, 0, 255).astype(np.uint8) return binary def _ensure_dark_text(self, binary: np.ndarray) -> np.ndarray: black_pixels = np.sum(binary == 0) white_pixels = np.sum(binary == 255) if black_pixels > white_pixels: binary = cv2.bitwise_not(binary) return binary def _deskew(self, binary: np.ndarray) -> np.ndarray: h, w = binary.shape[:2] if h < 20 or w < 100: return binary try: inv = cv2.bitwise_not(binary) edges = cv2.Canny(inv, 50, 150, apertureSize=3) lines = cv2.HoughLinesP( edges, 1, np.pi/180, threshold=max(30, w//10), minLineLength=w//4, maxLineGap=20 ) if lines is None or len(lines) < 3: return binary # Tính góc từ các đường nằm ngang angles = [] for line in lines: x1, y1, x2, y2 = line[0] if abs(x2 - x1) > abs(y2 - y1): # đường nằm ngang angle = np.degrees(np.arctan2(y2 - y1, x2 - x1)) if abs(angle) < 10: # chỉ lấy góc nhỏ angles.append(angle) if not angles: return binary # Lấy median angle angle = float(np.median(angles)) # Chỉ sửa nếu nghiêng đáng kể (>0.5°) và nhỏ (<5°) if abs(angle) < 0.5 or abs(angle) > 5: return binary # Xoay ảnh center = (w // 2, h // 2) M = cv2.getRotationMatrix2D(center, angle, 1.0) rotated = cv2.warpAffine( binary, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_CONSTANT, borderValue=255 ) return rotated except Exception: return binary def _morph_clean(self, binary: np.ndarray) -> np.ndarray: h = binary.shape[0] if h < 30: return binary kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2, 2)) cleaned = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel) text_before = np.sum(binary == 0) text_after = np.sum(cleaned == 0) if text_after < text_before * 0.6: return binary return cleaned def get_steps(self) -> dict: return getattr(self, '_steps', {}) def compare_steps(image: np.ndarray, preprocessor: ImagePreprocessor, save_path: str = None, title: str = "") -> np.ndarray: result = preprocessor.process(image, verbose=True) steps = preprocessor.get_steps() n = len(steps) + 1 names = list(steps.keys()) + ["final"] imgs = list(steps.values()) + [cv2.cvtColor(result, cv2.COLOR_BGR2GRAY)] target_h = 80 resized = [] for img in imgs: if len(img.shape) == 3: img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) h, w = img.shape scale = target_h / h new_w = max(1, int(w * scale)) r = cv2.resize(img, (new_w, target_h)) resized.append(r) max_w = max(r.shape[1] for r in resized) label_h = 20 cell_h = target_h + label_h + 4 grid_rows = (n + 3) // 4 grid_cols = min(n, 4) canvas_h = grid_rows * cell_h + 30 canvas_w = grid_cols * (max_w + 10) + 10 canvas = np.ones((canvas_h, canvas_w), dtype=np.uint8) * 200 for idx, (name, img) in enumerate(zip(names, resized)): row = idx // 4 col = idx % 4 x = col * (max_w + 10) + 5 y = row * cell_h + 25 if name == "final": canvas[y-2:y+target_h+2, x-2:x+img.shape[1]+2] = 0 canvas[y-1:y+target_h+1, x-1:x+img.shape[1]+1] = 200 canvas[y:y+target_h, x:x+img.shape[1]] = img cv2.putText(canvas, name, (x, y-3), cv2.FONT_HERSHEY_SIMPLEX, 0.35, 30, 1) if title: cv2.putText(canvas, title, (5, 15), cv2.FONT_HERSHEY_SIMPLEX, 0.5, 0, 1) if save_path: cv2.imwrite(save_path, canvas) print(f" Saved: {save_path}") return canvas def test_with_rec(image: np.ndarray, preprocessor: ImagePreprocessor, rec_model_dir: str = "./models/inference_rec"): try: from paddlex import create_predictor predictor = create_predictor( model_name='latin_PP-OCRv5_mobile_rec', model_dir=rec_model_dir ) # Before results_before = list(predictor.predict(image)) text_before = results_before[0].get('rec_text', '') if results_before else '' score_before = results_before[0].get('rec_score', 0) if results_before else 0 # After preprocessing processed = preprocessor.process(image) results_after = list(predictor.predict(processed)) text_after = results_after[0].get('rec_text', '') if results_after else '' score_after = results_after[0].get('rec_score', 0) if results_after else 0 print(f"\n Before: [{score_before:.4f}] {text_before}") print(f" After: [{score_after:.4f}] {text_after}") improvement = score_after - score_before if improvement > 0.01: print(f" ✅ Improved: +{improvement:.4f}") elif improvement < -0.01: print(f" ⚠️ Worse: {improvement:.4f}") else: print(f" ➡️ Similar: {improvement:.4f}") return text_before, score_before, text_after, score_after except Exception as e: print(f" Rec test skipped: {e}") return None, None, None, None def main(): parser = argparse.ArgumentParser( description="Image Preprocessor - Test tiền xử lý ảnh OCR" ) parser.add_argument("--image", required=True) parser.add_argument("--output", default="./output/preprocess_test") parser.add_argument("--method", default="adaptive", choices=["otsu", "adaptive", "sauvola"], help="Binarization method") parser.add_argument("--denoise", default="gaussian", choices=["gaussian", "nlm"]) parser.add_argument("--no_clahe", action="store_true") parser.add_argument("--no_sharpen", action="store_true") parser.add_argument("--no_deskew", action="store_true") parser.add_argument("--no_morph", action="store_true") parser.add_argument("--compare", action="store_true", help="Lưu ảnh so sánh từng bước") parser.add_argument("--test_rec", action="store_true", help="Test rec trước/sau preprocessing") parser.add_argument("--rec_model", default="./models/inference_rec") args = parser.parse_args() os.makedirs(args.output, exist_ok=True) # Khởi tạo preprocessor preprocessor = ImagePreprocessor( to_grayscale=True, denoise=True, denoise_method=args.denoise, clahe=not args.no_clahe, binarize=True, binarize_method=args.method, deskew=not args.no_deskew, sharpen=not args.no_sharpen, morph_clean=not args.no_morph, padding=4, ) print(f"Config: binarize={args.method}, denoise={args.denoise}, " f"clahe={not args.no_clahe}, sharpen={not args.no_sharpen}") # Load ảnh image = cv2.imread(args.image) if image is None: print(f"Cannot read: {args.image}") return h, w = image.shape[:2] print(f"Image: {w}x{h}") stem = Path(args.image).stem # Nếu ảnh nhỏ (crop image) → xử lý trực tiếp is_crop = h < 100 or (h < 200 and w / h > 5) if is_crop: print(" [CROP] Processing single crop image...") # Process result = preprocessor.process(image, verbose=True) out_path = os.path.join(args.output, f"{stem}_processed.jpg") cv2.imwrite(out_path, result) print(f" Saved: {out_path}") # So sánh if args.compare: compare_path = os.path.join(args.output, f"{stem}_compare.jpg") compare_steps(image, preprocessor, compare_path, title=stem) # Test rec if args.test_rec: test_with_rec(image, preprocessor, args.rec_model) else: # Ảnh lớn → detect → crop từng region → process print(" [PAGE] Detecting text regions...") gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU) kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (50, 5)) dilated = cv2.dilate(thresh, kernel, iterations=2) contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # Sắp xếp top→bottom bboxes = [] for c in contours: x, y, cw, ch = cv2.boundingRect(c) if cw > 50 and ch > 10: bboxes.append((x, y, cw, ch)) bboxes.sort(key=lambda b: (b[1], b[0])) print(f" Found {len(bboxes)} regions") # Tạo visualization vis = image.copy() results_log = [] for idx, (x, y, cw, ch) in enumerate(bboxes[:20]): # max 20 pad = 5 x1 = max(0, x - pad) y1 = max(0, y - pad) x2 = min(w, x + cw + pad) y2 = min(h, y + ch + pad) crop = image[y1:y2, x1:x2] # Process processed_crop = preprocessor.process(crop) # Lưu crop trước/sau crop_dir = os.path.join(args.output, "crops") os.makedirs(crop_dir, exist_ok=True) cv2.imwrite(os.path.join(crop_dir, f"{idx:03d}_before.jpg"), crop) cv2.imwrite(os.path.join(crop_dir, f"{idx:03d}_after.jpg"), processed_crop) # So sánh if args.compare: cmp_path = os.path.join(crop_dir, f"{idx:03d}_compare.jpg") compare_steps(crop, ImagePreprocessor( to_grayscale=True, denoise=True, denoise_method=args.denoise, clahe=not args.no_clahe, binarize=True, binarize_method=args.method, deskew=not args.no_deskew, sharpen=not args.no_sharpen, morph_clean=not args.no_morph, ), cmp_path) # Test rec if args.test_rec: print(f"\n Region {idx+1}:") test_with_rec(crop, preprocessor, args.rec_model) # Vẽ bbox cv2.rectangle(vis, (x1, y1), (x2, y2), (0, 200, 0), 2) cv2.putText(vis, str(idx+1), (x1, y1-3), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0,200,0), 1) vis_path = os.path.join(args.output, f"{stem}_vis.jpg") cv2.imwrite(vis_path, vis) print(f"\n Saved visualization: {vis_path}") print(f" Crops saved in: {os.path.join(args.output, 'crops')}/") print("\nDone!") if __name__ == "__main__": main()