Deploy pdf_extractor.py to backend/ directory

backend/pdf_extractor.py

@@ -0,0 +1,670 @@
+"""
+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 = "<s_rvlcdip>"
+        
+        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": "<s_radiology_report>",
+                "laboratory": "<s_laboratory_report>",
+                "clinical_notes": "<s_clinical_note>",
+                "ecg_report": "<s_ecg_report>",
+                "unknown": "<s_medical_document>"
+            }
+            
+            task_prompt = task_prompts.get(document_type, "<s_medical_document>")
+            
+            # 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"
+]