gauravmeena0708/epfo-circulars
0
1# pdf_parser.py2import os3import re4try:5 import pymupdf as fitz6except ImportError:7 import fitz8from pdf2image import convert_from_path9from PIL import Image10import numpy as np11import cv212import logging13from concurrent.futures import ThreadPoolExecutor, as_completed14 15import config16 17# Configure logging18logger = logging.getLogger(__name__)19logging.basicConfig(level=config.LOG_LEVEL, format=config.LOG_FORMAT)20 21 22def convert_pdf_page_to_image(pdf_path, page_num, dpi=config.PDF_TO_IMAGE_DPI):23 """Converts a single page of a PDF to a PIL Image."""24 try:25 images = convert_from_path(pdf_path, dpi=dpi, first_page=page_num + 1, last_page=page_num + 1)26 if images:27 return images[0]28 except Exception as e:29 logger.error(f"Error converting page {page_num} of PDF '{pdf_path}' to image: {e}")30 return None31 32 33def pil_to_cv2(pil_image):34 """Converts a PIL Image to an OpenCV image (BGR format)."""35 return cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)36 37 38def _extract_native_page_text(doc, page_num, source_pdf):39 """40 Attempts native text extraction from a PDF page using PyMuPDF.41 Includes a safety net to verify text sufficiency and quality.42 43 Returns:44 tuple: (success (bool), page_data (dict or None))45 """46 try:47 page = doc[page_num]48 raw_text = page.get_text("text") or ""49 words = raw_text.split()50 51 # Safety Net Checks:52 # 1. Minimum word count threshold53 if len(words) < config.NATIVE_TEXT_MIN_WORDS:54 return False, None55 56 # 2. Printable character ratio (detect broken/garbled font mappings)57 printable_count = sum(1 for c in raw_text if c.isprintable() or c.isspace())58 if len(raw_text) > 0 and (printable_count / len(raw_text)) < 0.80:59 logger.info(f"Page {page_num + 1} of '{source_pdf}' failed printable character ratio ({printable_count}/{len(raw_text)}). Falling back to OCR.")60 return False, None61 62 # Extract structured blocks from PyMuPDF63 # format: (x0, y0, x1, y1, "lines of text", block_no, block_type)64 blocks = page.get_text("blocks")65 content_items = []66 67 # Check if PyMuPDF table extraction is available68 tables = []69 try:70 if hasattr(page, "find_tables"):71 tab_finder = page.find_tables()72 if tab_finder and tab_finder.tables:73 for tab in tab_finder.tables:74 tab_bbox = [int(v) for v in tab.bbox]75 extracted_rows = tab.extract()76 flat_cells = [str(c).strip() for row in extracted_rows for c in row if c is not None and str(c).strip()]77 if flat_cells:78 tables.append({79 "bbox": tab_bbox,80 "cells": flat_cells81 })82 except Exception as tab_e:83 logger.debug(f"PyMuPDF find_tables not available or failed: {tab_e}")84 85 for b in blocks:86 if b[6] == 0: # Text block87 b_text = b[4].strip()88 if not b_text:89 continue90 x0, y0, x1, y1 = int(b[0]), int(b[1]), int(b[2]), int(b[3])91 92 # Check if this text is already inside an extracted table93 inside_table = False94 for tbl in tables:95 tx0, ty0, tx1, ty1 = tbl["bbox"]96 if x0 >= tx0 - 5 and y0 >= ty0 - 5 and x1 <= tx1 + 5 and y1 <= ty1 + 5:97 inside_table = True98 break99 100 if not inside_table:101 content_items.append((102 y0,103 {104 "type": "plain_text",105 "text": b_text,106 "bbox_pil": [x0, y0, x1, y1]107 }108 ))109 110 for tbl in tables:111 x0, y0, x1, y1 = tbl["bbox"]112 is_at_top = (y0 <= 100)113 is_at_bottom = (y1 >= page.rect.height - 100)114 content_items.append((115 y0,116 {117 "type": "table",118 "bbox_pil": [x0, y0, x1, y1],119 "extracted_text_list": tbl["cells"],120 "raw_cells": [{"text": c, "bbox_pil": [x0, y0, x1, y1]} for c in tbl["cells"]],121 "is_at_page_top": is_at_top,122 "is_at_page_bottom": is_at_bottom123 }124 ))125 126 content_items.sort(key=lambda x: x[0])127 page_data = {128 "page_number": page_num + 1,129 "source_pdf": source_pdf,130 "content": [item[1] for item in content_items]131 }132 return True, page_data133 134 except Exception as e:135 logger.warning(f"Error during native text extraction on page {page_num + 1} of '{source_pdf}': {e}. Falling back to OCR.")136 return False, None137 138 139def _process_page_ocr(pdf_path, page_num, table_detector_pipeline, ocr_reader):140 """Helper function to process a single page using Table Transformer + EasyOCR."""141 page_content_blocks = []142 current_page_data = {"page_number": page_num + 1, "source_pdf": os.path.basename(pdf_path), "content": []}143 144 try:145 pil_image = convert_pdf_page_to_image(pdf_path, page_num, dpi=config.PDF_TO_IMAGE_DPI)146 if not pil_image:147 logger.warning(f"Could not convert page {page_num} of '{pdf_path}' to image. Skipping page.")148 return current_page_data149 150 table_boxes_pil = []151 if table_detector_pipeline:152 table_detections = table_detector_pipeline(pil_image)153 for detection in table_detections:154 if detection['label'] == 'table':155 box = detection['box']156 x0 = max(0, int(box['xmin']))157 y0 = max(0, int(box['ymin']))158 x1 = min(pil_image.width, int(box['xmax']))159 y1 = min(pil_image.height, int(box['ymax']))160 if x1 > x0 and y1 > y0:161 table_boxes_pil.append((x0, y0, x1, y1))162 163 img_np_rgb = np.array(pil_image)164 non_table_mask = np.ones(img_np_rgb.shape[:2], dtype=np.uint8) * 255165 for x0, y0, x1, y1 in table_boxes_pil:166 non_table_mask[y0:y1, x0:x1] = 0167 168 non_table_img_np = cv2.bitwise_and(img_np_rgb, img_np_rgb, mask=non_table_mask)169 non_table_ocr_results = ocr_reader.readtext(non_table_img_np, paragraph=True) if ocr_reader else []170 171 for ocr_result in non_table_ocr_results:172 if len(ocr_result) == 2:173 bbox, text = ocr_result174 elif len(ocr_result) == 3:175 bbox, text, _ = ocr_result176 else:177 continue178 179 if bbox and isinstance(bbox, list) and len(bbox) > 0 and \180 isinstance(bbox[0], (list, tuple)) and len(bbox[0]) == 2:181 pos_y = int(bbox[0][1])182 page_content_blocks.append(183 (pos_y, {"type": "plain_text", "text": text, "bbox_pil": [int(c) for pt in bbox for c in pt]})184 )185 186 page_height = pil_image.height187 tolerance = 100 # Pixels tolerance for checking if table is at edge188 189 for x0_tbl, y0_tbl, x1_tbl, y1_tbl in table_boxes_pil:190 table_pil_image_crop = pil_image.crop((x0_tbl, y0_tbl, x1_tbl, y1_tbl))191 table_ocr_results = ocr_reader.readtext(np.array(table_pil_image_crop)) if ocr_reader else []192 193 table_cells_text = []194 raw_table_cells = []195 for item in table_ocr_results:196 if len(item) == 3:197 bbox_cell, text_cell, prob_cell = item198 elif len(item) == 2:199 bbox_cell, text_cell = item200 else:201 continue202 table_cells_text.append(text_cell)203 adjusted_bbox_cell = [[int(pt[0] + x0_tbl), int(pt[1] + y0_tbl)] for pt in bbox_cell]204 raw_table_cells.append({"text": text_cell, "bbox_pil": [int(c) for pt in adjusted_bbox_cell for c in pt]})205 206 if table_cells_text:207 is_at_top = (y0_tbl <= tolerance)208 is_at_bottom = (y1_tbl >= page_height - tolerance)209 page_content_blocks.append(210 (y0_tbl, {211 "type": "table",212 "bbox_pil": [x0_tbl, y0_tbl, x1_tbl, y1_tbl],213 "extracted_text_list": table_cells_text,214 "raw_cells": raw_table_cells,215 "is_at_page_top": is_at_top,216 "is_at_page_bottom": is_at_bottom217 })218 )219 220 page_content_blocks.sort(key=lambda x: x[0])221 current_page_data["content"] = [block[1] for block in page_content_blocks]222 223 except Exception as e:224 logger.error(f"Error OCR processing page {page_num} of PDF '{pdf_path}': {e}", exc_info=True)225 226 return current_page_data227 228 229def _process_single_page_hybrid(pdf_path, page_num, doc, table_detector_pipeline, ocr_reader):230 """231 Executes native PyMuPDF text extraction first with safety net;232 falls back to OCR if word density is low or page is scanned.233 """234 source_pdf = os.path.basename(pdf_path)235 236 if getattr(config, "USE_NATIVE_PDF_TEXT", True) and doc is not None:237 success, native_data = _extract_native_page_text(doc, page_num, source_pdf)238 if success and native_data and native_data.get("content"):239 logger.debug(f"Native extraction successful for page {page_num + 1} of '{source_pdf}'.")240 return native_data241 242 logger.info(f"Using OCR fallback for page {page_num + 1} of '{source_pdf}'.")243 return _process_page_ocr(pdf_path, page_num, table_detector_pipeline, ocr_reader)244 245 246def extract_content_from_pdf(pdf_path, table_detector_pipeline=None, ocr_reader=None, max_workers=4):247 """248 Extracts structured content (text and tables) from a PDF file.249 Uses native text extraction where available and falls back to OCR.250 251 Args:252 pdf_path (str): Path to the PDF file.253 table_detector_pipeline: Initialized table detection pipeline (optional for native-only).254 ocr_reader: Initialized EasyOCR reader (optional for native-only).255 max_workers (int): Maximum number of threads for parallel processing.256 257 Returns:258 list: A list of dicts with 'page_number' and 'content' list.259 """260 extracted_pdf_data = []261 try:262 doc = fitz.open(pdf_path)263 num_pages = len(doc)264 except Exception as e:265 logger.error(f"Error opening PDF '{pdf_path}': {e}")266 return extracted_pdf_data267 268 # Check how many pages can be natively parsed269 all_pages_native = True270 page_results = [None] * num_pages271 272 if getattr(config, "USE_NATIVE_PDF_TEXT", True):273 for page_num in range(num_pages):274 success, native_data = _extract_native_page_text(doc, page_num, os.path.basename(pdf_path))275 if success and native_data and native_data.get("content"):276 page_results[page_num] = native_data277 else:278 all_pages_native = False279 280 if all_pages_native and all(p is not None for p in page_results):281 doc.close()282 logger.info(f"All {num_pages} pages of '{os.path.basename(pdf_path)}' extracted natively via PyMuPDF.")283 return sorted(page_results, key=lambda x: x['page_number'])284 285 # For pages requiring OCR, run OCR in thread pool286 missing_page_nums = [i for i, r in enumerate(page_results) if r is None]287 logger.info(f"PDF '{os.path.basename(pdf_path)}': {len(page_results) - len(missing_page_nums)}/{num_pages} pages parsed natively. Running OCR fallback on {len(missing_page_nums)} pages.")288 289 with ThreadPoolExecutor(max_workers=max_workers) as executor:290 futures = {291 executor.submit(_process_page_ocr, pdf_path, page_num, table_detector_pipeline, ocr_reader): page_num 292 for page_num in missing_page_nums293 }294 for future in as_completed(futures):295 p_num = futures[future]296 page_results[p_num] = future.result()297 298 doc.close()299 extracted_pdf_data = [p for p in page_results if p is not None]300 extracted_pdf_data = sorted(extracted_pdf_data, key=lambda x: x['page_number'])301 return extracted_pdf_data