import os import gc import cv2 import numpy as np from PIL import Image, ImageFilter, ImageEnhance from pdf2image import convert_from_path, pdfinfo_from_path from paddleocr import PaddleOCR from vietocr.tool.predictor import Predictor from vietocr.tool.config import Cfg import torch # ============================================================ # CONFIGURATION - Tuned for MAXIMUM ACCURACY # ============================================================ INPUT_FOLDER = "input_pdf_folder" OUTPUT_FOLDER = "output_txt_folder" MODEL_FOLDER = "models" VIETOCR_MODEL_PATH = os.path.join(MODEL_FOLDER, "vgg_transformer.pth") # --- Quality Settings --- DPI = 600 # Maximum practical DPI for OCR PADDING = 8 # Pixels of padding around each text box MIN_BOX_WIDTH = 15 # Minimum box width in pixels (filter noise) MIN_BOX_HEIGHT = 10 # Minimum box height in pixels (filter noise) MIN_BOX_AREA = 300 # Minimum box area in pixels² (filter tiny artifacts) Y_TOLERANCE = 20 # Pixels tolerance for same-line grouping (higher DPI = more tolerance) # --- PaddleOCR Detection Tuning --- # Lower thresholds = catch more text, but may also catch noise # Higher thresholds = miss some text, but results are cleaner DET_DB_THRESH = 0.25 # Binarization threshold (default 0.3, lower = more sensitive) DET_DB_BOX_THRESH = 0.45 # Box confidence threshold (default 0.6, lower = more boxes) DET_DB_UNCLIP_RATIO = 1.8 # Box expansion ratio (default 1.5, higher = larger boxes) DET_LIMIT_SIDE_LEN = 2560 # Max image side length for detection (default 960, higher = more detail) def setup_directories(): """Create necessary directories if they don't exist.""" os.makedirs(INPUT_FOLDER, exist_ok=True) os.makedirs(OUTPUT_FOLDER, exist_ok=True) os.makedirs(MODEL_FOLDER, exist_ok=True) def download_vietocr_weights(): """Download VietOCR weights if not present with corruption protection.""" if not os.path.exists(VIETOCR_MODEL_PATH): print("Downloading VietOCR weights (vgg_transformer)...") import requests temp_path = VIETOCR_MODEL_PATH + ".tmp" url = "https://vocr.vn/data/vietocr/vgg_transformer.pth" try: response = requests.get(url, stream=True) if response.status_code == 200: with open(temp_path, 'wb') as f: for chunk in response.iter_content(chunk_size=8192): f.write(chunk) # Verify the file is a valid torch checkpoint try: torch.load(temp_path, map_location='cpu') except Exception: os.remove(temp_path) raise Exception("Downloaded file is corrupted. Retrying on next run.") os.rename(temp_path, VIETOCR_MODEL_PATH) print("Download completed and verified successfully.") else: raise Exception(f"Failed to download weights. Status code: {response.status_code}") except Exception as e: if os.path.exists(temp_path): os.remove(temp_path) raise e def preprocess_for_detection(pil_img): """ Station 1a: Prepare image for PaddleOCR text detection. Uses CLAHE + denoising on grayscale for optimal box detection. Returns a 3-channel image ready for PaddleOCR. """ img = np.array(pil_img) if len(img.shape) == 3 and img.shape[2] == 3: gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) else: gray = img # CLAHE for contrast enhancement clahe = cv2.createCLAHE(clipLimit=2.5, tileGridSize=(8, 8)) enhanced = clahe.apply(gray) # Light denoising (preserve edges) denoised = cv2.fastNlMeansDenoising(enhanced, None, 8, 7, 21) # Convert back to 3-channel for PaddleOCR return cv2.cvtColor(denoised, cv2.COLOR_GRAY2RGB) def preprocess_for_recognition(pil_img): """ Station 1b: Prepare the ORIGINAL high-res image for VietOCR recognition. Uses the original color/grayscale image (not the detection-preprocessed one) with mild sharpening for best character recognition. Returns a numpy array used for cropping individual text boxes. """ # Sharpen the original PIL image mildly sharpened = pil_img.filter(ImageFilter.SHARPEN) # Boost contrast slightly enhancer = ImageEnhance.Contrast(sharpened) enhanced = enhancer.enhance(1.3) return np.array(enhanced) def get_padded_crop(img, box, padding=PADDING): """ Extract a padded crop from the image based on the bounding box. Uses perspective transform for rotated boxes to get a clean crop. """ box_arr = np.array(box).astype(np.float32) # Calculate width and height of the oriented bounding box width = int(max( np.linalg.norm(box_arr[0] - box_arr[1]), np.linalg.norm(box_arr[2] - box_arr[3]) )) height = int(max( np.linalg.norm(box_arr[0] - box_arr[3]), np.linalg.norm(box_arr[1] - box_arr[2]) )) if width < MIN_BOX_WIDTH or height < MIN_BOX_HEIGHT: return None if width * height < MIN_BOX_AREA: return None # Add padding to the destination dimensions dst_width = width + 2 * padding dst_height = height + 2 * padding # Define destination points with padding offset dst_pts = np.array([ [padding, padding], [padding + width, padding], [padding + width, padding + height], [padding, padding + height] ], dtype=np.float32) # Perspective transform to handle rotated text M = cv2.getPerspectiveTransform(box_arr, dst_pts) cropped = cv2.warpPerspective( img, M, (dst_width, dst_height), borderMode=cv2.BORDER_REPLICATE ) return cropped def sort_boxes(boxes): """ Logical Sorting & Column Handling: Sort bounding boxes primarily by Y-coordinate and secondarily by X-coordinate to ensure correct reading order for multi-column medical leaflets. """ if not boxes: return [] # First sort purely by y-coordinate (top-left y) boxes.sort(key=lambda b: b[0][1]) sorted_boxes = [] current_line = [] current_y = boxes[0][0][1] for box in boxes: y1 = box[0][1] if abs(y1 - current_y) <= Y_TOLERANCE: current_line.append(box) else: current_line.sort(key=lambda b: b[0][0]) sorted_boxes.extend(current_line) current_line = [box] current_y = y1 if current_line: current_line.sort(key=lambda b: b[0][0]) sorted_boxes.extend(current_line) return sorted_boxes def process_pdf(pdf_path, detector, recognizer): """Process a single PDF file through the OCR pipeline.""" pdf_name = os.path.basename(pdf_path) txt_filename = os.path.splitext(pdf_name)[0] + ".txt" txt_filepath = os.path.join(OUTPUT_FOLDER, txt_filename) try: info = pdfinfo_from_path(pdf_path) total_pages = info["Pages"] with open(txt_filepath, 'w', encoding='utf-8') as f_out: for page_num in range(1, total_pages + 1): print(f"Processing [{pdf_name}] - Page [{page_num}/{total_pages}]...") # Station 0: PDF to Image Conversion (600 DPI for maximum detail) images = convert_from_path( pdf_path, dpi=DPI, first_page=page_num, last_page=page_num ) if not images: continue pil_img = images[0] # Station 1a: Preprocess for detection (grayscale, denoised) detection_img = preprocess_for_detection(pil_img) # Station 1b: Preprocess for recognition (original, sharpened) recognition_img = preprocess_for_recognition(pil_img) # Station 2: Text Detection (PaddleOCR) dt_boxes = detector.ocr(detection_img, rec=False, cls=False) if dt_boxes and len(dt_boxes) > 0 and dt_boxes[0] is not None: boxes = dt_boxes[0] boxes = sort_boxes(boxes) recognized_count = 0 for box in boxes: # Station 3: Crop from the RECOGNITION image (not detection image) # This gives VietOCR the highest quality input cropped_cv = get_padded_crop(recognition_img, box) if cropped_cv is None: continue # Convert to PIL for VietOCR cropped_pil = Image.fromarray(cropped_cv) if cropped_pil.mode != 'RGB': cropped_pil = cropped_pil.convert('RGB') # Recognize text text = recognizer.predict(cropped_pil) text = text.strip() # Filter out very short garbage text (single chars that are likely noise) if len(text) >= 1: f_out.write(text + "\n") recognized_count += 1 print(f" → Detected {len(boxes)} boxes, recognized {recognized_count} text lines.") # Memory cleanup del images, pil_img, detection_img, recognition_img gc.collect() print(f"Processing [{pdf_name}] - Page [{page_num}/{total_pages}]... Done.") except Exception as e: print(f"Error processing {pdf_name}: {e}") import traceback traceback.print_exc() def main(): setup_directories() # Initialize PaddleOCR detector (CPU ONLY, tuned for maximum accuracy) print("Initializing Station 2: PaddleOCR Detector (CPU, MAX ACCURACY mode)...") detector = PaddleOCR( use_angle_cls=True, # Enable angle classification for rotated text lang='vi', use_gpu=False, show_log=False, det_db_thresh=DET_DB_THRESH, det_db_box_thresh=DET_DB_BOX_THRESH, det_db_unclip_ratio=DET_DB_UNCLIP_RATIO, det_limit_side_len=DET_LIMIT_SIDE_LEN, use_dilation=True, # Dilate text regions for better detection ) # Initialize VietOCR recognizer (CPU ONLY, beam search for max accuracy) print("Initializing Station 3: VietOCR Recognizer (MAX ACCURACY mode)...") download_vietocr_weights() config = Cfg.load_config_from_name('vgg_transformer') config['weights'] = VIETOCR_MODEL_PATH config['device'] = 'cpu' config['predictor']['beamsearch'] = True # Beam search is slower but more accurate recognizer = Predictor(config) # Scan input folder for PDFs pdf_files = [f for f in os.listdir(INPUT_FOLDER) if f.lower().endswith('.pdf')] if not pdf_files: print(f"No PDF files found in {INPUT_FOLDER}/. Please add PDF files and run again.") return print(f"Found {len(pdf_files)} PDF files to process.") print(f"Settings: DPI={DPI}, Padding={PADDING}px, BeamSearch=ON") print(f"Detection: thresh={DET_DB_THRESH}, box_thresh={DET_DB_BOX_THRESH}, " f"unclip={DET_DB_UNCLIP_RATIO}, max_side={DET_LIMIT_SIDE_LEN}") print("=" * 60) for pdf_file in pdf_files: pdf_path = os.path.join(INPUT_FOLDER, pdf_file) process_pdf(pdf_path, detector, recognizer) print("=" * 60) print("Batch processing completed.") if __name__ == "__main__": main()