medical_leaflets_detection/ocr_pipeline.py

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