""" Main OCR Processor Orchestrates the complete OCR processing pipeline """ import time from datetime import datetime import tempfile import os import fitz # PyMuPDF from pathlib import Path from typing import Optional, Callable from dataclasses import dataclass import threading from logger import ProcessingLogger from ocr_engine import OCREngine from text_processor import ContentFormatter from utils import validate_page_ranges @dataclass class ProcessingResult: """Container for processing results.""" content: str output_file: Optional[str] status: str logs: str success: bool vision_calls_used: int = 0 processing_time: float = 0.0 pages_processed: int = 0 class DocumentProcessor: """Main document processor orchestrating the OCR pipeline.""" def __init__(self): self.logger = ProcessingLogger() self.ocr_engine = OCREngine(self.logger) self.content_formatter = ContentFormatter(self.logger) # --- Abort state (instance-scoped) --- self._abort_flag = threading.Event() # ---------- Abort control ---------- def abort_processing(self) -> None: """Signal that processing should be aborted.""" self._abort_flag.set() self.logger.log_section("Processing Aborted") self.logger.log_metric("Status", "Aborted by user") def clear_abort(self) -> None: """Clear the abort flag (start fresh for a new run).""" self._abort_flag.clear() def is_abort_requested(self) -> bool: """Check whether the user has requested an abort.""" return self._abort_flag.is_set() # ----------------------------------- # Logs as a single string (Textbox-safe) def _logs_text(self) -> str: logs = self.logger.get_logs() if isinstance(logs, list): return "\n".join(str(x) for x in logs) return str(logs or "") def save_output(self, pdf_path: Path, content: str) -> Optional[str]: """Save processed content to file.""" try: base_name = pdf_path.stem now = datetime.now().strftime("%Y%m%d_%H%M%S") filename = f"{base_name}_{now}.md" temp_dir = tempfile.gettempdir() output_path = os.path.join(temp_dir, filename) with open(output_path, "w", encoding="utf-8") as f: f.write(content) self.logger.log_success(f"Saved: {filename}") return output_path except Exception as e: self.logger.log_error(f"Save failed: {e}") return None def process_document(self, uploaded_file, page_ranges_str: Optional[str] = None, progress_callback: Optional[Callable] = None) -> ProcessingResult: """ Process a PDF or Markdown document through the processing pipeline. """ start_time = time.time() # Fresh run => ensure abort flag is clear self.clear_abort() try: if not uploaded_file: return ProcessingResult( content="Please upload a PDF, Markdown, or TXT file.", output_file=None, status="No file", logs=self._logs_text(), success=False ) file_path = Path(uploaded_file.name) self.logger.log_section(f"Processing: {file_path.name}") if progress_callback: progress_callback(f"📄 Processing: {file_path.name}") # Markdown / Text if file_path.suffix.lower() in ['.md', '.markdown', '.txt']: if file_path.suffix.lower() == '.txt': return self._process_txt_file(file_path, progress_callback, start_time) else: return self._process_markdown_file(file_path, progress_callback, start_time) # PDF path with fitz.open(file_path) as doc: total_pages = len(doc) self.logger.log_metric("Total pages in document", total_pages) # Page ranges if page_ranges_str and page_ranges_str.strip(): is_valid, error_msg, pages_to_process = validate_page_ranges(page_ranges_str, total_pages) if not is_valid: return ProcessingResult( content=f"Invalid page ranges: {error_msg}", output_file=None, status="Invalid page ranges", logs=self._logs_text(), success=False ) page_numbers = [p + 1 for p in pages_to_process] # back to 1-indexed self.logger.log_metric("Pages to process", f"{len(page_numbers)} pages: {page_ranges_str}") else: page_numbers = list(range(1, total_pages + 1)) self.logger.log_metric("Pages to process", f"All {total_pages} pages") # Extract text page_texts = {} for i, page_no in enumerate(page_numbers): if self.is_abort_requested(): return ProcessingResult( content="Processing was aborted by user", output_file=None, status="Aborted", logs=self._logs_text(), success=False, vision_calls_used=self.ocr_engine.get_vision_calls_used(), processing_time=time.time() - start_time, pages_processed=i ) if progress_callback: progress_callback(f"📖 Processing page {page_no} ({i+1}/{len(page_numbers)})") page = doc[page_no - 1] # fitz is 0-indexed text = self.ocr_engine.extract_page_text(page, page_no) page_texts[page_no] = text # Formatting self.logger.log_section("Content Formatting") if progress_callback: progress_callback("📝 Formatting content...") formatted_pages = [] document_title = self._extract_document_title(file_path.name) for page_no in sorted(page_texts.keys()): if self.is_abort_requested(): return ProcessingResult( content="Processing was aborted by user", output_file=None, status="Aborted", logs=self._logs_text(), success=False, vision_calls_used=self.ocr_engine.get_vision_calls_used(), processing_time=time.time() - start_time, pages_processed=len(page_texts) ) if len(page_texts[page_no].strip()) >= 10: formatted = self.content_formatter.format_content( page_texts[page_no], page_no, document_title ) formatted_pages.append(formatted) # Assemble self.logger.log_section("Document Assembly") header = self.content_formatter.build_document_header(document_title) if page_ranges_str and page_ranges_str.strip(): header += f"\n\n**Pages Processed:** {page_ranges_str}" final_content = f"{header}\n\n---\n\n" + "\n\n---\n\n".join(formatted_pages) # Save output_file = self.save_output(file_path, final_content) # Metrics processing_time = time.time() - start_time vision_calls = self.ocr_engine.get_vision_calls_used() self.logger.log_section("Processing Complete") self.logger.log_metric("Processing time", f"{processing_time:.1f}s") self.logger.log_metric("Vision calls used", vision_calls) self.logger.log_metric("Pages processed", len(formatted_pages)) self.logger.log_metric("Total words", len(final_content.split())) if progress_callback: progress_callback("✅ Complete!") return ProcessingResult( content=final_content, output_file=output_file, status="Complete", logs=self._logs_text(), success=True, vision_calls_used=vision_calls, processing_time=processing_time, pages_processed=len(formatted_pages) ) except Exception as e: processing_time = time.time() - start_time error_msg = f"Processing error: {str(e)}" self.logger.log_error(error_msg) return ProcessingResult( content=error_msg, output_file=None, status="Error", logs=self._logs_text(), success=False, vision_calls_used=self.ocr_engine.get_vision_calls_used(), processing_time=processing_time, pages_processed=0 ) def _process_markdown_file(self, file_path: Path, progress_callback: Optional[Callable], start_time: float) -> ProcessingResult: """Process a markdown file - apply formatting only, no OCR needed.""" try: self.logger.log_section("Markdown Processing") if progress_callback: progress_callback("📝 Reading markdown file...") with open(file_path, 'r', encoding='utf-8') as f: markdown_content = f.read() self.logger.log_metric("File size", f"{len(markdown_content)} characters") self.logger.log_metric("File type", "Markdown") document_title = self._extract_document_title(file_path.name) self.logger.log_section("Content Formatting") MAX_CHARS_PER_CHUNK = 100000 # ~25k tokens chunks_processed = 1 if len(markdown_content) <= MAX_CHARS_PER_CHUNK: if progress_callback: progress_callback("🎨 Formatting content...") formatted_content = self.content_formatter.format_content( markdown_content, 1, document_title ) else: chunks = self._split_text_into_chunks(markdown_content, MAX_CHARS_PER_CHUNK) chunks_processed = len(chunks) self.logger.log_metric("Chunks to process", chunks_processed) formatted_chunks = [] for i, chunk in enumerate(chunks, 1): if progress_callback: progress_callback(f"🎨 Formatting chunk {i}/{chunks_processed}...") formatted_chunks.append( self.content_formatter.format_content(chunk, i, document_title) ) formatted_content = "\n\n---\n\n".join(formatted_chunks) header = self.content_formatter.build_document_header(document_title) final_content = f"{header}\n\n---\n\n{formatted_content}" output_file = self.save_output(file_path, final_content) processing_time = time.time() - start_time self.logger.log_section("Processing Complete") self.logger.log_metric("Processing time", f"{processing_time:.1f}s") self.logger.log_metric("Total characters", len(final_content)) if progress_callback: progress_callback("✅ Complete!") return ProcessingResult( content=final_content, output_file=output_file, status="Complete", logs=self._logs_text(), success=True, vision_calls_used=0, processing_time=processing_time, pages_processed=chunks_processed ) except Exception as e: return ProcessingResult( content=f"Markdown processing error: {str(e)}", output_file=None, status="Error", logs=self._logs_text(), success=False, vision_calls_used=0, processing_time=time.time() - start_time, pages_processed=0 ) def _process_txt_file(self, file_path: Path, progress_callback: Optional[Callable], start_time: float) -> ProcessingResult: """Process a text file - apply formatting only, no OCR needed.""" try: self.logger.log_section("Text Processing") if progress_callback: progress_callback("📝 Reading text file...") with open(file_path, 'r', encoding='utf-8') as f: text_content = f.read() self.logger.log_metric("File size", f"{len(text_content)} characters") self.logger.log_metric("File type", "Text") document_title = self._extract_document_title(file_path.name) self.logger.log_section("Content Formatting") MAX_CHARS_PER_CHUNK = 100000 chunks_processed = 1 if len(text_content) <= MAX_CHARS_PER_CHUNK: if progress_callback: progress_callback("🎨 Formatting content...") formatted_content = self.content_formatter.format_content( text_content, 1, document_title ) else: chunks = self._split_text_into_chunks(text_content, MAX_CHARS_PER_CHUNK) chunks_processed = len(chunks) self.logger.log_metric("Chunks to process", chunks_processed) formatted_chunks = [] for i, chunk in enumerate(chunks, 1): if progress_callback: progress_callback(f"🎨 Formatting chunk {i}/{chunks_processed}...") formatted_chunks.append( self.content_formatter.format_content(chunk, i, document_title) ) formatted_content = "\n\n---\n\n".join(formatted_chunks) header = self.content_formatter.build_document_header(document_title) final_content = f"{header}\n\n---\n\n{formatted_content}" output_file = self.save_output(file_path, final_content) processing_time = time.time() - start_time self.logger.log_section("Processing Complete") self.logger.log_metric("Processing time", f"{processing_time:.1f}s") self.logger.log_metric("Total characters", len(final_content)) if progress_callback: progress_callback("✅ Complete!") return ProcessingResult( content=final_content, output_file=output_file, status="Complete", logs=self._logs_text(), success=True, vision_calls_used=0, processing_time=processing_time, pages_processed=chunks_processed ) except Exception as e: return ProcessingResult( content=f"Text processing error: {str(e)}", output_file=None, status="Error", logs=self._logs_text(), success=False, vision_calls_used=0, processing_time=time.time() - start_time, pages_processed=0 ) def _split_text_into_chunks(self, text: str, max_chunk_size: int) -> list: """Split text into chunks at logical boundaries.""" chunks = [] paragraphs = text.split('\n\n') current_chunk = "" for para in paragraphs: if len(current_chunk) + len(para) + 2 > max_chunk_size: if current_chunk: chunks.append(current_chunk.strip()) current_chunk = para else: # Break a huge paragraph by sentences sentences = para.replace('. ', '.\n').split('\n') for sent in sentences: if len(current_chunk) + len(sent) + 1 > max_chunk_size: if current_chunk: chunks.append(current_chunk.strip()) current_chunk = sent else: current_chunk += (" " + sent) if current_chunk else sent else: current_chunk += ("\n\n" + para) if current_chunk else para if current_chunk: chunks.append(current_chunk.strip()) return chunks def _extract_document_title(self, filename: str) -> str: """Extract a clean document title from filename.""" title = os.path.splitext(os.path.basename(filename))[0] title = title.replace('_', ' ').replace('-', ' ') title = ' '.join(word.capitalize() for word in title.split()) return title if title else "Document" def add_log_callback(self, callback: Callable[[str], None]) -> None: """Add a callback for real-time log updates.""" self.logger.add_callback(callback) def clear_logs(self) -> None: """Clear all logs and reset counters.""" self.logger.clear() self.ocr_engine.reset_vision_counter()