""" PDF Medical Extractor - Phase 2 Structured PDF extraction using Donut/LayoutLMv3 for medical documents. This module provides specialized extraction for medical PDFs including radiology reports, laboratory results, clinical notes, and ECG reports. Author: MiniMax Agent Date: 2025-10-29 Version: 1.0.0 """ import os import json import io import logging from typing import Dict, List, Optional, Any, Tuple from dataclasses import dataclass from pathlib import Path import numpy as np from PIL import Image import fitz # PyMuPDF import pytesseract from transformers import DonutProcessor, VisionEncoderDecoderModel import torch from tqdm import tqdm from medical_schemas import ( MedicalDocumentMetadata, ConfidenceScore, RadiologyAnalysis, LaboratoryResults, ClinicalNotesAnalysis, ValidationResult, validate_document_schema ) logger = logging.getLogger(__name__) @dataclass class ExtractionResult: """Result of PDF extraction with confidence scoring""" raw_text: str structured_data: Dict[str, Any] confidence_scores: Dict[str, float] extraction_method: str # "donut", "ocr", "hybrid" processing_time: float tables_extracted: List[Dict[str, Any]] images_extracted: List[str] metadata: Dict[str, Any] class DonutMedicalExtractor: """Medical PDF extraction using Donut model for structured output""" def __init__(self, model_name: str = "naver-clova-ix/donut-base-finetuned-rvlcdip"): self.model_name = model_name self.processor = None self.model = None self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self._load_model() def _load_model(self): """Load Donut model and processor""" try: logger.info(f"Loading Donut model: {self.model_name}") self.processor = DonutProcessor.from_pretrained(self.model_name) self.model = VisionEncoderDecoderModel.from_pretrained(self.model_name) self.model.to(self.device) self.model.eval() logger.info("Donut model loaded successfully") except Exception as e: logger.error(f"Failed to load Donut model: {str(e)}") raise def extract_from_image(self, image: Image.Image, task_prompt: str = None) -> Dict[str, Any]: """Extract structured data from image using Donut""" if task_prompt is None: task_prompt = "" try: # Prepare image for Donut pixel_values = self.processor(images=image, return_tensors="pt").pixel_values pixel_values = pixel_values.to(self.device) # Generate structured output task_prompt_ids = self.processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids task_prompt_ids = task_prompt_ids.to(self.device) with torch.no_grad(): outputs = self.model.generate( task_prompt_ids, pixel_values, max_length=512, early_stopping=False, pad_token_id=self.processor.tokenizer.pad_token_id, eos_token_id=self.processor.tokenizer.eos_token_id, use_cache=True, ) # Decode output output_sequence = outputs.cpu().numpy()[0] decoded_output = self.processor.tokenizer.decode(output_sequence, skip_special_tokens=True) # Parse JSON from decoded output json_start = decoded_output.find('{') json_end = decoded_output.rfind('}') + 1 if json_start != -1 and json_end != -1: json_str = decoded_output[json_start:json_end] structured_data = json.loads(json_str) else: structured_data = {"raw_text": decoded_output} return structured_data except Exception as e: logger.error(f"Donut extraction error: {str(e)}") return {"raw_text": "", "error": str(e)} class MedicalPDFProcessor: """Medical PDF processing with multiple extraction methods""" def __init__(self): self.donut_extractor = None self.ocr_enabled = True # Initialize Donut extractor try: self.donut_extractor = DonutMedicalExtractor() except Exception as e: logger.warning(f"Donut extractor not available: {str(e)}") self.donut_extractor = None def process_pdf(self, pdf_path: str, document_type: str = "unknown") -> ExtractionResult: """ Process medical PDF with multiple extraction methods Args: pdf_path: Path to PDF file document_type: Type of medical document Returns: ExtractionResult with structured data """ import time start_time = time.time() try: # Open PDF and extract basic info doc = fitz.open(pdf_path) page_count = len(doc) metadata = { "page_count": page_count, "pdf_metadata": doc.metadata, "file_size": os.path.getsize(pdf_path) } # Extract text using multiple methods raw_text = "" tables = [] images = [] for page_num in range(page_count): page = doc.load_page(page_num) # Extract text page_text = page.get_text() raw_text += f"\n--- Page {page_num + 1} ---\n{page_text}" # Extract tables using different methods page_tables = self._extract_tables(page) tables.extend(page_tables) # Extract images page_images = self._extract_images(page, pdf_path, page_num) images.extend(page_images) doc.close() # Determine extraction method based on content extraction_method = self._determine_extraction_method(raw_text, document_type) # Extract structured data based on document type if extraction_method == "donut" and self.donut_extractor: structured_data = self._extract_with_donut(pdf_path, document_type) else: structured_data = self._extract_with_fallback(raw_text, document_type) # Calculate confidence scores confidence_scores = self._calculate_extraction_confidence( raw_text, structured_data, tables, images ) processing_time = time.time() - start_time return ExtractionResult( raw_text=raw_text, structured_data=structured_data, confidence_scores=confidence_scores, extraction_method=extraction_method, processing_time=processing_time, tables_extracted=tables, images_extracted=images, metadata=metadata ) except Exception as e: logger.error(f"PDF processing error: {str(e)}") return ExtractionResult( raw_text="", structured_data={"error": str(e)}, confidence_scores={"overall": 0.0}, extraction_method="error", processing_time=time.time() - start_time, tables_extracted=[], images_extracted=[], metadata={"error": str(e)} ) def _determine_extraction_method(self, text: str, document_type: str) -> str: """Determine best extraction method based on content and type""" # High confidence cases for Donut if document_type in ["radiology", "ecg_report"] and len(text) > 500: return "donut" # Check for structured content indicators structured_indicators = [ "findings:", "impression:", "technique:", "results:", "normal ranges:", "reference values:", "patient information:" ] indicator_count = sum(1 for indicator in structured_indicators if indicator.lower() in text.lower()) if indicator_count >= 3 and len(text) > 1000: return "donut" # Fallback to text-based extraction return "fallback" def _extract_with_donut(self, pdf_path: str, document_type: str) -> Dict[str, Any]: """Extract structured data using Donut model""" if not self.donut_extractor: return self._extract_with_fallback("", document_type) try: # Convert PDF to images (first page for now, can be extended) images = self._pdf_to_images(pdf_path) if not images: return self._extract_with_fallback("", document_type) # Define task prompt based on document type task_prompts = { "radiology": "", "laboratory": "", "clinical_notes": "", "ecg_report": "", "unknown": "" } task_prompt = task_prompts.get(document_type, "") # Extract using Donut structured_data = self.donut_extractor.extract_from_image(images[0], task_prompt) # Post-process based on document type if document_type == "radiology": structured_data = self._postprocess_radiology(structured_data) elif document_type == "laboratory": structured_data = self._postprocess_laboratory(structured_data) elif document_type == "clinical_notes": structured_data = self._postprocess_clinical_notes(structured_data) elif document_type == "ecg_report": structured_data = self._postprocess_ecg(structured_data) return structured_data except Exception as e: logger.error(f"Donut extraction error: {str(e)}") return self._extract_with_fallback("", document_type) def _extract_with_fallback(self, text: str, document_type: str) -> Dict[str, Any]: """Fallback extraction using text processing and OCR if needed""" try: # Basic text cleaning cleaned_text = text.strip() # Document-type specific extraction if document_type == "radiology": return self._extract_radiology_from_text(cleaned_text) elif document_type == "laboratory": return self._extract_laboratory_from_text(cleaned_text) elif document_type == "clinical_notes": return self._extract_clinical_notes_from_text(cleaned_text) elif document_type == "ecg_report": return self._extract_ecg_from_text(cleaned_text) else: return { "raw_text": cleaned_text, "document_type": document_type, "extraction_method": "fallback_text" } except Exception as e: logger.error(f"Fallback extraction error: {str(e)}") return {"raw_text": text, "error": str(e), "extraction_method": "fallback"} def _extract_radiology_from_text(self, text: str) -> Dict[str, Any]: """Extract radiology information from text""" lines = text.split('\n') findings = [] impression = [] technique = [] current_section = None for line in lines: line = line.strip() if not line: continue line_lower = line.lower() if any(keyword in line_lower for keyword in ["findings:", "findings"]): current_section = "findings" continue elif any(keyword in line_lower for keyword in ["impression:", "impression", "conclusion:"]): current_section = "impression" continue elif any(keyword in line_lower for keyword in ["technique:", "protocol:"]): current_section = "technique" continue if current_section == "findings": findings.append(line) elif current_section == "impression": impression.append(line) elif current_section == "technique": technique.append(line) return { "findings": " ".join(findings), "impression": " ".join(impression), "technique": " ".join(technique), "document_type": "radiology", "extraction_method": "text_pattern_matching" } def _extract_laboratory_from_text(self, text: str) -> Dict[str, Any]: """Extract laboratory results from text""" lines = text.split('\n') tests = [] for line in lines: line = line.strip() if not line: continue # Look for test patterns # Pattern: Test Name Value Units Reference Range Flag parts = line.split() if len(parts) >= 3: # Try to identify test components test_data = { "raw_line": line, "potential_test": parts[0] if len(parts) > 0 else "", "potential_value": parts[1] if len(parts) > 1 else "", "potential_unit": parts[2] if len(parts) > 2 else "", } tests.append(test_data) return { "tests": tests, "document_type": "laboratory", "extraction_method": "text_pattern_matching" } def _extract_clinical_notes_from_text(self, text: str) -> Dict[str, Any]: """Extract clinical notes sections from text""" lines = text.split('\n') sections = {} current_section = "general" for line in lines: line = line.strip() if not line: continue line_lower = line.lower() # Identify section headers if any(keyword in line_lower for keyword in ["chief complaint:", "chief complaint", "cc:"]): current_section = "chief_complaint" continue elif any(keyword in line_lower for keyword in ["history of present illness:", "hpi:", "history:"]): current_section = "history_present_illness" continue elif any(keyword in line_lower for keyword in ["assessment:", "diagnosis:", "impression:"]): current_section = "assessment" continue elif any(keyword in line_lower for keyword in ["plan:", "treatment:", "recommendations:"]): current_section = "plan" continue # Add line to current section if current_section not in sections: sections[current_section] = [] sections[current_section].append(line) # Convert lists to text for section in sections: sections[section] = " ".join(sections[section]) return { "sections": sections, "document_type": "clinical_notes", "extraction_method": "text_pattern_matching" } def _extract_ecg_from_text(self, text: str) -> Dict[str, Any]: """Extract ECG information from text""" lines = text.split('\n') ecg_data = {} for line in lines: line = line.strip().lower() # Extract ECG measurements if "heart rate" in line or "hr" in line: import re hr_match = re.search(r'(\d+)', line) if hr_match: ecg_data["heart_rate"] = int(hr_match.group(1)) if "rhythm" in line: ecg_data["rhythm"] = line if any(interval in line for interval in ["pr interval", "qrs", "qt"]): ecg_data[line.split(':')[0]] = line return { "ecg_data": ecg_data, "document_type": "ecg_report", "extraction_method": "text_pattern_matching" } def _postprocess_radiology(self, data: Dict[str, Any]) -> Dict[str, Any]: """Post-process radiology extraction results""" # Ensure required fields exist if "findings" not in data: data["findings"] = "" if "impression" not in data: data["impression"] = "" data["document_type"] = "radiology" return data def _postprocess_laboratory(self, data: Dict[str, Any]) -> Dict[str, Any]: """Post-process laboratory extraction results""" # Ensure tests array exists if "tests" not in data: data["tests"] = [] data["document_type"] = "laboratory" return data def _postprocess_clinical_notes(self, data: Dict[str, Any]) -> Dict[str, Any]: """Post-process clinical notes extraction results""" # Ensure sections exist if "sections" not in data: data["sections"] = {} data["document_type"] = "clinical_notes" return data def _postprocess_ecg(self, data: Dict[str, Any]) -> Dict[str, Any]: """Post-process ECG extraction results""" # Ensure ecg_data exists if "ecg_data" not in data: data["ecg_data"] = {} data["document_type"] = "ecg_report" return data def _pdf_to_images(self, pdf_path: str) -> List[Image.Image]: """Convert PDF pages to images for Donut processing""" images = [] try: doc = fitz.open(pdf_path) for page_num in range(min(3, len(doc))): # Process first 3 pages page = doc.load_page(page_num) mat = fitz.Matrix(2.0, 2.0) # 2x zoom for better OCR pix = page.get_pixmap(matrix=mat) img_data = pix.tobytes("png") image = Image.open(io.BytesIO(img_data)) images.append(image) doc.close() except Exception as e: logger.error(f"PDF to image conversion error: {str(e)}") return images def _extract_tables(self, page) -> List[Dict[str, Any]]: """Extract tables from PDF page""" tables = [] try: # Use PyMuPDF table extraction if available tables_data = page.find_tables() for table in tables_data: table_dict = table.extract() tables.append({ "rows": len(table_dict), "columns": len(table_dict[0]) if table_dict else 0, "data": table_dict }) except Exception as e: logger.debug(f"Table extraction failed: {str(e)}") return tables def _extract_images(self, page, pdf_path: str, page_num: int) -> List[str]: """Extract images from PDF page""" images = [] try: image_list = page.get_images() for img_index, img in enumerate(image_list): xref = img[0] pix = fitz.Pixmap(page.parent, xref) if pix.n - pix.alpha < 4: # GRAY or RGB img_path = f"{Path(pdf_path).stem}_page{page_num+1}_img{img_index+1}.png" pix.save(img_path) images.append(img_path) pix = None except Exception as e: logger.debug(f"Image extraction failed: {str(e)}") return images def _calculate_extraction_confidence(self, raw_text: str, structured_data: Dict[str, Any], tables: List[Dict], images: List[str]) -> Dict[str, float]: """Calculate confidence scores for extraction quality""" confidence_scores = {} # Text extraction confidence text_length = len(raw_text.strip()) confidence_scores["text_extraction"] = min(1.0, text_length / 1000) if text_length > 0 else 0.0 # Structured data completeness required_fields = 0 present_fields = 0 if "findings" in structured_data or "impression" in structured_data: required_fields += 1 if structured_data.get("findings") or structured_data.get("impression"): present_fields += 1 if "tests" in structured_data: required_fields += 1 if structured_data.get("tests"): present_fields += 1 if "sections" in structured_data: required_fields += 1 if structured_data.get("sections"): present_fields += 1 confidence_scores["structural_completeness"] = present_fields / max(required_fields, 1) # Table extraction confidence confidence_scores["table_extraction"] = min(1.0, len(tables) * 0.3) # Image extraction confidence confidence_scores["image_extraction"] = min(1.0, len(images) * 0.2) # Overall confidence (weighted average) overall = ( 0.4 * confidence_scores["text_extraction"] + 0.4 * confidence_scores["structural_completeness"] + 0.1 * confidence_scores["table_extraction"] + 0.1 * confidence_scores["image_extraction"] ) confidence_scores["overall"] = overall return confidence_scores def convert_to_schema_format(self, extraction_result: ExtractionResult, document_type: str) -> Optional[Dict[str, Any]]: """Convert extraction result to canonical schema format""" try: # Create metadata metadata = MedicalDocumentMetadata( source_type=document_type, data_completeness=extraction_result.confidence_scores.get("overall", 0.0) ) # Create confidence score confidence = ConfidenceScore( extraction_confidence=extraction_result.confidence_scores.get("overall", 0.0), model_confidence=0.8, # Default assumption data_quality=extraction_result.confidence_scores.get("text_extraction", 0.0) ) # Convert based on document type if document_type == "radiology": return self._convert_to_radiology_schema(extraction_result, metadata, confidence) elif document_type == "laboratory": return self._convert_to_laboratory_schema(extraction_result, metadata, confidence) elif document_type == "clinical_notes": return self._convert_to_clinical_notes_schema(extraction_result, metadata, confidence) else: return None except Exception as e: logger.error(f"Schema conversion error: {str(e)}") return None def _convert_to_radiology_schema(self, result: ExtractionResult, metadata: MedicalDocumentMetadata, confidence: ConfidenceScore) -> Dict[str, Any]: """Convert to radiology schema format""" data = result.structured_data return { "metadata": metadata.dict(), "image_references": [], "findings": { "findings_text": data.get("findings", ""), "impression_text": data.get("impression", ""), "technique_description": data.get("technique", "") }, "segmentations": [], "metrics": {}, "confidence": confidence.dict(), "criticality_level": "routine", "follow_up_recommendations": [] } def _convert_to_laboratory_schema(self, result: ExtractionResult, metadata: MedicalDocumentMetadata, confidence: ConfidenceScore) -> Dict[str, Any]: """Convert to laboratory schema format""" data = result.structured_data return { "metadata": metadata.dict(), "tests": data.get("tests", []), "confidence": confidence.dict(), "critical_values": [], "abnormal_count": 0, "critical_count": 0 } def _convert_to_clinical_notes_schema(self, result: ExtractionResult, metadata: MedicalDocumentMetadata, confidence: ConfidenceScore) -> Dict[str, Any]: """Convert to clinical notes schema format""" data = result.structured_data sections = data.get("sections", {}) return { "metadata": metadata.dict(), "sections": [{"section_type": k, "content": v, "confidence": 0.8} for k, v in sections.items()], "entities": [], "confidence": confidence.dict() } # Export main classes __all__ = [ "MedicalPDFProcessor", "DonutMedicalExtractor", "ExtractionResult" ]