mirror of
https://git.victorphan.net/basketballcantho/medical_leaflets_detection.git
synced 2026-08-05 06:33:11 +07:00
303 lines
12 KiB
Python
303 lines
12 KiB
Python
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()
|