init

ocr_pipeline.py

@@ -0,0 +1,302 @@
+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()