from __future__ import annotations import logging from pathlib import Path from qtpy.QtCore import Qt, QThreadPool, QTimer from qtpy.QtGui import QAction, QKeySequence from qtpy.QtWidgets import ( QButtonGroup, QDialog, QFileDialog, QHBoxLayout, QLabel, QMainWindow, QMessageBox, QProgressBar, QPushButton, QSplitter, QToolBar, QToolButton, QVBoxLayout, QWidget, ) from .dialogs import ( EditTextDialog, FontDialog, SettingsDialog, load_qsettings, save_qsettings, ) from .models import ( DocumentElement, DocumentModel, ElementType, ) from .services import PDFService from .utils import discover_tesseract, exception_text, free_space from .widgets import PageEditor, PreviewWidget from .workers import ( BaseWorker, DocxWorker, ImageWorker, OCRWorker, RenderWorker, TableWorker, ) logger = logging.getLogger(__name__) class MainWindow(QMainWindow): def __init__(self) -> None: super().__init__() self.setWindowTitle("PDF → DOCX OCR Editor") self.resize(1450, 950) self.thread_pool = QThreadPool.globalInstance() ( font, tesseract_path, language, dpi, preprocess, ) = load_qsettings() if not tesseract_path: tesseract_path = discover_tesseract() self.model = DocumentModel( font=font, tesseract_path=tesseract_path, ocr_language=language, render_dpi=dpi, preprocess=preprocess, ) self._current_render_worker: RenderWorker | None = None self._processing_workers: set[BaseWorker] = set() self._moving_old_rects: dict[str, tuple] = {} self._build_ui() self._build_menu() self._build_toolbar() self._build_shortcuts() self.statusBar().showMessage( "Select a PDF to begin." ) QTimer.singleShot(0, self.open_pdf) # ------------------------------------------------------------------ # UI # ------------------------------------------------------------------ def _build_ui(self) -> None: central = QWidget() self.setCentralWidget(central) main_layout = QVBoxLayout(central) main_layout.setContentsMargins(8, 8, 8, 8) # -------------------------------------------------------------- # PDF editor — full width, no preview panel and no title # -------------------------------------------------------------- self.pdf_editor = PageEditor( self.model, self, ) main_layout.addWidget(self.pdf_editor, 1) # -------------------------------------------------------------- # Bottom controls # -------------------------------------------------------------- bottom = QHBoxLayout() self.page_label = QLabel("Page 0 / 0") bottom.addWidget(self.page_label) self.progress = QProgressBar() self.progress.setRange(0, 100) self.progress.setValue(0) bottom.addWidget(self.progress, 1) self.prev_button = QPushButton("Previous") self.next_button = QPushButton("Next") self.generate_button = QPushButton("Generate DOCX") self.prev_button.clicked.connect( lambda: self.change_page(-1) ) self.next_button.clicked.connect( lambda: self.change_page(1) ) self.generate_button.clicked.connect( self.generate_docx ) bottom.addWidget(self.prev_button) bottom.addWidget(self.next_button) bottom.addWidget(self.generate_button) main_layout.addLayout(bottom) # -------------------------------------------------------------- # Editor signals # -------------------------------------------------------------- self.pdf_editor.selection_created.connect( self.on_selection_created ) self.pdf_editor.selection_changed.connect( self.on_selection_changed ) self.pdf_editor.selection_deleted.connect( self.delete_element ) self.pdf_editor.element_double_clicked.connect( self.edit_element ) self._apply_style() # self.preview = PreviewWidget( # self.model, # self, # ) def _apply_style(self) -> None: self.setStyleSheet( """ QMainWindow { background: #ffffff; } QWidget { color: #202124; font-size: 13px; } QToolBar { background: #ffffff; border: none; border-bottom: 1px solid #dadce0; spacing: 6px; padding: 6px; } QToolButton { background: #ffffff; border: 1px solid #dadce0; border-radius: 6px; padding: 8px 14px; color: #202124; } QToolButton:checked { background: #e8f0fe; border-color: #1a73e8; color: #174ea6; } QToolButton:hover { background: #f1f3f4; } QPushButton { background: #ffffff; border: 1px solid #dadce0; border-radius: 5px; padding: 7px 12px; color: #202124; } QPushButton:hover { background: #f1f3f4; } QProgressBar { border: 1px solid #dadce0; border-radius: 4px; text-align: center; background: #f1f3f4; color: #202124; } QProgressBar::chunk { background: #1a73e8; } QLabel#pageTitle { color: #5f6368; font-weight: bold; padding: 4px; } QStatusBar { background: #ffffff; color: #5f6368; border-top: 1px solid #dadce0; } QMenuBar { background: #ffffff; color: #202124; } QMenuBar::item { background: transparent; padding: 5px 8px; } QMenuBar::item:selected { background: #f1f3f4; } QMenu { background: #ffffff; color: #202124; border: 1px solid #dadce0; } QMenu::item { padding: 6px 24px; } QMenu::item:selected { background: #e8f0fe; color: #202124; } QLineEdit, QTextEdit, QPlainTextEdit, QSpinBox, QDoubleSpinBox, QComboBox { background: #ffffff; color: #202124; border: 1px solid #dadce0; border-radius: 4px; padding: 5px; } QLineEdit:focus, QTextEdit:focus, QPlainTextEdit:focus, QSpinBox:focus, QDoubleSpinBox:focus, QComboBox:focus { border: 1px solid #1a73e8; } QComboBox QAbstractItemView { background: #ffffff; color: #202124; selection-background-color: #e8f0fe; selection-color: #202124; } QCheckBox, QRadioButton { color: #202124; } QScrollBar:vertical { background: #f1f3f4; width: 12px; margin: 0; } QScrollBar::handle:vertical { background: #bdc1c6; border-radius: 6px; min-height: 30px; } QScrollBar::handle:vertical:hover { background: #9aa0a6; } QScrollBar:horizontal { background: #f1f3f4; height: 12px; margin: 0; } QScrollBar::handle:horizontal { background: #bdc1c6; border-radius: 6px; min-width: 30px; } QScrollBar::handle:horizontal:hover { background: #9aa0a6; } """ ) def _build_toolbar(self) -> None: """Build the main application toolbar.""" toolbar = QToolBar("Tools", self) toolbar.setMovable(False) toolbar.setToolButtonStyle( Qt.ToolButtonStyle.ToolButtonTextOnly ) self.addToolBar(toolbar) # ------------------------------------------------------------------ # Selection tools # ------------------------------------------------------------------ self.tool_group = QButtonGroup(self) self.tool_group.setExclusive(True) self.text_action = QAction("Text", self) self.text_action.setCheckable(True) self.text_action.setToolTip( "1 — Draw a text region and OCR it." ) self.image_action = QAction("Image", self) self.image_action.setCheckable(True) self.image_action.setToolTip( "2 — Draw an image region." ) self.table_action = QAction("Table", self) self.table_action.setCheckable(True) self.table_action.setToolTip( "3 — Draw a table region and recognize its cells." ) self.text_button = QToolButton(self) self.text_button.setDefaultAction(self.text_action) self.image_button = QToolButton(self) self.image_button.setDefaultAction(self.image_action) self.table_button = QToolButton(self) self.table_button.setDefaultAction(self.table_action) self.tool_group.addButton(self.text_button, 1) self.tool_group.addButton(self.image_button, 2) self.tool_group.addButton(self.table_button, 3) toolbar.addWidget(self.text_button) toolbar.addWidget(self.image_button) toolbar.addWidget(self.table_button) # ------------------------------------------------------------------ # Font / OCR settings # ------------------------------------------------------------------ toolbar.addSeparator() self.font_action = QAction("Font", self) self.font_action.setToolTip( "F — Select OCR output font." ) self.settings_action = QAction("OCR Settings", self) self.settings_action.setToolTip( "Configure OCR language, Tesseract path and processing options." ) toolbar.addAction(self.font_action) toolbar.addAction(self.settings_action) # ------------------------------------------------------------------ # File actions # ------------------------------------------------------------------ toolbar.addSeparator() open_action = QAction("Open PDF", self) open_action.setToolTip("Ctrl+O — Open a PDF document.") open_action.triggered.connect(self.open_pdf) save_action = QAction("Save Project", self) save_action.setToolTip("Ctrl+S — Save the current project.") save_action.triggered.connect(self.save_project) toolbar.addAction(open_action) toolbar.addAction(save_action) # ------------------------------------------------------------------ # Tool selection # ------------------------------------------------------------------ self.text_action.triggered.connect( lambda checked: self.set_tool( ElementType.TEXT if checked else None ) ) self.image_action.triggered.connect( lambda checked: self.set_tool( ElementType.IMAGE if checked else None ) ) self.table_action.triggered.connect( lambda checked: self.set_tool( ElementType.TABLE if checked else None ) ) # ------------------------------------------------------------------ # Other actions # ------------------------------------------------------------------ self.font_action.triggered.connect(self.choose_font) self.settings_action.triggered.connect(self.open_settings) # ------------------------------------------------------------------ # Default tool # ------------------------------------------------------------------ self.text_action.setChecked(True) self.set_tool(ElementType.TEXT) def _build_menu(self) -> None: menu = self.menuBar() file_menu = menu.addMenu("File") open_action = QAction("Open PDF", self) open_action.setShortcut(QKeySequence("Ctrl+O")) open_action.triggered.connect(self.open_pdf) file_menu.addAction(open_action) open_project = QAction("Open Project", self) open_project.triggered.connect(self.open_project) file_menu.addAction(open_project) save_project = QAction("Save Project", self) save_project.setShortcut(QKeySequence("Ctrl+S")) save_project.triggered.connect(self.save_project) file_menu.addAction(save_project) save_as = QAction("Save DOCX As…", self) save_as.setShortcut( QKeySequence("Ctrl+Shift+S") ) save_as.triggered.connect(self.save_docx_as) file_menu.addAction(save_as) file_menu.addSeparator() generate = QAction("Generate DOCX", self) generate.triggered.connect(self.generate_docx) file_menu.addAction(generate) file_menu.addSeparator() exit_action = QAction("Exit", self) exit_action.triggered.connect(self.close) file_menu.addAction(exit_action) edit_menu = menu.addMenu("Edit") delete_action = QAction("Delete Selected", self) delete_action.setShortcut(QKeySequence("Delete")) delete_action.triggered.connect( self.delete_selected ) edit_menu.addAction(delete_action) def _build_shortcuts(self) -> None: # Tool shortcuts are implemented through actions. self.text_action.setShortcut(QKeySequence("1")) self.image_action.setShortcut(QKeySequence("2")) self.table_action.setShortcut(QKeySequence("3")) self.font_action.setShortcut(QKeySequence("F")) from qtpy.QtGui import QShortcut QShortcut( QKeySequence(Qt.Key.Key_Left), self, activated=lambda: self.change_page(-1), ) QShortcut( QKeySequence(Qt.Key.Key_Right), self, activated=lambda: self.change_page(1), ) # ------------------------------------------------------------------ # PDF/project lifecycle # ------------------------------------------------------------------ def open_pdf(self) -> None: path, _ = QFileDialog.getOpenFileName( self, "Open PDF", str(Path.home()), "PDF files (*.pdf);;All files (*.*)", ) if not path: return self.load_pdf(path) def load_pdf(self, path: str) -> None: valid, error = PDFService.validate(path) if not valid: QMessageBox.critical( self, "Cannot Open PDF", error, ) return try: count = PDFService.page_count(path) pages = [] for page_number in range(1, count + 1): width, height = PDFService.page_size( path, page_number, ) from .models import PageState pages.append( PageState( page_number=page_number, width_pt=width, height_pt=height, ) ) output = str( Path(path).with_suffix(".docx") ) self.model = DocumentModel( source_pdf=path, output_docx=output, current_page=1, pages=pages, font=self.model.font, ocr_language=self.model.ocr_language, render_dpi=self.model.render_dpi, preprocess=self.model.preprocess, tesseract_path=self.model.tesseract_path, ) self.pdf_editor.model = self.model # self.preview.model = self.model self.progress.setValue(0) self.load_page(1) self.statusBar().showMessage( f"Loaded {count} pages." ) self._save_settings() except Exception as exc: logger.exception("Unable to load PDF") QMessageBox.critical( self, "PDF Error", exception_text(exc), ) def save_project(self) -> None: if not self.model.source_pdf: QMessageBox.information( self, "No Project", "Open a PDF first.", ) return default = ( Path(self.model.source_pdf) .with_suffix(".pdfocrproject") ) path, _ = QFileDialog.getSaveFileName( self, "Save OCR Project", str(default), "PDF OCR Project (*.pdfocrproject);;JSON (*.json)", ) if not path: return try: self.model.save_project(path) self.statusBar().showMessage( f"Project saved: {path}" ) except Exception as exc: logger.exception("Project save failed") QMessageBox.critical( self, "Save Error", exception_text(exc), ) def open_project(self) -> None: path, _ = QFileDialog.getOpenFileName( self, "Open OCR Project", str(Path.home()), "PDF OCR Project (*.pdfocrproject *.json)", ) if not path: return try: model = DocumentModel.load_project(path) if not Path(model.source_pdf).is_file(): QMessageBox.warning( self, "Source PDF Missing", ( "The project was loaded, but its source PDF " "could not be found:\n\n" f"{model.source_pdf}" ), ) return self.model = model self.pdf_editor.model = model # self.preview.model = model page = min( max(1, model.current_page), len(model.pages), ) self.load_page(page) except Exception as exc: logger.exception("Project load failed") QMessageBox.critical( self, "Project Error", exception_text(exc), ) # ------------------------------------------------------------------ # Page navigation # ------------------------------------------------------------------ def load_page(self, page_number: int) -> None: if not self.model.pages: return page_number = min( max(1, page_number), len(self.model.pages), ) self.model.current_page = page_number self.pdf_editor.set_page(page_number) # self.preview.set_page(page_number) self.page_label.setText( f"Page {page_number} / {len(self.model.pages)}" ) self.prev_button.setEnabled(page_number > 1) self.next_button.setEnabled( page_number < len(self.model.pages) ) self.progress.setValue( int( 100 * sum( page.processed for page in self.model.pages ) / max(1, len(self.model.pages)) ) ) self.statusBar().showMessage( f"Loading page {page_number}…" ) if self.model.source_pdf: worker = RenderWorker( self.model.source_pdf, page_number, 144, ) self._current_render_worker = worker worker.signals.result.connect( lambda data, p=page_number: ( self.on_page_rendered(p, data) ) ) worker.signals.error.connect( self.on_worker_error ) self.thread_pool.start(worker) def on_page_rendered( self, page_number: int, data: bytes, ) -> None: if page_number != self.model.current_page: return from qtpy.QtGui import QImage, QPixmap image = QImage.fromData(data) self.pdf_editor.set_pixmap( QPixmap.fromImage(image) ) self.statusBar().showMessage( f"Page {page_number} ready." ) def change_page(self, delta: int) -> None: if not self.model.pages: return target = self.model.current_page + delta if 1 <= target <= len(self.model.pages): self.load_page(target) # ------------------------------------------------------------------ # Tools # ------------------------------------------------------------------ def set_tool( self, tool: ElementType | None, ) -> None: self.pdf_editor.set_tool(tool) if tool: self.statusBar().showMessage( f"{tool.value.title()} tool selected." ) else: self.statusBar().showMessage( "Selection tool disabled." ) def choose_font(self) -> None: dialog = FontDialog( self.model.font, self, ) if dialog.exec() != QDialog.DialogCode.Accepted: return self.model.font = dialog.selected_font() # self.preview.update() self.pdf_editor.update() self._save_settings() def open_settings(self) -> None: dialog = SettingsDialog( self.model.tesseract_path, self.model.ocr_language, self.model.render_dpi, self.model.preprocess, self, ) if dialog.exec() != QDialog.DialogCode.Accepted: return ( path, language, dpi, preprocess, ) = dialog.values() self.model.tesseract_path = path self.model.ocr_language = language self.model.render_dpi = dpi self.model.preprocess = preprocess self._save_settings() self.statusBar().showMessage( "OCR settings updated." ) def _save_settings(self) -> None: save_qsettings( self.model.font, self.model.tesseract_path, self.model.ocr_language, self.model.render_dpi, self.model.preprocess, ) # ------------------------------------------------------------------ # Elements # ------------------------------------------------------------------ def on_selection_created( self, element: DocumentElement, ) -> None: page = self.model.page(element.page_number) page.elements.append(element) page.processed = True self.pdf_editor.selected_id = element.id # self.preview.update() self.pdf_editor.update() if element.element_type == ElementType.TEXT: self.start_ocr(element) elif element.element_type == ElementType.IMAGE: self.start_image(element) elif element.element_type == ElementType.TABLE: self.start_table(element) self._update_progress() def start_ocr( self, element: DocumentElement, ) -> None: if not self.model.tesseract_path: element.processing = False element.error = ( "Tesseract was not found. " "Configure it in OCR Settings." ) # self.preview.update() self.pdf_editor.update() QMessageBox.warning( self, "Tesseract Not Found", ( "Tesseract OCR was not found.\n\n" "Install Tesseract or configure its executable " "under OCR Settings." ), ) return worker = OCRWorker( self.model.source_pdf, element.page_number, element.rect(), self.model.render_dpi, self.model.tesseract_path, self.model.ocr_language, self.model.preprocess, ) self._register_worker(worker) worker.signals.result.connect( lambda text, eid=element.id: ( self.on_ocr_result(eid, text) ) ) worker.signals.error.connect( lambda error, eid=element.id: ( self.on_element_error(eid, error) ) ) self.statusBar().showMessage( "Processing OCR…" ) self.thread_pool.start(worker) def start_image( self, element: DocumentElement, ) -> None: worker = ImageWorker( self.model.source_pdf, element.page_number, element.rect(), max(200, self.model.render_dpi), ) self._register_worker(worker) worker.signals.result.connect( lambda data, eid=element.id: ( self.on_image_result(eid, data) ) ) worker.signals.error.connect( lambda error, eid=element.id: ( self.on_element_error(eid, error) ) ) self.statusBar().showMessage( "Extracting image…" ) self.thread_pool.start(worker) def start_table( self, element: DocumentElement, ) -> None: if not self.model.tesseract_path: element.processing = False element.error = "Tesseract is required for table OCR." # self.preview.update() self.pdf_editor.update() QMessageBox.warning( self, "Tesseract Not Found", "Configure Tesseract under OCR Settings.", ) return worker = TableWorker( self.model.source_pdf, element.page_number, element.rect(), max(220, self.model.render_dpi), self.model.tesseract_path, self.model.ocr_language, self.model.preprocess, ) self._register_worker(worker) worker.signals.result.connect( lambda result, eid=element.id: ( self.on_table_result(eid, result) ) ) worker.signals.error.connect( lambda error, eid=element.id: ( self.on_element_error(eid, error) ) ) self.statusBar().showMessage( "Recognizing table…" ) self.thread_pool.start(worker) def _register_worker(self, worker: BaseWorker) -> None: self._processing_workers.add(worker) worker.signals.finished.connect( lambda w=worker: self._processing_workers.discard(w) ) def on_ocr_result( self, element_id: str, text: str, ) -> None: element = self.model.element(element_id) if element is None: return element.processing = False element.error = "" element.text = text # self.preview.update() self.pdf_editor.update() self.statusBar().showMessage( "OCR completed." ) def on_image_result( self, element_id: str, data: bytes, ) -> None: element = self.model.element(element_id) if element is None: return element.processing = False element.error = "" element.image_png_b64 = ( DocumentModel.image_to_b64(data) ) # self.preview.update() self.pdf_editor.update() self.statusBar().showMessage( "Image extraction completed." ) def on_table_result( self, element_id: str, result, ) -> None: element = self.model.element(element_id) if element is None: return table, image_data = result element.processing = False element.error = "" if table is not None: element.table = table element.table_fallback_image = False element.image_png_b64 = "" self.statusBar().showMessage( "Table recognized." ) else: element.table = None element.table_fallback_image = True element.image_png_b64 = ( DocumentModel.image_to_b64(image_data) ) self.statusBar().showMessage( "Table recognition failed — inserted as image." ) # self.preview.update() self.pdf_editor.update() def on_element_error( self, element_id: str, error: str, ) -> None: element = self.model.element(element_id) if element is None: return element.processing = False element.error = error.splitlines()[0] # self.preview.update() self.pdf_editor.update() QMessageBox.warning( self, "Processing Error", error, ) def on_selection_changed( self, element_id: str, change, ) -> None: # During dragging, PageEditor sends the current rectangle. # On release it sends {"old": ..., "new": ...}. # self.preview.update() pass def edit_element( self, element_id: str, ) -> None: element = self.model.element(element_id) if element is None: return if element.element_type != ElementType.TEXT: return dialog = EditTextDialog( element.text, self, ) if dialog.exec() != QDialog.DialogCode.Accepted: return element.text = dialog.text() # self.preview.update() self.pdf_editor.update() def delete_selected(self) -> None: element_id = self.pdf_editor.selected_id if element_id: self.delete_element(element_id) def delete_element( self, element_id: str, ) -> None: for page in self.model.pages: for index, element in enumerate(page.elements): if element.id == element_id: del page.elements[index] if not page.elements: # A page with no elements is intentionally # represented by processed=False. page.processed = False self.pdf_editor.selected_id = "" self.pdf_editor.update() # self.preview.update() return # ------------------------------------------------------------------ # DOCX # ------------------------------------------------------------------ def default_output_path(self) -> Path: """Return a sensible automatic DOCX output path.""" if self.model.output_docx: output = Path(self.model.output_docx) if output.suffix.lower() == ".docx": return output if self.model.source_pdf: pdf_path = Path(self.model.source_pdf) return pdf_path.with_suffix(".docx") return Path.home() / "output.docx" def generate_docx(self) -> None: if not self.model.source_pdf: QMessageBox.information( self, "No PDF", "Open a PDF first.", ) return # Automatically determine output filename. output = self.default_output_path() # Remember the output path. self.model.output_docx = str(output) if output.exists(): answer = QMessageBox.question( self, "Overwrite DOCX?", ( f"The file already exists:\n\n{output}\n\n" "Overwrite it?" ), QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, QMessageBox.StandardButton.No, ) if answer != QMessageBox.StandardButton.Yes: return try: if free_space(output) < 50 * 1024 * 1024: QMessageBox.warning( self, "Low Disk Space", "Less than 50 MB of free disk space is available.", ) # Make sure the parent directory exists. output.parent.mkdir(parents=True, exist_ok=True) # Snapshot the model so editing remains independent # from the worker generating the document. snapshot = self.model.clone() worker = DocxWorker( snapshot, str(output), ) self._register_worker(worker) worker.signals.started.connect( lambda: self.statusBar().showMessage( f"Generating DOCX: {output}" ) ) worker.signals.result.connect( lambda _: self.on_docx_generated(output) ) worker.signals.error.connect( self.on_worker_error ) self.generate_button.setEnabled(False) worker.signals.finished.connect( lambda: self.generate_button.setEnabled(True) ) self.thread_pool.start(worker) except Exception as exc: logger.exception("DOCX generation setup failed") QMessageBox.critical( self, "DOCX Error", exception_text(exc), ) def on_docx_generated( self, output: Path, ) -> None: self.statusBar().showMessage( f"DOCX generated: {output}" ) QMessageBox.information( self, "DOCX Complete", f"Document generated successfully:\n\n{output}", ) def save_docx_as(self) -> None: default = self.default_output_path() path, _ = QFileDialog.getSaveFileName( self, "Save DOCX As", str(default), "Word document (*.docx)", ) if not path: return self.model.output_docx = path self.generate_docx() # ------------------------------------------------------------------ # Workers/errors # ------------------------------------------------------------------ def on_worker_error( self, message: str, ) -> None: logger.error("Worker error: %s", message) short = message.splitlines()[0] if message else "Unknown error" self.statusBar().showMessage( f"Background operation failed: {short}" ) QMessageBox.critical( self, "Background Processing Error", message, ) def _update_progress(self) -> None: if not self.model.pages: self.progress.setValue(0) return processed = sum( page.processed for page in self.model.pages ) self.progress.setValue( int(100 * processed / len(self.model.pages)) ) def closeEvent(self, event) -> None: if self._processing_workers: answer = QMessageBox.question( self, "Operations Running", ( "Background operations are still running.\n\n" "Exit anyway?" ), QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, QMessageBox.StandardButton.No, ) if answer != QMessageBox.StandardButton.Yes: event.ignore() return for worker in list( self._processing_workers ): try: worker.cancel() except Exception: logger.exception( "Worker cancellation failed" ) self._save_settings() event.accept()