import sys import os import numpy as np import cv2 import gc from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QPushButton, QLabel, QFileDialog, QTableWidget, QTableWidgetItem, QHeaderView, QSplitter, QFrame, QProgressBar, QMessageBox, QComboBox, QScrollArea, QCheckBox, QGroupBox, QRadioButton, QButtonGroup, QSlider, QDialog, QListWidget, QDialogButtonBox, QTabWidget, QGridLayout, QLineEdit, QSpinBox, QFormLayout, QColorDialog) from PyQt5.QtCore import Qt, QThread, pyqtSignal, QSize, QPoint, QRect, pyqtSlot from PyQt5.QtGui import QImage, QPixmap, QFont, QColor, QPainter, QPen, QBrush, QWheelEvent, QMouseEvent from paddleocr import LayoutDetection import fitz # PyMuPDF import tempfile import shutil from concurrent.futures import ThreadPoolExecutor import queue import time class ScrollableImageLabel(QScrollArea): """可缩放和拖动的图像显示控件""" def __init__(self, parent=None): super().__init__(parent) self.image_label = QLabel() self.image_label.setAlignment(Qt.AlignCenter) self.setWidget(self.image_label) self.setWidgetResizable(True) self.setStyleSheet("border: 1px solid #ddd; border-radius: 6px;") # 图像和缩放相关变量 self.original_pixmap = None self.scale_factor = 1.0 self.drag_position = QPoint() self.is_dragging = False # 设置鼠标跟踪 self.setMouseTracking(True) self.image_label.setMouseTracking(True) self.setCursor(Qt.OpenHandCursor) # 连接鼠标事件 self.image_label.mousePressEvent = self.mousePressEvent self.image_label.mouseReleaseEvent = self.mouseReleaseEvent self.image_label.mouseMoveEvent = self.mouseMoveEvent self.image_label.wheelEvent = self.wheelEvent def set_image(self, img_array): """设置图像""" # 确保图像是RGB格式 if len(img_array.shape) == 2: # 灰度图 img_array = cv2.cvtColor(img_array, cv2.COLOR_GRAY2RGB) elif img_array.shape[2] == 4: # RGBA img_array = cv2.cvtColor(img_array, cv2.COLOR_RGBA2RGB) elif img_array.shape[2] == 3: # BGR (OpenCV默认) img_array = cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB) height, width, channel = img_array.shape bytes_per_line = 3 * width q_img = QImage(img_array.data, width, height, bytes_per_line, QImage.Format_RGB888) # 保存原始图像 self.original_pixmap = QPixmap.fromImage(q_img) # 适应窗口大小显示 self.fit_to_window() def fit_to_window(self): """适应窗口大小显示图像""" if not self.original_pixmap: return # 计算缩放比例以适应窗口 window_size = self.viewport().size() pixmap_size = self.original_pixmap.size() width_ratio = window_size.width() / pixmap_size.width() height_ratio = window_size.height() / pixmap_size.height() # 取较小的比例，确保图像完全显示在窗口内 self.scale_factor = min(width_ratio, height_ratio) # 缩放图像 self.update_scaled_pixmap() def update_scaled_pixmap(self): """更新缩放后的图像""" if not self.original_pixmap: return # 计算新尺寸 new_size = self.original_pixmap.size() * self.scale_factor # 缩放图像 scaled_pixmap = self.original_pixmap.scaled( new_size, Qt.KeepAspectRatio, Qt.SmoothTransformation ) # 更新显示 self.image_label.setPixmap(scaled_pixmap) self.image_label.setFixedSize(scaled_pixmap.size()) def set_scale(self, scale): """设置缩放比例""" if not self.original_pixmap: return # 限制缩放范围 self.scale_factor = max(0.1, min(5.0, scale)) # 更新缩放后的图像 self.update_scaled_pixmap() def zoom_in(self): """放大图像""" self.set_scale(self.scale_factor * 1.25) def zoom_out(self): """缩小图像""" self.set_scale(self.scale_factor * 0.8) def mousePressEvent(self, event: QMouseEvent): """鼠标按下事件""" if event.button() == Qt.LeftButton: # 开始拖动 self.is_dragging = True self.drag_position = event.pos() self.setCursor(Qt.ClosedHandCursor) def mouseReleaseEvent(self, event: QMouseEvent): """鼠标释放事件""" if event.button() == Qt.LeftButton: # 结束拖动 self.is_dragging = False self.setCursor(Qt.OpenHandCursor) def mouseMoveEvent(self, event: QMouseEvent): """鼠标移动事件""" if self.is_dragging: # 计算移动距离 delta = event.pos() - self.drag_position # 移动滚动条 self.horizontalScrollBar().setValue(self.horizontalScrollBar().value() - delta.x()) self.verticalScrollBar().setValue(self.verticalScrollBar().value() - delta.y()) # 更新拖动位置 self.drag_position = event.pos() def wheelEvent(self, event: QWheelEvent): """鼠标滚轮事件，用于缩放""" # 获取滚轮角度 angle_delta = event.angleDelta().y() # 根据滚轮方向调整缩放比例 if angle_delta > 0: self.zoom_in() else: self.zoom_out() class PageSelectionDialog(QDialog): """页面选择对话框""" def __init__(self, total_pages, current_page, parent=None): super().__init__(parent) self.setWindowTitle("选择页面") self.setMinimumWidth(400) self.setMinimumHeight(300) layout = QVBoxLayout() # 选择模式 mode_group = QGroupBox("选择模式") mode_layout = QVBoxLayout() self.all_pages_radio = QRadioButton("所有页面") self.all_pages_radio.setChecked(True) mode_layout.addWidget(self.all_pages_radio) self.current_page_radio = QRadioButton("当前页面") mode_layout.addWidget(self.current_page_radio) self.select_pages_radio = QRadioButton("选择页面") mode_layout.addWidget(self.select_pages_radio) mode_group.setLayout(mode_layout) layout.addWidget(mode_group) # 页面列表 list_widget = QWidget() list_layout = QVBoxLayout() # 页面选择网格 self.page_grid = QGridLayout() self.page_buttons = [] # 添加页面按钮 rows = (total_pages + 4) // 5 # 每行5个按钮 for i in range(total_pages): btn = QPushButton(f"第 {i + 1} 页") btn.setCheckable(True) btn.setEnabled(False) # 默认选中当前页 if i == current_page: btn.setChecked(True) self.page_buttons.append(btn) row = i // 5 col = i % 5 self.page_grid.addWidget(btn, row, col) list_layout.addLayout(self.page_grid) list_widget.setLayout(list_layout) # 创建滚动区域 scroll = QScrollArea() scroll.setWidget(list_widget) scroll.setWidgetResizable(True) layout.addWidget(scroll) # 全选/取消全选按钮 select_buttons = QWidget() select_layout = QHBoxLayout() self.select_all_btn = QPushButton("全选") self.select_all_btn.setEnabled(False) self.select_all_btn.clicked.connect(self.select_all_pages) select_layout.addWidget(self.select_all_btn) self.deselect_all_btn = QPushButton("取消全选") self.deselect_all_btn.setEnabled(False) self.deselect_all_btn.clicked.connect(self.deselect_all_pages) select_layout.addWidget(self.deselect_all_btn) select_layout.addStretch() select_buttons.setLayout(select_layout) layout.addWidget(select_buttons) # 连接信号 self.all_pages_radio.toggled.connect(self.on_mode_changed) self.current_page_radio.toggled.connect(self.on_mode_changed) self.select_pages_radio.toggled.connect(self.on_mode_changed) # 按钮 buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel) buttons.accepted.connect(self.accept) buttons.rejected.connect(self.reject) layout.addWidget(buttons) self.setLayout(layout) def on_mode_changed(self): """选择模式改变""" is_select_mode = self.select_pages_radio.isChecked() # 启用/禁用页面按钮 for btn in self.page_buttons: btn.setEnabled(is_select_mode) # 启用/禁用全选按钮 self.select_all_btn.setEnabled(is_select_mode) self.deselect_all_btn.setEnabled(is_select_mode) def select_all_pages(self): """全选页面""" for btn in self.page_buttons: btn.setChecked(True) def deselect_all_pages(self): """取消全选页面""" for btn in self.page_buttons: btn.setChecked(False) def get_selected_pages(self): """获取选中的页面""" if self.all_pages_radio.isChecked(): return "all" elif self.current_page_radio.isChecked(): return "current" else: selected_pages = [] for i, btn in enumerate(self.page_buttons): if btn.isChecked(): selected_pages.append(i) return selected_pages if selected_pages else None class ImageExportOptionsDialog(QDialog): """图像导出选项对话框""" def __init__(self, block_types, parent=None): super().__init__(parent) self.setWindowTitle("导出区块图像选项") self.setMinimumWidth(400) layout = QVBoxLayout() # 区块类型选择 type_group = QGroupBox("选择要导出的区块类型") type_layout = QVBoxLayout() # 全选/取消全选 select_layout = QHBoxLayout() self.select_all_btn = QPushButton("全选") self.select_all_btn.clicked.connect(self.select_all) select_layout.addWidget(self.select_all_btn) self.deselect_all_btn = QPushButton("取消全选") self.deselect_all_btn.clicked.connect(self.deselect_all) select_layout.addWidget(self.deselect_all_btn) select_layout.addStretch() type_layout.addLayout(select_layout) # 区块类型复选框 self.type_checkboxes = [] for block_type in block_types: checkbox = QCheckBox(block_type) checkbox.setChecked(True) # 默认全选 type_layout.addWidget(checkbox) self.type_checkboxes.append(checkbox) type_group.setLayout(type_layout) layout.addWidget(type_group) # 页面选择 page_group = QGroupBox("页面选择") page_layout = QVBoxLayout() self.all_pages_radio = QRadioButton("所有页面") self.all_pages_radio.setChecked(True) page_layout.addWidget(self.all_pages_radio) self.current_page_radio = QRadioButton("当前页面") page_layout.addWidget(self.current_page_radio) self.select_pages_radio = QRadioButton("选择页面") page_layout.addWidget(self.select_pages_radio) page_group.setLayout(page_layout) layout.addWidget(page_group) # 文本布局选项 layout_group = QGroupBox("文本布局") layout_layout = QVBoxLayout() self.single_column_radio = QRadioButton("单栏文本") self.single_column_radio.setChecked(True) layout_layout.addWidget(self.single_column_radio) self.double_column_radio = QRadioButton("双栏文本") layout_layout.addWidget(self.double_column_radio) self.triple_column_radio = QRadioButton("三栏文本") layout_layout.addWidget(self.triple_column_radio) layout_group.setLayout(layout_layout) layout.addWidget(layout_group) # 导出格式选项 format_group = QGroupBox("导出格式") format_layout = QVBoxLayout() self.export_image_radio = QRadioButton("图像格式") self.export_image_radio.setChecked(True) format_layout.addWidget(self.export_image_radio) self.export_pdf_radio = QRadioButton("PDF格式") format_layout.addWidget(self.export_pdf_radio) format_group.setLayout(format_layout) layout.addWidget(format_group) # 图像格式选项（仅当导出为图像时显示） self.image_format_group = QGroupBox("图像格式") image_format_layout = QFormLayout() self.format_combo = QComboBox() self.format_combo.addItems(["JPEG", "PNG"]) image_format_layout.addRow("图像格式:", self.format_combo) # 图像质量选项 self.quality_slider = QSlider(Qt.Horizontal) self.quality_slider.setMinimum(1) self.quality_slider.setMaximum(100) self.quality_slider.setValue(90) self.quality_label = QLabel("90%") quality_layout = QHBoxLayout() quality_layout.addWidget(self.quality_slider) quality_layout.addWidget(self.quality_label) self.quality_slider.valueChanged.connect(lambda v: self.quality_label.setText(f"{v}%")) image_format_layout.addRow("图像质量:", quality_layout) # 尺寸限制选项 self.limit_size_checkbox = QCheckBox("限制最大尺寸") self.limit_size_checkbox.setChecked(False) image_format_layout.addRow(self.limit_size_checkbox) size_layout = QHBoxLayout() self.max_width_spin = QSpinBox() self.max_width_spin.setRange(100, 5000) self.max_width_spin.setValue(1920) self.max_width_spin.setEnabled(False) size_layout.addWidget(QLabel("最大宽度:")) size_layout.addWidget(self.max_width_spin) size_layout.addWidget(QLabel("像素")) self.max_height_spin = QSpinBox() self.max_height_spin.setRange(100, 5000) self.max_height_spin.setValue(1080) self.max_height_spin.setEnabled(False) size_layout.addWidget(QLabel("最大高度:")) size_layout.addWidget(self.max_height_spin) size_layout.addWidget(QLabel("像素")) image_format_layout.addRow(size_layout) self.limit_size_checkbox.toggled.connect( lambda checked: ( self.max_width_spin.setEnabled(checked), self.max_height_spin.setEnabled(checked) ) ) self.image_format_group.setLayout(image_format_layout) layout.addWidget(self.image_format_group) # PDF压缩选项（仅当导出为PDF时显示） self.pdf_format_group = QGroupBox("PDF压缩选项") pdf_format_layout = QFormLayout() # PDF图像质量选项 self.pdf_quality_slider = QSlider(Qt.Horizontal) self.pdf_quality_slider.setMinimum(1) self.pdf_quality_slider.setMaximum(100) self.pdf_quality_slider.setValue(75) self.pdf_quality_label = QLabel("75%") pdf_quality_layout = QHBoxLayout() pdf_quality_layout.addWidget(self.pdf_quality_slider) pdf_quality_layout.addWidget(self.pdf_quality_label) self.pdf_quality_slider.valueChanged.connect(lambda v: self.pdf_quality_label.setText(f"{v}%")) pdf_format_layout.addRow("图像质量:", pdf_quality_layout) # PDF尺寸限制选项 self.pdf_limit_size_checkbox = QCheckBox("限制最大尺寸") self.pdf_limit_size_checkbox.setChecked(False) pdf_format_layout.addRow(self.pdf_limit_size_checkbox) pdf_size_layout = QHBoxLayout() self.pdf_max_width_spin = QSpinBox() self.pdf_max_width_spin.setRange(100, 5000) self.pdf_max_width_spin.setValue(1920) self.pdf_max_width_spin.setEnabled(False) pdf_size_layout.addWidget(QLabel("最大宽度:")) pdf_size_layout.addWidget(self.pdf_max_width_spin) pdf_size_layout.addWidget(QLabel("像素")) self.pdf_max_height_spin = QSpinBox() self.pdf_max_height_spin.setRange(100, 5000) self.pdf_max_height_spin.setValue(1080) self.pdf_max_height_spin.setEnabled(False) pdf_size_layout.addWidget(QLabel("最大高度:")) pdf_size_layout.addWidget(self.pdf_max_height_spin) pdf_size_layout.addWidget(QLabel("像素")) pdf_format_layout.addRow(pdf_size_layout) self.pdf_limit_size_checkbox.toggled.connect( lambda checked: ( self.pdf_max_width_spin.setEnabled(checked), self.pdf_max_height_spin.setEnabled(checked) ) ) self.pdf_format_group.setLayout(pdf_format_layout) layout.addWidget(self.pdf_format_group) # 连接信号，根据导出格式显示/隐藏相应选项 self.export_image_radio.toggled.connect(self.image_format_group.setVisible) self.export_pdf_radio.toggled.connect(self.pdf_format_group.setVisible) self.export_image_radio.toggled.connect(lambda checked: self.pdf_format_group.setHidden(checked)) self.export_pdf_radio.toggled.connect(lambda checked: self.image_format_group.setHidden(checked)) # 输出目录选择 dir_group = QGroupBox("输出目录") dir_layout = QVBoxLayout() path_layout = QHBoxLayout() self.output_dir_edit = QLineEdit() path_layout.addWidget(self.output_dir_edit) self.browse_dir_btn = QPushButton("浏览...") self.browse_dir_btn.clicked.connect(self.browse_output_dir) path_layout.addWidget(self.browse_dir_btn) dir_layout.addLayout(path_layout) dir_group.setLayout(dir_layout) layout.addWidget(dir_group) # 按钮 buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel) buttons.accepted.connect(self.accept) buttons.rejected.connect(self.reject) layout.addWidget(buttons) self.setLayout(layout) # 初始化变量 self.output_dir = "" # 初始状态设置 self.image_format_group.setVisible(True) self.pdf_format_group.setVisible(False) def select_all(self): """全选""" for checkbox in self.type_checkboxes: checkbox.setChecked(True) def deselect_all(self): """取消全选""" for checkbox in self.type_checkboxes: checkbox.setChecked(False) def browse_output_dir(self): """选择输出目录""" dir_path = QFileDialog.getExistingDirectory(self, "选择输出目录") if dir_path: self.output_dir = dir_path self.output_dir_edit.setText(dir_path) def get_export_options(self): """获取导出选项""" # 获取选中的区块类型 selected_types = [] for checkbox in self.type_checkboxes: if checkbox.isChecked(): selected_types.append(checkbox.text()) # 获取页面选择 if self.all_pages_radio.isChecked(): page_selection = "all" elif self.current_page_radio.isChecked(): page_selection = "current" else: page_selection = "select" # 实际页面选择将在主窗口中处理 # 获取文本布局 if self.single_column_radio.isChecked(): layout_type = "single" elif self.double_column_radio.isChecked(): layout_type = "double" else: layout_type = "triple" # 获取导出格式 export_format = "image" if self.export_image_radio.isChecked() else "pdf" # 获取图像格式和质量（仅当导出为图像时） image_format = None image_quality = None limit_size = False max_width = None max_height = None if export_format == "image": image_format = self.format_combo.currentText().lower() image_quality = self.quality_slider.value() limit_size = self.limit_size_checkbox.isChecked() max_width = self.max_width_spin.value() if limit_size else None max_height = self.max_height_spin.value() if limit_size else None else: # PDF image_quality = self.pdf_quality_slider.value() limit_size = self.pdf_limit_size_checkbox.isChecked() max_width = self.pdf_max_width_spin.value() if limit_size else None max_height = self.pdf_max_height_spin.value() if limit_size else None # 获取输出目录 output_dir = self.output_dir_edit.text() return { "selected_types": selected_types, "page_selection": page_selection, "layout_type": layout_type, "export_format": export_format, "image_format": image_format, "image_quality": image_quality, "limit_size": limit_size, "max_width": max_width, "max_height": max_height, "output_dir": output_dir } class PDFExportOptionsDialog(QDialog): """PDF导出选项对话框""" def __init__(self, parent=None): super().__init__(parent) self.setWindowTitle("导出标注PDF选项") self.setMinimumWidth(400) layout = QVBoxLayout() # 页面选择 page_group = QGroupBox("页面选择") page_layout = QVBoxLayout() self.all_pages_radio = QRadioButton("所有页面") self.all_pages_radio.setChecked(True) page_layout.addWidget(self.all_pages_radio) self.current_page_radio = QRadioButton("当前页面") page_layout.addWidget(self.current_page_radio) self.select_pages_radio = QRadioButton("选择页面") page_layout.addWidget(self.select_pages_radio) page_group.setLayout(page_layout) layout.addWidget(page_group) # 图像压缩选项 compression_group = QGroupBox("图像压缩选项") compression_layout = QFormLayout() # 图像质量选项 self.quality_slider = QSlider(Qt.Horizontal) self.quality_slider.setMinimum(1) self.quality_slider.setMaximum(100) self.quality_slider.setValue(75) self.quality_label = QLabel("75%") quality_layout = QHBoxLayout() quality_layout.addWidget(self.quality_slider) quality_layout.addWidget(self.quality_label) self.quality_slider.valueChanged.connect(lambda v: self.quality_label.setText(f"{v}%")) compression_layout.addRow("图像质量:", quality_layout) # 尺寸限制选项 self.limit_size_checkbox = QCheckBox("限制最大尺寸") self.limit_size_checkbox.setChecked(False) compression_layout.addRow(self.limit_size_checkbox) size_layout = QHBoxLayout() self.max_width_spin = QSpinBox() self.max_width_spin.setRange(100, 5000) self.max_width_spin.setValue(1920) self.max_width_spin.setEnabled(False) size_layout.addWidget(QLabel("最大宽度:")) size_layout.addWidget(self.max_width_spin) size_layout.addWidget(QLabel("像素")) self.max_height_spin = QSpinBox() self.max_height_spin.setRange(100, 5000) self.max_height_spin.setValue(1080) self.max_height_spin.setEnabled(False) size_layout.addWidget(QLabel("最大高度:")) size_layout.addWidget(self.max_height_spin) size_layout.addWidget(QLabel("像素")) compression_layout.addRow(size_layout) self.limit_size_checkbox.toggled.connect( lambda checked: ( self.max_width_spin.setEnabled(checked), self.max_height_spin.setEnabled(checked) ) ) compression_group.setLayout(compression_layout) layout.addWidget(compression_group) # 输出文件选择 file_group = QGroupBox("输出文件") file_layout = QVBoxLayout() path_layout = QHBoxLayout() self.output_file_edit = QLineEdit() path_layout.addWidget(self.output_file_edit) self.browse_file_btn = QPushButton("浏览...") self.browse_file_btn.clicked.connect(self.browse_output_file) path_layout.addWidget(self.browse_file_btn) file_layout.addLayout(path_layout) file_group.setLayout(file_layout) layout.addWidget(file_group) # 按钮 buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel) buttons.accepted.connect(self.accept) buttons.rejected.connect(self.reject) layout.addWidget(buttons) self.setLayout(layout) # 初始化变量 self.output_file = "" def browse_output_file(self): """选择输出文件""" file_path, _ = QFileDialog.getSaveFileName( self, "选择输出文件", "", "PDF文件 (*.pdf)" ) if file_path: self.output_file = file_path self.output_file_edit.setText(file_path) def get_export_options(self): """获取导出选项""" # 获取页面选择 if self.all_pages_radio.isChecked(): page_selection = "all" elif self.current_page_radio.isChecked(): page_selection = "current" else: page_selection = "select" # 实际页面选择将在主窗口中处理 # 获取图像质量 image_quality = self.quality_slider.value() # 获取尺寸限制 limit_size = self.limit_size_checkbox.isChecked() max_width = self.max_width_spin.value() if limit_size else None max_height = self.max_height_spin.value() if limit_size else None # 获取输出文件 output_file = self.output_file_edit.text() return { "page_selection": page_selection, "image_quality": image_quality, "limit_size": limit_size, "max_width": max_width, "max_height": max_height, "output_file": output_file } class ColumnCropDialog(QDialog): """分栏裁切选项对话框""" def __init__(self, block_types, parent=None): super().__init__(parent) self.setWindowTitle("分栏裁切选项") self.setMinimumWidth(500) layout = QVBoxLayout() # 区块类型选择（用于马赛克处理） mosaic_group = QGroupBox("选择需要马赛克处理的区块类型") mosaic_layout = QVBoxLayout() # 全选/取消全选 select_layout = QHBoxLayout() self.select_all_btn = QPushButton("全选") self.select_all_btn.clicked.connect(self.select_all) select_layout.addWidget(self.select_all_btn) self.deselect_all_btn = QPushButton("取消全选") self.deselect_all_btn.clicked.connect(self.deselect_all) select_layout.addWidget(self.deselect_all_btn) select_layout.addStretch() mosaic_layout.addLayout(select_layout) # 区块类型复选框 self.type_checkboxes = [] for block_type in block_types: checkbox = QCheckBox(block_type) checkbox.setChecked(False) # 默认不选 mosaic_layout.addWidget(checkbox) self.type_checkboxes.append(checkbox) mosaic_group.setLayout(mosaic_layout) layout.addWidget(mosaic_group) # 马赛克颜色选择 color_group = QGroupBox("马赛克颜色") color_layout = QVBoxLayout() self.white_radio = QRadioButton("纯白色") self.white_radio.setChecked(True) color_layout.addWidget(self.white_radio) self.gray_radio = QRadioButton("纯灰色") color_layout.addWidget(self.gray_radio) self.custom_radio = QRadioButton("自定义颜色") color_layout.addWidget(self.custom_radio) # 自定义颜色选择 custom_color_layout = QHBoxLayout() self.color_label = QLabel("颜色:") custom_color_layout.addWidget(self.color_label) self.color_button = QPushButton() self.color_button.setStyleSheet("background-color: #000000;") self.color_button.setFixedSize(30, 20) self.color_button.clicked.connect(self.choose_color) custom_color_layout.addWidget(self.color_button) custom_color_layout.addStretch() color_layout.addLayout(custom_color_layout) color_group.setLayout(color_layout) layout.addWidget(color_group) # 分栏选项 column_group = QGroupBox("分栏设置") column_layout = QVBoxLayout() self.double_column_radio = QRadioButton("双栏") self.double_column_radio.setChecked(True) column_layout.addWidget(self.double_column_radio) self.triple_column_radio = QRadioButton("三栏") column_layout.addWidget(self.triple_column_radio) column_group.setLayout(column_layout) layout.addWidget(column_group) # 导出选项 export_group = QGroupBox("导出选项") export_layout = QVBoxLayout() self.export_image_radio = QRadioButton("导出为图像") self.export_image_radio.setChecked(True) export_layout.addWidget(self.export_image_radio) self.export_pdf_radio = QRadioButton("导出为PDF") export_layout.addWidget(self.export_pdf_radio) export_group.setLayout(export_layout) layout.addWidget(export_group) # 图像格式选项（仅当导出为图像时显示） self.format_group = QGroupBox("图像格式") format_layout = QFormLayout() self.format_combo = QComboBox() self.format_combo.addItems(["JPEG", "PNG"]) format_layout.addRow("图像格式:", self.format_combo) # 图像质量选项 self.quality_slider = QSlider(Qt.Horizontal) self.quality_slider.setMinimum(1) self.quality_slider.setMaximum(100) self.quality_slider.setValue(90) self.quality_label = QLabel("90%") quality_layout = QHBoxLayout() quality_layout.addWidget(self.quality_slider) quality_layout.addWidget(self.quality_label) self.quality_slider.valueChanged.connect(lambda v: self.quality_label.setText(f"{v}%")) format_layout.addRow("图像质量:", quality_layout) self.format_group.setLayout(format_layout) layout.addWidget(self.format_group) # 连接信号，根据导出格式显示/隐藏图像格式选项 self.export_image_radio.toggled.connect(self.format_group.setVisible) # 输出目录选择 dir_group = QGroupBox("输出目录") dir_layout = QVBoxLayout() path_layout = QHBoxLayout() self.output_dir_edit = QLineEdit() path_layout.addWidget(self.output_dir_edit) self.browse_dir_btn = QPushButton("浏览...") self.browse_dir_btn.clicked.connect(self.browse_output_dir) path_layout.addWidget(self.browse_dir_btn) dir_layout.addLayout(path_layout) dir_group.setLayout(dir_layout) layout.addWidget(dir_group) # 按钮 buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel) buttons.accepted.connect(self.accept) buttons.rejected.connect(self.reject) layout.addWidget(buttons) self.setLayout(layout) # 初始化变量 self.output_dir = "" self.custom_color = QColor(0, 0, 0) # 默认黑色 def select_all(self): """全选""" for checkbox in self.type_checkboxes: checkbox.setChecked(True) def deselect_all(self): """取消全选""" for checkbox in self.type_checkboxes: checkbox.setChecked(False) def choose_color(self): """选择自定义颜色""" color = QColorDialog.getColor(self.custom_color, self) if color.isValid(): self.custom_color = color self.color_button.setStyleSheet(f"background-color: {color.name()};") def browse_output_dir(self): """选择输出目录""" dir_path = QFileDialog.getExistingDirectory(self, "选择输出目录") if dir_path: self.output_dir = dir_path self.output_dir_edit.setText(dir_path) def get_crop_options(self): """获取分栏裁切选项""" # 获取选中的区块类型（用于马赛克） mosaic_types = [] for checkbox in self.type_checkboxes: if checkbox.isChecked(): mosaic_types.append(checkbox.text()) # 获取马赛克颜色 if self.white_radio.isChecked(): mosaic_color = "white" elif self.gray_radio.isChecked(): mosaic_color = "gray" else: mosaic_color = self.custom_color # 获取分栏数 column_count = 2 if self.double_column_radio.isChecked() else 3 # 获取导出格式 export_format = "image" if self.export_image_radio.isChecked() else "pdf" # 获取图像格式和质量（仅当导出为图像时） image_format = None image_quality = None if export_format == "image": image_format = self.format_combo.currentText().lower() image_quality = self.quality_slider.value() # 获取输出目录 output_dir = self.output_dir_edit.text() return { "mosaic_types": mosaic_types, "mosaic_color": mosaic_color, "column_count": column_count, "export_format": export_format, "image_format": image_format, "image_quality": image_quality, "output_dir": output_dir } class ImageMaskDialog(QDialog): """图像掩码选项对话框""" def __init__(self, block_types, parent=None): super().__init__(parent) self.setWindowTitle("图像掩码选项") self.setMinimumWidth(500) layout = QVBoxLayout() # 区块类型选择（用于马赛克处理） mosaic_group = QGroupBox("选择需要马赛克处理的区块类型") mosaic_layout = QVBoxLayout() # 全选/取消全选 select_layout = QHBoxLayout() self.select_all_btn = QPushButton("全选") self.select_all_btn.clicked.connect(self.select_all) select_layout.addWidget(self.select_all_btn) self.deselect_all_btn = QPushButton("取消全选") self.deselect_all_btn.clicked.connect(self.deselect_all) select_layout.addWidget(self.deselect_all_btn) select_layout.addStretch() mosaic_layout.addLayout(select_layout) # 区块类型复选框 self.type_checkboxes = [] for block_type in block_types: checkbox = QCheckBox(block_type) checkbox.setChecked(False) # 默认不选 mosaic_layout.addWidget(checkbox) self.type_checkboxes.append(checkbox) mosaic_group.setLayout(mosaic_layout) layout.addWidget(mosaic_group) # 马赛克颜色选择 color_group = QGroupBox("马赛克颜色") color_layout = QVBoxLayout() self.white_radio = QRadioButton("纯白色") self.white_radio.setChecked(True) color_layout.addWidget(self.white_radio) self.gray_radio = QRadioButton("纯灰色") color_layout.addWidget(self.gray_radio) self.custom_radio = QRadioButton("自定义颜色") color_layout.addWidget(self.custom_radio) # 自定义颜色选择 custom_color_layout = QHBoxLayout() self.color_label = QLabel("颜色:") custom_color_layout.addWidget(self.color_label) self.color_button = QPushButton() self.color_button.setStyleSheet("background-color: #000000;") self.color_button.setFixedSize(30, 20) self.color_button.clicked.connect(self.choose_color) custom_color_layout.addWidget(self.color_button) custom_color_layout.addStretch() color_layout.addLayout(custom_color_layout) color_group.setLayout(color_layout) layout.addWidget(color_group) # 页面选择 page_group = QGroupBox("页面选择") page_layout = QVBoxLayout() self.all_pages_radio = QRadioButton("所有页面") self.all_pages_radio.setChecked(True) page_layout.addWidget(self.all_pages_radio) self.current_page_radio = QRadioButton("当前页面") page_layout.addWidget(self.current_page_radio) self.select_pages_radio = QRadioButton("选择页面") page_layout.addWidget(self.select_pages_radio) page_group.setLayout(page_layout) layout.addWidget(page_group) # 导出选项 export_group = QGroupBox("导出选项") export_layout = QVBoxLayout() self.export_image_radio = QRadioButton("导出为图像") self.export_image_radio.setChecked(True) export_layout.addWidget(self.export_image_radio) self.export_pdf_radio = QRadioButton("导出为PDF") export_layout.addWidget(self.export_pdf_radio) export_group.setLayout(export_layout) layout.addWidget(export_group) # 图像格式选项（仅当导出为图像时显示） self.format_group = QGroupBox("图像格式") format_layout = QFormLayout() self.format_combo = QComboBox() self.format_combo.addItems(["JPEG", "PNG"]) format_layout.addRow("图像格式:", self.format_combo) # 图像质量选项 self.quality_slider = QSlider(Qt.Horizontal) self.quality_slider.setMinimum(1) self.quality_slider.setMaximum(100) self.quality_slider.setValue(90) self.quality_label = QLabel("90%") quality_layout = QHBoxLayout() quality_layout.addWidget(self.quality_slider) quality_layout.addWidget(self.quality_label) self.quality_slider.valueChanged.connect(lambda v: self.quality_label.setText(f"{v}%")) format_layout.addRow("图像质量:", quality_layout) # 尺寸限制选项 self.limit_size_checkbox = QCheckBox("限制最大尺寸") self.limit_size_checkbox.setChecked(False) format_layout.addRow(self.limit_size_checkbox) size_layout = QHBoxLayout() self.max_width_spin = QSpinBox() self.max_width_spin.setRange(100, 5000) self.max_width_spin.setValue(1920) self.max_width_spin.setEnabled(False) size_layout.addWidget(QLabel("最大宽度:")) size_layout.addWidget(self.max_width_spin) size_layout.addWidget(QLabel("像素")) self.max_height_spin = QSpinBox() self.max_height_spin.setRange(100, 5000) self.max_height_spin.setValue(1080) self.max_height_spin.setEnabled(False) size_layout.addWidget(QLabel("最大高度:")) size_layout.addWidget(self.max_height_spin) size_layout.addWidget(QLabel("像素")) format_layout.addRow(size_layout) self.limit_size_checkbox.toggled.connect( lambda checked: ( self.max_width_spin.setEnabled(checked), self.max_height_spin.setEnabled(checked) ) ) self.format_group.setLayout(format_layout) layout.addWidget(self.format_group) # PDF压缩选项（仅当导出为PDF时显示） self.pdf_format_group = QGroupBox("PDF压缩选项") pdf_format_layout = QFormLayout() # PDF图像质量选项 self.pdf_quality_slider = QSlider(Qt.Horizontal) self.pdf_quality_slider.setMinimum(1) self.pdf_quality_slider.setMaximum(100) self.pdf_quality_slider.setValue(75) self.pdf_quality_label = QLabel("75%") pdf_quality_layout = QHBoxLayout() pdf_quality_layout.addWidget(self.pdf_quality_slider) pdf_quality_layout.addWidget(self.pdf_quality_label) self.pdf_quality_slider.valueChanged.connect(lambda v: self.pdf_quality_label.setText(f"{v}%")) pdf_format_layout.addRow("图像质量:", pdf_quality_layout) # PDF尺寸限制选项 self.pdf_limit_size_checkbox = QCheckBox("限制最大尺寸") self.pdf_limit_size_checkbox.setChecked(False) pdf_format_layout.addRow(self.pdf_limit_size_checkbox) pdf_size_layout = QHBoxLayout() self.pdf_max_width_spin = QSpinBox() self.pdf_max_width_spin.setRange(100, 5000) self.pdf_max_width_spin.setValue(1920) self.pdf_max_width_spin.setEnabled(False) pdf_size_layout.addWidget(QLabel("最大宽度:")) pdf_size_layout.addWidget(self.pdf_max_width_spin) pdf_size_layout.addWidget(QLabel("像素")) self.pdf_max_height_spin = QSpinBox() self.pdf_max_height_spin.setRange(100, 5000) self.pdf_max_height_spin.setValue(1080) self.pdf_max_height_spin.setEnabled(False) pdf_size_layout.addWidget(QLabel("最大高度:")) pdf_size_layout.addWidget(self.pdf_max_height_spin) pdf_size_layout.addWidget(QLabel("像素")) pdf_format_layout.addRow(pdf_size_layout) self.pdf_limit_size_checkbox.toggled.connect( lambda checked: ( self.pdf_max_width_spin.setEnabled(checked), self.pdf_max_height_spin.setEnabled(checked) ) ) self.pdf_format_group.setLayout(pdf_format_layout) layout.addWidget(self.pdf_format_group) # 连接信号，根据导出格式显示/隐藏相应选项 self.export_image_radio.toggled.connect(self.format_group.setVisible) self.export_pdf_radio.toggled.connect(self.pdf_format_group.setVisible) self.export_image_radio.toggled.connect(lambda checked: self.pdf_format_group.setHidden(checked)) self.export_pdf_radio.toggled.connect(lambda checked: self.format_group.setHidden(checked)) # 输出目录选择 dir_group = QGroupBox("输出目录") dir_layout = QVBoxLayout() path_layout = QHBoxLayout() self.output_dir_edit = QLineEdit() path_layout.addWidget(self.output_dir_edit) self.browse_dir_btn = QPushButton("浏览...") self.browse_dir_btn.clicked.connect(self.browse_output_dir) path_layout.addWidget(self.browse_dir_btn) dir_layout.addLayout(path_layout) dir_group.setLayout(dir_layout) layout.addWidget(dir_group) # 按钮 buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel) buttons.accepted.connect(self.accept) buttons.rejected.connect(self.reject) layout.addWidget(buttons) self.setLayout(layout) # 初始化变量 self.output_dir = "" self.custom_color = QColor(0, 0, 0) # 默认黑色 # 初始状态设置 self.format_group.setVisible(True) self.pdf_format_group.setVisible(False) def select_all(self): """全选""" for checkbox in self.type_checkboxes: checkbox.setChecked(True) def deselect_all(self): """取消全选""" for checkbox in self.type_checkboxes: checkbox.setChecked(False) def choose_color(self): """选择自定义颜色""" color = QColorDialog.getColor(self.custom_color, self) if color.isValid(): self.custom_color = color self.color_button.setStyleSheet(f"background-color: {color.name()};") def browse_output_dir(self): """选择输出目录""" dir_path = QFileDialog.getExistingDirectory(self, "选择输出目录") if dir_path: self.output_dir = dir_path self.output_dir_edit.setText(dir_path) def get_mask_options(self): """获取图像掩码选项""" # 获取选中的区块类型（用于马赛克） mosaic_types = [] for checkbox in self.type_checkboxes: if checkbox.isChecked(): mosaic_types.append(checkbox.text()) # 获取马赛克颜色 if self.white_radio.isChecked(): mosaic_color = "white" elif self.gray_radio.isChecked(): mosaic_color = "gray" else: mosaic_color = self.custom_color # 获取页面选择 if self.all_pages_radio.isChecked(): page_selection = "all" elif self.current_page_radio.isChecked(): page_selection = "current" else: page_selection = "select" # 实际页面选择将在主窗口中处理 # 获取导出格式 export_format = "image" if self.export_image_radio.isChecked() else "pdf" # 获取图像格式和质量（仅当导出为图像时） image_format = None image_quality = None limit_size = False max_width = None max_height = None if export_format == "image": image_format = self.format_combo.currentText().lower() image_quality = self.quality_slider.value() limit_size = self.limit_size_checkbox.isChecked() max_width = self.max_width_spin.value() if limit_size else None max_height = self.max_height_spin.value() if limit_size else None else: # PDF image_quality = self.pdf_quality_slider.value() limit_size = self.pdf_limit_size_checkbox.isChecked() max_width = self.pdf_max_width_spin.value() if limit_size else None max_height = self.pdf_max_height_spin.value() if limit_size else None # 获取输出目录 output_dir = self.output_dir_edit.text() return { "mosaic_types": mosaic_types, "mosaic_color": mosaic_color, "page_selection": page_selection, "export_format": export_format, "image_format": image_format, "image_quality": image_quality, "limit_size": limit_size, "max_width": max_width, "max_height": max_height, "output_dir": output_dir } class LayoutAnalyzerThread(QThread): """后台线程处理版面分析任务""" progress_updated = pyqtSignal(int) analysis_complete = pyqtSignal(dict, list) error_occurred = pyqtSignal(str) def __init__(self, pdf_path, model_name, model_dir=None, img_size=None, threshold=None, layout_unclip_ratio=None, layout_merge_bboxes_mode=None): super().__init__() self.pdf_path = pdf_path self.model_name = model_name self.model_dir = model_dir self.img_size = img_size # 保留参数以备将来使用 self.threshold = threshold self.layout_unclip_ratio = layout_unclip_ratio self.layout_merge_bboxes_mode = layout_merge_bboxes_mode self._is_running = True def convert_pdf_to_images(self, pdf_path): """使用PyMuPDF将PDF转换为图像列表""" try: doc = fitz.open(pdf_path) images = [] total_pages = len(doc) for page_num in range(total_pages): if not self._is_running: break # 加载页面 page = doc.load_page(page_num) # 获取页面图像，提高分辨率以获得更好的OCR效果 zoom = 2 # 缩放因子 mat = fitz.Matrix(zoom, zoom) pix = page.get_pixmap(matrix=mat) # 转换为numpy数组 img = np.frombuffer(pix.samples, dtype=np.uint8).reshape((pix.height, pix.width, pix.n)) # 如果是RGBA格式，转换为RGB if pix.n == 4: img = img[:, :, :3] # 如果是灰度图，转换为RGB elif pix.n == 1: img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB) images.append(img) doc.close() return images except Exception as e: raise Exception(f"PDF转换失败: {str(e)}") def run(self): try: self.progress_updated.emit(10) try: if self.model_dir and os.path.exists(os.path.join(self.model_dir, self.model_name)): model_path = os.path.join(self.model_dir, self.model_name) model = LayoutDetection( model_name=self.model_name, model_dir=model_path, threshold=self.threshold, layout_unclip_ratio=self.layout_unclip_ratio, layout_merge_bboxes_mode=self.layout_merge_bboxes_mode ) print(f"使用本地模型: {model_path}") else: model = LayoutDetection( model_name=self.model_name, threshold=self.threshold, layout_unclip_ratio=self.layout_unclip_ratio, layout_merge_bboxes_mode=self.layout_merge_bboxes_mode ) print(f"使用在线模型: {self.model_name}") except Exception as e: print(f"本地模型加载失败，尝试使用在线模型: {str(e)}") model = LayoutDetection( model_name=self.model_name, threshold=self.threshold, layout_unclip_ratio=self.layout_unclip_ratio, layout_merge_bboxes_mode=self.layout_merge_bboxes_mode ) # 将PDF转换为图像 self.progress_updated.emit(30) images = self.convert_pdf_to_images(self.pdf_path) total_pages = len(images) if total_pages == 0: raise Exception("PDF文件为空或无法读取") results = [] all_detections = [] for i, image in enumerate(images): if not self._is_running: break # 更新进度 progress = 30 + int(60 * (i + 1) / total_pages) self.progress_updated.emit(progress) # 确保图像是numpy数组格式 if not isinstance(image, np.ndarray): img_np = np.array(image) else: img_np = image # 执行版面分析 output = model.predict(img_np, batch_size=1, layout_nms=True) # 处理结果 page_result = { 'page_index': i, 'image': img_np, 'detections': [] } # 正确处理PaddleOCR的返回结果 if output and len(output) > 0: # 获取第一个结果 result = output[0] # 检查结果类型并提取数据 if hasattr(result, 'data'): # 如果是Result对象，获取data属性 data = result.data if isinstance(data, dict) and 'boxes' in data: boxes = data['boxes'] else: boxes = [] elif isinstance(result, dict): # 如果是字典，直接获取boxes boxes = result.get('boxes', []) else: # 其他情况，尝试转换为字典 try: result_dict = dict(result) boxes = result_dict.get('boxes', []) except: boxes = [] # 处理每个检测框 for box in boxes: if isinstance(box, dict): detection = { 'label': box.get('label', 'unknown'), 'confidence': box.get('score', 0.0), 'coordinates': box.get('coordinate', [0, 0, 0, 0]), 'cls_id': box.get('cls_id', -1) } page_result['detections'].append(detection) all_detections.append(detection) results.append(page_result) self.progress_updated.emit(100) self.analysis_complete.emit({'pages': results}, all_detections) except Exception as e: self.error_occurred.emit(f"分析过程中出错: {str(e)}") def stop(self): self._is_running = False class ImageExportWorker: """图像导出工作器 - 使用多线程处理""" def __init__(self, progress_callback): self.progress_callback = progress_callback self.stop_flag = False def process_page(self, page_data, page_idx, export_options): """处理单个页面""" if self.stop_flag: return None try: # 裁切区块图像 self.crop_blocks_from_page(page_data, page_idx, export_options) return (page_idx, True) # 返回成功状态 except Exception as e: print(f"处理第 {page_idx + 1} 页时出错: {str(e)}") return (page_idx, False) def crop_blocks_from_page(self, page_data, page_idx, export_options): """从页面裁切区块图像""" # 获取导出选项 selected_types = export_options.get("selected_types", []) layout_type = export_options.get("layout_type", "single") export_format = export_options.get("export_format", "image") image_format = export_options.get("image_format", "jpeg") image_quality = export_options.get("image_quality", 90) limit_size = export_options.get("limit_size", False) max_width = export_options.get("max_width", None) max_height = export_options.get("max_height", None) output_dir = export_options.get("output_dir", "") if not output_dir or not os.path.exists(output_dir): return # 创建类别目录 type_dirs = {} for detection in page_data['detections']: block_type = detection['label'] if block_type in selected_types and block_type not in type_dirs: type_dir = os.path.join(output_dir, block_type) os.makedirs(type_dir, exist_ok=True) type_dirs[block_type] = type_dir # 获取页面图像 img = page_data['image'] # 确保图像是RGB格式 if len(img.shape) == 2: # 灰度图 img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB) elif img.shape[2] == 4: # RGBA img = cv2.cvtColor(img, cv2.COLOR_RGBA2RGB) elif img.shape[2] == 3: # BGR (OpenCV默认) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 按区块类型分组 blocks_by_type = {} for detection in page_data['detections']: block_type = detection['label'] if block_type in selected_types: if block_type not in blocks_by_type: blocks_by_type[block_type] = [] blocks_by_type[block_type].append(detection) # 对每种类型的区块进行处理 for block_type, blocks in blocks_by_type.items(): # 按位置排序区块：先按Y坐标（从上到下），再按X坐标（从左到右） blocks.sort(key=lambda b: (b['coordinates'][1], b['coordinates'][0])) # 确定栏位 img_width = img.shape[1] columns = [] if layout_type == "single": # 单栏，只有一个栏位 columns.append(("single", 0, img_width)) elif layout_type == "double": # 双栏，平均分为左右两栏 mid_x = img_width // 2 columns.append(("left", 0, mid_x)) columns.append(("right", mid_x, img_width)) elif layout_type == "triple": # 三栏，平均分为三栏 one_third = img_width // 3 two_thirds = 2 * img_width // 3 columns.append(("left", 0, one_third)) columns.append(("middle", one_third, two_thirds)) columns.append(("right", two_thirds, img_width)) # 为每栏创建区块计数器 column_counters = {} for name, _, _ in columns: column_counters[name] = 1 # 为每个区块裁切图像并保存 for block in blocks: coords = block['coordinates'] x1, y1, x2, y2 = map(int, coords) # 裁切区块 block_img = img[y1:y2, x1:x2] # 确定栏位 block_center_x = (x1 + x2) // 2 column_name = "single" for name, col_start, col_end in columns: if col_start <= block_center_x < col_end: column_name = name break # 应用尺寸限制（如果需要） if limit_size and (max_width or max_height): h, w = block_img.shape[:2] scale = 1.0 if max_width and w > max_width: scale = max_width / w if max_height and h > max_height: scale = min(scale, max_height / h) if scale != 1.0: new_width = int(w * scale) new_height = int(h * scale) block_img = cv2.resize(block_img, (new_width, new_height), interpolation=cv2.INTER_AREA) # 获取当前栏的序号 sequence_num = column_counters[column_name] column_counters[column_name] += 1 # 生成文件名：页码_栏位_序号 if layout_type == "single": filename = f"page_{page_idx + 1:03d}_single_{sequence_num:02d}.{image_format}" else: filename = f"page_{page_idx + 1:03d}_{column_name}_{sequence_num:02d}.{image_format}" filepath = os.path.join(type_dirs[block_type], filename) # 根据导出格式保存 if export_format == "image": # 根据格式保存图像 if image_format.lower() == "jpeg" or image_format.lower() == "jpg": cv2.imwrite(filepath, cv2.cvtColor(block_img, cv2.COLOR_RGB2BGR), [int(cv2.IMWRITE_JPEG_QUALITY), image_quality]) elif image_format.lower() == "png": cv2.imwrite(filepath, cv2.cvtColor(block_img, cv2.COLOR_RGB2BGR), [int(cv2.IMWRITE_PNG_COMPRESSION), 10 - (image_quality // 10)]) else: # PDF # 保存为PDF pdf_path = os.path.join(type_dirs[block_type], f"page_{page_idx + 1:03d}_{column_name}_{sequence_num:02d}.pdf") # 创建PDF文档 doc = fitz.open() page = doc.new_page(width=block_img.shape[1], height=block_img.shape[0]) # 将图像转换为字节流 img_bytes = cv2.imencode('.jpg', block_img, [int(cv2.IMWRITE_JPEG_QUALITY), image_quality])[ 1].tobytes() # 插入图像 rect = fitz.Rect(0, 0, block_img.shape[1], block_img.shape[0]) page.insert_image(rect, stream=img_bytes) # 保存PDF doc.save(pdf_path) doc.close() class ImageExportThread(QThread): """图像导出线程""" progress_updated = pyqtSignal(int, str) export_complete = pyqtSignal(str) export_failed = pyqtSignal(str) def __init__(self, analysis_results, export_options): super().__init__() self.analysis_results = analysis_results self.export_options = export_options self._is_running = True self.worker = None def run(self): try: # 确定要处理的页面 if self.export_options["page_selection"] == "all": pages_to_process = range(len(self.analysis_results)) elif self.export_options["page_selection"] == "current": pages_to_process = [self.export_options.get("current_page", 0)] else: pages_to_process = self.export_options.get("selected_pages", []) total_pages = len(pages_to_process) if total_pages == 0: self.export_failed.emit("没有选择要处理的页面") return # 创建工作器 self.worker = ImageExportWorker(self.progress_updated.emit) # 使用线程池处理页面 with ThreadPoolExecutor(max_workers=4) as executor: # 提交所有页面处理任务 future_to_page = { executor.submit( self.worker.process_page, self.analysis_results[page_idx], page_idx, self.export_options ): page_idx for page_idx in pages_to_process } # 收集处理结果 completed = 0 for future in future_to_page: if not self._is_running: break page_idx = future_to_page[future] try: result = future.result(timeout=30) # 设置超时 completed += 1 progress = int(completed / total_pages * 100) self.progress_updated.emit(progress, f"已处理 {completed}/{total_pages} 页...") except Exception as e: print(f"处理第 {page_idx + 1} 页时出错: {str(e)}") if not self._is_running: return # 导出完成 output_dir = self.export_options.get("output_dir", "") self.export_complete.emit(f"区块图像已成功导出到:\n{output_dir}") except Exception as e: self.export_failed.emit(str(e)) def stop(self): self._is_running = False if self.worker: self.worker.stop_flag = True class PDFExportWorker: """PDF导出工作器 - 使用多线程处理""" def __init__(self, progress_callback): self.progress_callback = progress_callback self.stop_flag = False def process_page(self, page_data, page_idx, export_options): """处理单个页面""" if self.stop_flag: return None try: # 获取图像质量选项 image_quality = export_options.get("image_quality", 75) limit_size = export_options.get("limit_size", False) max_width = export_options.get("max_width", None) max_height = export_options.get("max_height", None) # 创建标注图像 img = self.create_annotated_image(page_data['image'], page_data['detections']) # 应用尺寸限制（如果需要） if limit_size and (max_width or max_height): h, w = img.shape[:2] scale = 1.0 if max_width and w > max_width: scale = max_width / w if max_height and h > max_height: scale = min(scale, max_height / h) if scale != 1.0: new_width = int(w * scale) new_height = int(h * scale) img = cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_AREA) return (page_idx, img, image_quality) except Exception as e: print(f"处理第 {page_idx + 1} 页时出错: {str(e)}") return None def create_annotated_image(self, img_array, detections): # 创建图像副本 annotated_img = img_array.copy() # 确保图像是BGR格式（OpenCV默认） if len(annotated_img.shape) == 2: # 灰度图 annotated_img = cv2.cvtColor(annotated_img, cv2.COLOR_GRAY2BGR) elif annotated_img.shape[2] == 4: # RGBA annotated_img = cv2.cvtColor(annotated_img, cv2.COLOR_RGBA2BGR) elif annotated_img.shape[2] == 3 and not np.array_equal(annotated_img[0, 0], img_array[0, 0]): # 可能是RGB annotated_img = cv2.cvtColor(annotated_img, cv2.COLOR_RGB2BGR) # 为不同区块类型定义颜色 colors = { # 基础文本元素 'doc_title': (0, 0, 255), # 红色 - 文档标题 'paragraph_title': (255, 0, 128), # 玫瑰红 - 段落标题 'text': (0, 255, 0), # 绿色 - 文本 'number': (128, 128, 128), # 灰色 - 页码 # 文档结构元素 'abstract': (128, 128, 0), # 橄榄色 - 摘要 'content': (64, 128, 0), # 深橄榄绿 - 目录 'reference': (42, 42, 165), # 深蓝 - 参考文献 'footnote': (192, 192, 192), # 银色 - 脚注 'header': (255, 192, 203), # 粉色 - 页眉 'footer': (173, 216, 230), # 浅蓝 - 页脚 # 特殊内容元素 'algorithm': (75, 0, 130), # 靛蓝 - 算法 'formula': (255, 255, 0), # 黄色 - 公式 'formula_number': (128, 0, 255), # 紫罗兰 - 公式编号 # 图像和表格元素 'image': (255, 128, 0), # 深橙 - 图像 'figure_title': (255, 0, 64), # 深玫瑰红 - 图像标题 'chart': (0, 128, 255), # 天蓝 - 图表 'chart_title': (0, 255, 128), # 春绿色 - 图表标题 'table': (0, 0, 192), # 深红色 - 表格 'table_title': (128, 255, 255), # 浅青色 - 表格标题 # 其他元素 'seal': (128, 0, 128), # 深紫色 - 印章 'header_image': (192, 255, 62), # 酸橙绿 - 页眉图像 'footer_image': (62, 255, 192), # 蓝绿 - 页脚图像 'aside_text': (128, 64, 0), # 棕色 - 侧栏文本 } # 绘制边界框 for detection in detections: label = detection['label'] coords = detection['coordinates'] x1, y1, x2, y2 = map(int, coords) # 获取颜色，如果不在字典中则使用白色 color = colors.get(label, (255, 255, 255)) # 绘制矩形 cv2.rectangle(annotated_img, (x1, y1), (x2, y2), color, 2) # 添加标签 label_text = f"{label} ({detection['confidence']:.2f})" # 确保文本不会超出图像边界 text_y = max(y1 - 10, 20) cv2.putText( annotated_img, label_text, (x1, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1 ) return annotated_img class PDFExportThread(QThread): """PDF导出线程""" progress_updated = pyqtSignal(int, str) export_complete = pyqtSignal(str) export_failed = pyqtSignal(str) def __init__(self, analysis_results, export_options, file_path=None): super().__init__() self.analysis_results = analysis_results self.export_options = export_options self.file_path = file_path # PDF输出文件路径 self._is_running = True self.worker = None def run(self): try: # 确定要处理的页面 if self.export_options["page_selection"] == "all": pages_to_process = range(len(self.analysis_results)) elif self.export_options["page_selection"] == "current": pages_to_process = [self.export_options.get("current_page", 0)] else: pages_to_process = self.export_options.get("selected_pages", []) total_pages = len(pages_to_process) if total_pages == 0: self.export_failed.emit("没有选择要处理的页面") return # 创建工作器 self.worker = PDFExportWorker(self.progress_updated.emit) # 使用线程池处理页面 with ThreadPoolExecutor(max_workers=4) as executor: # 提交所有页面处理任务 future_to_page = { executor.submit( self.worker.process_page, self.analysis_results[page_idx], page_idx, self.export_options ): page_idx for page_idx in pages_to_process } # 收集处理结果 processed_pages = {} completed = 0 for future in future_to_page: if not self._is_running: break page_idx = future_to_page[future] try: result = future.result(timeout=30) # 设置超时 if result is not None: processed_pages[result[0]] = (result[1], result[2]) # (图像, 质量) completed += 1 progress = int(completed / total_pages * 100) self.progress_updated.emit(progress, f"已处理 {completed}/{total_pages} 页...") except Exception as e: print(f"处理第 {page_idx + 1} 页时出错: {str(e)}") if not self._is_running: return if not processed_pages: self.export_failed.emit("没有可导出的内容") return # 按页面顺序排序 sorted_pages = sorted(processed_pages.items()) # 创建最终PDF文档 self.progress_updated.emit(100, "正在生成PDF文件...") final_doc = fitz.open() # 分批处理页面，避免内存问题 batch_size = 10 for i in range(0, len(sorted_pages), batch_size): if not self._is_running: break batch = sorted_pages[i:i + batch_size] for page_idx, (img, quality) in batch: if not self._is_running: break # 创建新页面 page = final_doc.new_page(width=img.shape[1], height=img.shape[0]) # 将图像转换为字节流 img_bytes = cv2.imencode('.jpg', img, [int(cv2.IMWRITE_JPEG_QUALITY), quality])[1].tobytes() # 插入图像 rect = fitz.Rect(0, 0, img.shape[1], img.shape[0]) page.insert_image(rect, stream=img_bytes) # 释放图像内存 del img img = None # 定期保存和清理内存 if i % 20 == 0: gc.collect() # 保存最终PDF final_doc.save(self.file_path, deflate=True, garbage=1) final_doc.close() # 检查文件大小 file_size = os.path.getsize(self.file_path) / (1024 * 1024) # MB self.export_complete.emit(f"{self.file_path}\n\n文件大小: {file_size:.2f} MB") except Exception as e: self.export_failed.emit(str(e)) def stop(self): self._is_running = False if self.worker: self.worker.stop_flag = True class ColumnCropWorker: """分栏裁切工作器""" def __init__(self, progress_callback): self.progress_callback = progress_callback self.stop_flag = False def process_page(self, page_data, page_idx, crop_options): """处理单个页面""" if self.stop_flag: return None try: # 获取页面图像 img = page_data['image'].copy() # 确保图像是RGB格式 if len(img.shape) == 2: # 灰度图 img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB) elif img.shape[2] == 4: # RGBA img = cv2.cvtColor(img, cv2.COLOR_RGBA2RGB) elif img.shape[2] == 3 and not np.array_equal(img[0, 0], page_data['image'][0, 0]): # 可能是BGR img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 对选中的区块进行马赛克处理 mosaic_types = crop_options.get("mosaic_types", []) mosaic_color = crop_options.get("mosaic_color", "white") if mosaic_types: for detection in page_data['detections']: if detection['label'] in mosaic_types: coords = detection['coordinates'] x1, y1, x2, y2 = map(int, coords) # 根据选择的颜色进行马赛克 if mosaic_color == "white": color = (255, 255, 255) elif mosaic_color == "gray": color = (128, 128, 128) else: # 自定义颜色 color = (mosaic_color.red(), mosaic_color.green(), mosaic_color.blue()) # 填充矩形 cv2.rectangle(img, (x1, y1), (x2, y2), color, -1) # 计算分栏位置 img_height, img_width = img.shape[:2] column_count = crop_options.get("column_count", 2) column_width = img_width // column_count # 裁切每栏 column_images = [] for col in range(column_count): x_start = col * column_width x_end = (col + 1) * column_width if col < column_count - 1 else img_width # 裁切当前栏 column_img = img[:, x_start:x_end] column_images.append(column_img) return (page_idx, column_images) except Exception as e: print(f"处理第 {page_idx + 1} 页时出错: {str(e)}") return None class ColumnCropThread(QThread): """分栏裁切线程""" progress_updated = pyqtSignal(int, str) crop_complete = pyqtSignal(str) crop_failed = pyqtSignal(str) def __init__(self, analysis_results, crop_options): super().__init__() self.analysis_results = analysis_results self.crop_options = crop_options self._is_running = True self.worker = None def run(self): try: total_pages = len(self.analysis_results) output_dir = self.crop_options["output_dir"] column_count = self.crop_options["column_count"] export_format = self.crop_options["export_format"] # 创建工作器 self.worker = ColumnCropWorker(self.progress_updated.emit) # 使用线程池处理页面 with ThreadPoolExecutor(max_workers=4) as executor: # 提交所有页面处理任务 future_to_page = { executor.submit( self.worker.process_page, self.analysis_results[page_idx], page_idx, self.crop_options ): page_idx for page_idx in range(total_pages) } # 收集处理结果 processed_pages = {} completed = 0 for future in future_to_page: if not self._is_running: break page_idx = future_to_page[future] try: result = future.result(timeout=30) # 设置超时 if result is not None: processed_pages[result[0]] = result[1] # 页面索引和分栏图像 completed += 1 progress = int(completed / total_pages * 100) self.progress_updated.emit(progress, f"已处理 {completed}/{total_pages} 页...") except Exception as e: print(f"处理第 {page_idx + 1} 页时出错: {str(e)}") if not self._is_running: return if not processed_pages: self.crop_failed.emit("没有可裁切的内容") return # 按页面顺序排序 sorted_pages = sorted(processed_pages.items()) # 导出结果 self.progress_updated.emit(100, "正在导出结果...") # 创建输出目录 os.makedirs(output_dir, exist_ok=True) # 为每页的每栏导出 for page_idx, column_images in sorted_pages: for col, column_img in enumerate(column_images): # 生成文件名：页码_栏位 filename = f"page_{page_idx + 1:03d}_col{col + 1}" if export_format == "image": # 保存为图像 image_format = self.crop_options.get("image_format", "jpeg") image_quality = self.crop_options.get("image_quality", 90) if image_format.lower() == "jpeg" or image_format.lower() == "jpg": output_path = os.path.join(output_dir, f"{filename}.jpg") cv2.imwrite(output_path, cv2.cvtColor(column_img, cv2.COLOR_RGB2BGR), [int(cv2.IMWRITE_JPEG_QUALITY), image_quality]) elif image_format.lower() == "png": output_path = os.path.join(output_dir, f"{filename}.png") cv2.imwrite(output_path, cv2.cvtColor(column_img, cv2.COLOR_RGB2BGR), [int(cv2.IMWRITE_PNG_COMPRESSION), 10 - (image_quality // 10)]) else: # PDF # 保存为PDF output_path = os.path.join(output_dir, f"{filename}.pdf") # 创建PDF文档 doc = fitz.open() page = doc.new_page(width=column_img.shape[1], height=column_img.shape[0]) # 将图像转换为字节流 img_bytes = cv2.imencode('.jpg', column_img)[1].tobytes() # 插入图像 rect = fitz.Rect(0, 0, column_img.shape[1], column_img.shape[0]) page.insert_image(rect, stream=img_bytes) # 保存PDF doc.save(output_path) doc.close() if self._is_running: self.crop_complete.emit(output_dir) except Exception as e: self.crop_failed.emit(str(e)) def stop(self): self._is_running = False if self.worker: self.worker.stop_flag = True class ImageMaskWorker: """图像掩码工作器""" def __init__(self, progress_callback): self.progress_callback = progress_callback self.stop_flag = False def process_page(self, page_data, page_idx, mask_options): """处理单个页面""" if self.stop_flag: return None try: # 获取页面图像 img = page_data['image'].copy() # 确保图像是RGB格式 if len(img.shape) == 2: # 灰度图 img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB) elif img.shape[2] == 4: # RGBA img = cv2.cvtColor(img, cv2.COLOR_RGBA2RGB) elif img.shape[2] == 3 and not np.array_equal(img[0, 0], page_data['image'][0, 0]): # 可能是BGR img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 对选中的区块进行马赛克处理 mosaic_types = mask_options.get("mosaic_types", []) mosaic_color = mask_options.get("mosaic_color", "white") if mosaic_types: for detection in page_data['detections']: if detection['label'] in mosaic_types: coords = detection['coordinates'] x1, y1, x2, y2 = map(int, coords) # 根据选择的颜色进行马赛克 if mosaic_color == "white": color = (255, 255, 255) elif mosaic_color == "gray": color = (128, 128, 128) else: # 自定义颜色 color = (mosaic_color.red(), mosaic_color.green(), mosaic_color.blue()) # 填充矩形 cv2.rectangle(img, (x1, y1), (x2, y2), color, -1) return (page_idx, img) except Exception as e: print(f"处理第 {page_idx + 1} 页时出错: {str(e)}") return None class ImageMaskThread(QThread): """图像掩码线程""" progress_updated = pyqtSignal(int, str) mask_complete = pyqtSignal(str) mask_failed = pyqtSignal(str) def __init__(self, analysis_results, mask_options): super().__init__() self.analysis_results = analysis_results self.mask_options = mask_options self._is_running = True self.worker = None def run(self): try: # 确定要处理的页面 if self.mask_options["page_selection"] == "all": pages_to_process = range(len(self.analysis_results)) elif self.mask_options["page_selection"] == "current": pages_to_process = [self.mask_options.get("current_page", 0)] else: pages_to_process = self.mask_options.get("selected_pages", []) total_pages = len(pages_to_process) if total_pages == 0: self.mask_failed.emit("没有选择要处理的页面") return # 创建工作器 self.worker = ImageMaskWorker(self.progress_updated.emit) # 使用线程池处理页面 with ThreadPoolExecutor(max_workers=4) as executor: # 提交所有页面处理任务 future_to_page = { executor.submit( self.worker.process_page, self.analysis_results[page_idx], page_idx, self.mask_options ): page_idx for page_idx in pages_to_process } # 收集处理结果 processed_pages = {} completed = 0 for future in future_to_page: if not self._is_running: break page_idx = future_to_page[future] try: result = future.result(timeout=30) # 设置超时 if result is not None: processed_pages[result[0]] = result[1] # 页面索引和图像 completed += 1 progress = int(completed / total_pages * 100) self.progress_updated.emit(progress, f"已处理 {completed}/{total_pages} 页...") except Exception as e: print(f"处理第 {page_idx + 1} 页时出错: {str(e)}") if not self._is_running: return if not processed_pages: self.mask_failed.emit("没有可处理的内容") return # 按页面顺序排序 sorted_pages = sorted(processed_pages.items()) # 导出结果 self.progress_updated.emit(100, "正在导出结果...") # 获取导出选项 export_format = self.mask_options.get("export_format", "image") image_format = self.mask_options.get("image_format", "jpeg") image_quality = self.mask_options.get("image_quality", 90) limit_size = self.mask_options.get("limit_size", False) max_width = self.mask_options.get("max_width", None) max_height = self.mask_options.get("max_height", None) output_dir = self.mask_options.get("output_dir", "") # 创建输出目录 os.makedirs(output_dir, exist_ok=True) # 导出每页 for page_idx, img in sorted_pages: # 应用尺寸限制（如果需要） if limit_size and (max_width or max_height): h, w = img.shape[:2] scale = 1.0 if max_width and w > max_width: scale = max_width / w if max_height and h > max_height: scale = min(scale, max_height / h) if scale != 1.0: new_width = int(w * scale) new_height = int(h * scale) img = cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_AREA) # 生成文件名 filename = f"masked_page_{page_idx + 1:03d}" if export_format == "image": # 保存为图像 if image_format.lower() == "jpeg" or image_format.lower() == "jpg": output_path = os.path.join(output_dir, f"{filename}.jpg") cv2.imwrite(output_path, cv2.cvtColor(img, cv2.COLOR_RGB2BGR), [int(cv2.IMWRITE_JPEG_QUALITY), image_quality]) elif image_format.lower() == "png": output_path = os.path.join(output_dir, f"{filename}.png") cv2.imwrite(output_path, cv2.cvtColor(img, cv2.COLOR_RGB2BGR), [int(cv2.IMWRITE_PNG_COMPRESSION), 10 - (image_quality // 10)]) else: # PDF # 保存为PDF output_path = os.path.join(output_dir, f"{filename}.pdf") # 创建PDF文档 doc = fitz.open() page = doc.new_page(width=img.shape[1], height=img.shape[0]) # 将图像转换为字节流 img_bytes = cv2.imencode('.jpg', img, [int(cv2.IMWRITE_JPEG_QUALITY), image_quality])[1].tobytes() # 插入图像 rect = fitz.Rect(0, 0, img.shape[1], img.shape[0]) page.insert_image(rect, stream=img_bytes) # 保存PDF doc.save(output_path) doc.close() if self._is_running: self.mask_complete.emit(output_dir) except Exception as e: self.mask_failed.emit(str(e)) def stop(self): self._is_running = False if self.worker: self.worker.stop_flag = True class PDFLayoutAnalyzer(QMainWindow): def __init__(self): super().__init__() self.setWindowTitle("PDF版面分析工作台 - PP-StructureV3") self.setGeometry(100, 100, 1400, 900) # 初始化变量 self.current_pdf_path = "" self.analysis_results = [] self.all_detections = [] self.current_page = 0 self.model = None self.worker_thread = None self.image_export_thread = None self.pdf_export_thread = None self.column_crop_thread = None self.image_mask_thread = None self.model_dir = "C:\\Users\\egico\\.paddlex\\official_models" # 本地模型目录 # 设置样式 self.setStyleSheet(""" QMainWindow { background-color: #f5f5f7; } QLabel { color: #333; font-size: 14px; } QPushButton { background-color: #007aff; color: white; border: none; padding: 8px 16px; border-radius: 6px; font-weight: bold; font-size: 14px; } QPushButton:hover { background-color: #0051d5; } QPushButton:pressed { background-color: #003d99; } QPushButton:disabled { background-color: #c7c7cc; color: #999; } QTableWidget { background-color: white; border: 1px solid #ddd; border-radius: 6px; gridline-color: #eee; } QHeaderView::section { background-color: #f5f5f7; padding: 5px; border: none; border-right: 1px solid #ddd; border-bottom: 1px solid #ddd; font-weight: bold; } QProgressBar { border: 1px solid #ddd; border-radius: 5px; text-align: center; font-weight: bold; } QProgressBar::chunk { background-color: #007aff; border-radius: 4px; } QComboBox { padding: 5px; border: 1px solid #ddd; border-radius: 5px; background-color: white; } QCheckBox { spacing: 8px; } QRadioButton { spacing: 8px; } QGroupBox { font-weight: bold; border: 1px solid #ddd; border-radius: 6px; margin-top: 10px; padding-top: 10px; } QGroupBox::title { subcontrol-origin: margin; left: 10px; padding: 0 5px 0 5px; } QListWidget { border: 1px solid #ddd; border-radius: 6px; } QTabWidget::pane { border: 1px solid #ddd; border-radius: 6px; } QTabBar::tab { padding: 8px 16px; background-color: #f5f5f7; border: 1px solid #ddd; border-bottom: none; border-top-left-radius: 6px; border-top-right-radius: 6px; } QTabBar::tab:selected { background-color: white; border-bottom: none; } QScrollArea { border: 1px solid #ddd; border-radius: 6px; } QSpinBox { padding: 5px; border: 1px solid #ddd; border-radius: 5px; } QLineEdit { padding: 5px; border: 1px solid #ddd; border-radius: 5px; } """) self.init_ui() def init_ui(self): # 主窗口部件 main_widget = QWidget() main_layout = QVBoxLayout() main_widget.setLayout(main_layout) self.setCentralWidget(main_widget) # 顶部控制面板 - 减小高度 control_panel = QFrame() control_panel.setFrameShape(QFrame.StyledPanel) control_panel.setMaximumHeight(60) # 限制最大高度 control_layout = QHBoxLayout() control_panel.setLayout(control_layout) control_layout.setContentsMargins(5, 5, 5, 5) # 减小内边距 # 文件选择按钮 self.file_btn = QPushButton("选择PDF文件") self.file_btn.clicked.connect(self.select_pdf_file) control_layout.addWidget(self.file_btn) # 文件路径显示 self.file_path_label = QLabel("未选择文件") self.file_path_label.setMinimumWidth(300) control_layout.addWidget(self.file_path_label) # 模型选择 control_layout.addWidget(QLabel("选择模型:")) self.model_combo = QComboBox() # 定义所有官方重点模型 self.official_models = { "PP-DocLayout_plus-L": { "name": "PP-DocLayout_plus-L (最高精度)", "description": "基于RT-DETR-L的高精度版面区域定位模型", "mAP": "83.2", "size": "126.01 M" }, "PP-DocBlockLayout": { "name": "PP-DocBlockLayout (子模块检测)", "description": "基于RT-DETR-L的文档图像版面子模块检测模型", "mAP": "95.9", "size": "123.92 M" }, "PP-DocLayout-L": { "name": "PP-DocLayout-L (高精度)", "description": "基于RT-DETR-L的高精度版面区域定位模型", "mAP": "90.4", "size": "123.76 M" }, "PP-DocLayout-M": { "name": "PP-DocLayout-M (平衡)", "description": "基于PicoDet-L的精度效率平衡模型", "mAP": "75.2", "size": "22.578 M" }, "PP-DocLayout-S": { "name": "PP-DocLayout-S (快速)", "description": "基于PicoDet-S的高效率版面区域定位模型", "mAP": "70.9", "size": "4.834 M" } } # 检查本地模型目录 self.available_models = {} if os.path.exists(self.model_dir): for model_id in self.official_models: model_path = os.path.join(self.model_dir, model_id) if os.path.exists(model_path): self.available_models[model_id] = self.official_models[model_id] self.available_models[model_id]["local"] = True self.available_models[model_id]["path"] = model_path else: self.available_models[model_id] = self.official_models[model_id] self.available_models[model_id]["local"] = False else: self.available_models = self.official_models.copy() for model_id in self.available_models: self.available_models[model_id]["local"] = False # 添加模型到下拉框 for model_id, model_info in self.available_models.items(): display_name = model_info["name"] if model_info["local"]: display_name += " [本地]" self.model_combo.addItem(display_name, model_id) self.model_combo.setCurrentIndex(0) control_layout.addWidget(self.model_combo) # 模型信息标签 self.model_info_label = QLabel("") self.model_info_label.setMinimumWidth(200) control_layout.addWidget(self.model_info_label) # 更新模型信息 self.update_model_info() self.model_combo.currentIndexChanged.connect(self.update_model_info) # 分析按钮 self.analyze_btn = QPushButton("开始分析") self.analyze_btn.clicked.connect(self.start_analysis) self.analyze_btn.setEnabled(False) control_layout.addWidget(self.analyze_btn) control_layout.addStretch() main_layout.addWidget(control_panel) # 进度条 self.progress_bar = QProgressBar() self.progress_bar.setVisible(False) main_layout.addWidget(self.progress_bar) # 主内容区域 content_splitter = QSplitter(Qt.Horizontal) # 左侧图像显示区域 image_widget = QWidget() image_layout = QVBoxLayout() image_widget.setLayout(image_layout) # 图像控制面板 image_control_panel = QFrame() image_control_layout = QHBoxLayout() image_control_panel.setLayout(image_control_layout) # 缩放控制 image_control_layout.addWidget(QLabel("缩放:")) self.zoom_out_btn = QPushButton("-") self.zoom_out_btn.setMaximumWidth(30) self.zoom_out_btn.clicked.connect(self.zoom_out) image_control_layout.addWidget(self.zoom_out_btn) self.zoom_slider = QSlider(Qt.Horizontal) self.zoom_slider.setMinimum(10) self.zoom_slider.setMaximum(500) self.zoom_slider.setValue(100) self.zoom_slider.setTickPosition(QSlider.TicksBelow) self.zoom_slider.setTickInterval(50) self.zoom_slider.valueChanged.connect(self.on_zoom_changed) image_control_layout.addWidget(self.zoom_slider) self.zoom_in_btn = QPushButton("+") self.zoom_in_btn.setMaximumWidth(30) self.zoom_in_btn.clicked.connect(self.zoom_in) image_control_layout.addWidget(self.zoom_in_btn) self.fit_window_btn = QPushButton("适应窗口") self.fit_window_btn.clicked.connect(self.fit_to_window) image_control_layout.addWidget(self.fit_window_btn) image_control_layout.addStretch() image_layout.addWidget(image_control_panel) # 标签页 self.image_tabs = QTabWidget() # 原始图像标签页 self.original_image_label = ScrollableImageLabel() self.image_tabs.addTab(self.original_image_label, "原始图像") # 标注图像标签页 self.annotated_image_label = ScrollableImageLabel() self.image_tabs.addTab(self.annotated_image_label, "版面标注图像") image_layout.addWidget(self.image_tabs) # 页面导航和导出控制面板 nav_export_panel = QFrame() nav_export_layout = QHBoxLayout() nav_export_panel.setLayout(nav_export_layout) # 页面导航 nav_export_layout.addWidget(QLabel("页面:")) self.page_combo = QComboBox() self.page_combo.setMinimumWidth(80) self.page_combo.currentIndexChanged.connect(self.on_page_changed) nav_export_layout.addWidget(self.page_combo) self.page_label = QLabel("第 0 页，共 0 页") nav_export_layout.addWidget(self.page_label) self.prev_btn = QPushButton("上一页") self.prev_btn.clicked.connect(self.prev_page) self.prev_btn.setEnabled(False) nav_export_layout.addWidget(self.prev_btn) self.next_btn = QPushButton("下一页") self.next_btn.clicked.connect(self.next_page) self.next_btn.setEnabled(False) nav_export_layout.addWidget(self.next_btn) nav_export_layout.addStretch() # 导出按钮 self.export_image_btn = QPushButton("导出区块图像") self.export_image_btn.clicked.connect(self.export_block_images) self.export_image_btn.setEnabled(False) nav_export_layout.addWidget(self.export_image_btn) self.export_pdf_btn = QPushButton("导出标注PDF") self.export_pdf_btn.clicked.connect(self.export_annotated_pdf) self.export_pdf_btn.setEnabled(False) nav_export_layout.addWidget(self.export_pdf_btn) # 分栏裁切按钮 self.column_crop_btn = QPushButton("分栏裁切") self.column_crop_btn.clicked.connect(self.column_crop) self.column_crop_btn.setEnabled(False) nav_export_layout.addWidget(self.column_crop_btn) # 图像掩码按钮 self.image_mask_btn = QPushButton("图像掩码") self.image_mask_btn.clicked.connect(self.image_mask) self.image_mask_btn.setEnabled(False) nav_export_layout.addWidget(self.image_mask_btn) image_layout.addWidget(nav_export_panel) content_splitter.addWidget(image_widget) # 右侧检测结果区域 result_widget = QWidget() result_layout = QVBoxLayout() result_widget.setLayout(result_layout) result_layout.addWidget(QLabel("检测结果")) # 检测结果表格 self.result_table = QTableWidget() self.result_table.setColumnCount(5) self.result_table.setHorizontalHeaderLabels(["区块类型", "置信度", "X坐标", "Y坐标", "宽高"]) self.result_table.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch) self.result_table.setAlternatingRowColors(True) self.result_table.setSelectionBehavior(QTableWidget.SelectRows) self.result_table.setEditTriggers(QTableWidget.NoEditTriggers) result_layout.addWidget(self.result_table) # 统计信息 self.stats_label = QLabel("统计信息: 未分析") result_layout.addWidget(self.stats_label) content_splitter.addWidget(result_widget) content_splitter.setSizes([800, 600]) main_layout.addWidget(content_splitter) # 设置默认文件路径 default_path = "E:\\新建12.pdf" if os.path.exists(default_path): self.current_pdf_path = default_path self.file_path_label.setText(f"已选择: {default_path}") self.analyze_btn.setEnabled(True) def update_model_info(self): """更新模型信息显示""" current_index = self.model_combo.currentIndex() if current_index >= 0: model_id = self.model_combo.itemData(current_index) model_info = self.available_models[model_id] info_text = f"mAP: {model_info['mAP']} | 大小: {model_info['size']}" if model_info["local"]: info_text += " | 已下载" else: info_text += " | 在线加载" self.model_info_label.setText(info_text) def select_pdf_file(self): file_path, _ = QFileDialog.getOpenFileName( self, "选择PDF文件", "", "PDF文件 (*.pdf);;所有文件 (*)" ) if file_path: self.current_pdf_path = file_path self.file_path_label.setText(f"已选择: {file_path}") self.analyze_btn.setEnabled(True) self.reset_ui() def start_analysis(self): if not self.current_pdf_path: QMessageBox.warning(self, "警告", "请先选择PDF文件") return self.reset_ui() self.progress_bar.setVisible(True) self.progress_bar.setValue(0) self.analyze_btn.setEnabled(False) self.file_btn.setEnabled(False) current_index = self.model_combo.currentIndex() model_id = self.model_combo.itemData(current_index) model_info = self.available_models[model_id] threshold = {0: 0.7, 2: 0.5} layout_unclip_ratio = {2: (1, 1.05)} layout_merge_bboxes_mode = {2: "large", 8: "small"} self.worker_thread = LayoutAnalyzerThread( self.current_pdf_path, model_id, model_info["path"] if model_info["local"] else None, threshold=threshold, layout_unclip_ratio=layout_unclip_ratio, layout_merge_bboxes_mode=layout_merge_bboxes_mode ) self.worker_thread.progress_updated.connect(self.update_progress) self.worker_thread.analysis_complete.connect(self.analysis_finished) self.worker_thread.error_occurred.connect(self.show_error) self.worker_thread.start() def update_progress(self, value): self.progress_bar.setValue(value) def analysis_finished(self, results, all_detections): self.analysis_results = results['pages'] self.all_detections = all_detections self.current_page = 0 # 更新UI self.progress_bar.setVisible(False) self.analyze_btn.setEnabled(True) self.file_btn.setEnabled(True) self.export_image_btn.setEnabled(True) self.export_pdf_btn.setEnabled(True) self.column_crop_btn.setEnabled(True) self.image_mask_btn.setEnabled(True) # 更新页面导航 total_pages = len(self.analysis_results) self.page_label.setText(f"第 {self.current_page + 1} 页，共 {total_pages} 页") # 更新页面下拉框 self.page_combo.clear() for i in range(total_pages): self.page_combo.addItem(f"第 {i + 1} 页") self.page_combo.setCurrentIndex(self.current_page) self.prev_btn.setEnabled(self.current_page > 0) self.next_btn.setEnabled(self.current_page < total_pages - 1) # 显示当前页 self.display_current_page() # 更新统计信息 self.update_statistics() def show_error(self, message): self.progress_bar.setVisible(False) self.analyze_btn.setEnabled(True) self.file_btn.setEnabled(True) QMessageBox.critical(self, "错误", message) def reset_ui(self): self.original_image_label.set_image(np.zeros((100, 100, 3), dtype=np.uint8)) self.annotated_image_label.set_image(np.zeros((100, 100, 3), dtype=np.uint8)) self.result_table.setRowCount(0) self.page_label.setText("第 0 页，共 0 页") self.page_combo.clear() self.prev_btn.setEnabled(False) self.next_btn.setEnabled(False) self.stats_label.setText("统计信息: 未分析") self.export_image_btn.setEnabled(False) self.export_pdf_btn.setEnabled(False) self.column_crop_btn.setEnabled(False) self.image_mask_btn.setEnabled(False) def display_current_page(self): if not self.analysis_results or self.current_page >= len(self.analysis_results): return page_data = self.analysis_results[self.current_page] # 显示原始图像 original_img = page_data['image'] self.original_image_label.set_image(original_img) # 显示标注图像 annotated_img = self.create_annotated_image(original_img, page_data['detections']) self.annotated_image_label.set_image(annotated_img) # 更新结果表格 self.update_result_table(page_data['detections']) def create_annotated_image(self, img_array, detections): # 创建图像副本 annotated_img = img_array.copy() # 确保图像是BGR格式（OpenCV默认） if len(annotated_img.shape) == 2: # 灰度图 annotated_img = cv2.cvtColor(annotated_img, cv2.COLOR_GRAY2BGR) elif annotated_img.shape[2] == 4: # RGBA annotated_img = cv2.cvtColor(annotated_img, cv2.COLOR_RGBA2BGR) elif annotated_img.shape[2] == 3 and not np.array_equal(annotated_img[0, 0], img_array[0, 0]): # 可能是RGB annotated_img = cv2.cvtColor(annotated_img, cv2.COLOR_RGB2BGR) # 为不同区块类型定义颜色 colors = { # 基础文本元素 'doc_title': (0, 0, 255), # 红色 - 文档标题 'paragraph_title': (255, 0, 128), # 玫瑰红 - 段落标题 'text': (0, 255, 0), # 绿色 - 文本 'number': (128, 128, 128), # 灰色 - 页码 # 文档结构元素 'abstract': (128, 128, 0), # 橄榄色 - 摘要 'content': (64, 128, 0), # 深橄榄绿 - 目录 'reference': (42, 42, 165), # 深蓝 - 参考文献 'footnote': (192, 192, 192), # 银色 - 脚注 'header': (255, 192, 203), # 粉色 - 页眉 'footer': (173, 216, 230), # 浅蓝 - 页脚 # 特殊内容元素 'algorithm': (75, 0, 130), # 靛蓝 - 算法 'formula': (255, 255, 0), # 黄色 - 公式 'formula_number': (128, 0, 255), # 紫罗兰 - 公式编号 # 图像和表格元素 'image': (255, 128, 0), # 深橙 - 图像 'figure_title': (255, 0, 64), # 深玫瑰红 - 图像标题 'chart': (0, 128, 255), # 天蓝 - 图表 'chart_title': (0, 255, 128), # 春绿色 - 图表标题 'table': (0, 0, 192), # 深红色 - 表格 'table_title': (128, 255, 255), # 浅青色 - 表格标题 # 其他元素 'seal': (128, 0, 128), # 深紫色 - 印章 'header_image': (192, 255, 62), # 酸橙绿 - 页眉图像 'footer_image': (62, 255, 192), # 蓝绿 - 页脚图像 'aside_text': (128, 64, 0), # 棕色 - 侧栏文本 } # 绘制边界框 for detection in detections: label = detection['label'] coords = detection['coordinates'] x1, y1, x2, y2 = map(int, coords) # 获取颜色 color = colors.get(label, (255, 255, 255)) # 绘制矩形 cv2.rectangle(annotated_img, (x1, y1), (x2, y2), color, 2) # 添加标签 label_text = f"{label} ({detection['confidence']:.2f})" # 确保文本不会超出图像边界 text_y = max(y1 - 10, 20) cv2.putText( annotated_img, label_text, (x1, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1 ) return annotated_img def update_result_table(self, detections): self.result_table.setRowCount(len(detections)) for i, detection in enumerate(detections): # 区块类型 self.result_table.setItem(i, 0, QTableWidgetItem(detection['label'])) # 置信度 conf_item = QTableWidgetItem(f"{detection['confidence']:.4f}") conf_item.setTextAlignment(Qt.AlignCenter) self.result_table.setItem(i, 1, conf_item) # 坐标 coords = detection['coordinates'] x1, y1, x2, y2 = map(int, coords) self.result_table.setItem(i, 2, QTableWidgetItem(f"{x1}, {x2}")) self.result_table.setItem(i, 3, QTableWidgetItem(f"{y1}, {y2}")) # 宽高 width = x2 - x1 height = y2 - y1 self.result_table.setItem(i, 4, QTableWidgetItem(f"{width} × {height}")) # 设置行高 self.result_table.setRowHeight(i, 25) def update_statistics(self): if not self.all_detections: self.stats_label.setText("统计信息: 未分析") return # 统计各类型区块数量 type_counts = {} for detection in self.all_detections: label = detection['label'] type_counts[label] = type_counts.get(label, 0) + 1 # 构建统计文本 stats_text = "统计信息: " stats_items = [] for label, count in sorted(type_counts.items()): stats_items.append(f"{label}: {count}") stats_text += " | ".join(stats_items) self.stats_label.setText(stats_text) def on_page_changed(self, index): """页面下拉框改变""" if index >= 0 and index < len(self.analysis_results): self.current_page = index self.update_page_display() def prev_page(self): if self.current_page > 0: self.current_page -= 1 self.update_page_display() def next_page(self): if self.current_page < len(self.analysis_results) - 1: self.current_page += 1 self.update_page_display() def update_page_display(self): # 更新页面标签 total_pages = len(self.analysis_results) self.page_label.setText(f"第 {self.current_page + 1} 页，共 {total_pages} 页") # 更新页面下拉框 self.page_combo.setCurrentIndex(self.current_page) # 更新按钮状态 self.prev_btn.setEnabled(self.current_page > 0) self.next_btn.setEnabled(self.current_page < total_pages - 1) # 显示当前页 self.display_current_page() # 图像缩放和拖动功能 def on_zoom_changed(self, value): """缩放滑块值改变""" scale = value / 100.0 # 获取当前标签页 current_tab = self.image_tabs.currentWidget() if current_tab == self.original_image_label: self.original_image_label.set_scale(scale) elif current_tab == self.annotated_image_label: self.annotated_image_label.set_scale(scale) def zoom_in(self): """放大图像""" # 获取当前标签页 current_tab = self.image_tabs.currentWidget() if current_tab == self.original_image_label: self.original_image_label.zoom_in() elif current_tab == self.annotated_image_label: self.annotated_image_label.zoom_in() # 更新滑块值 self.zoom_slider.setValue(int(self.zoom_slider.value() * 1.25)) def zoom_out(self): """缩小图像""" # 获取当前标签页 current_tab = self.image_tabs.currentWidget() if current_tab == self.original_image_label: self.original_image_label.zoom_out() elif current_tab == self.annotated_image_label: self.annotated_image_label.zoom_out() # 更新滑块值 self.zoom_slider.setValue(int(self.zoom_slider.value() * 0.8)) def fit_to_window(self): """适应窗口大小""" # 获取当前标签页 current_tab = self.image_tabs.currentWidget() if current_tab == self.original_image_label: self.original_image_label.fit_to_window() elif current_tab == self.annotated_image_label: self.annotated_image_label.fit_to_window() # 更新滑块值 self.zoom_slider.setValue(100) # 导出区块图像功能 def export_block_images(self): """导出区块图像""" if not self.analysis_results: QMessageBox.warning(self, "警告", "没有可导出的内容") return # 获取所有区块类型 block_types = set() for page in self.analysis_results: for detection in page['detections']: block_types.add(detection['label']) block_types = sorted(list(block_types)) # 显示图像导出选项对话框 export_dialog = ImageExportOptionsDialog(block_types, self) if export_dialog.exec_() != QDialog.Accepted: return export_options = export_dialog.get_export_options() # 检查是否选择了区块类型 if not export_options["selected_types"]: QMessageBox.warning(self, "警告", "请至少选择一种区块类型") return # 检查是否选择了输出目录 if not export_options["output_dir"]: QMessageBox.warning(self, "警告", "请选择输出目录") return # 确定要处理的页面 if export_options["page_selection"] == "all": pages_to_process = range(len(self.analysis_results)) elif export_options["page_selection"] == "current": pages_to_process = [self.current_page] else: # 显示页面选择对话框 page_dialog = PageSelectionDialog(len(self.analysis_results), self.current_page, self) if page_dialog.exec_() != QDialog.Accepted: return selected_pages = page_dialog.get_selected_pages() if selected_pages is None: QMessageBox.warning(self, "警告", "请至少选择一页") return if selected_pages == "all": pages_to_process = range(len(self.analysis_results)) elif selected_pages == "current": pages_to_process = [self.current_page] else: pages_to_process = selected_pages # 添加当前页面和选择的页面到导出选项 export_options["current_page"] = self.current_page export_options["selected_pages"] = pages_to_process # 创建输出目录结构 os.makedirs(export_options["output_dir"], exist_ok=True) # 为每种区块类型创建子目录 for block_type in export_options["selected_types"]: type_dir = os.path.join(export_options["output_dir"], block_type) os.makedirs(type_dir, exist_ok=True) # 显示进度对话框 progress_dialog = QDialog(self) progress_dialog.setWindowTitle("导出进度") progress_dialog.setFixedSize(400, 150) progress_layout = QVBoxLayout() progress_label = QLabel("准备导出区块图像...") progress_label.setAlignment(Qt.AlignCenter) progress_layout.addWidget(progress_label) progress_bar = QProgressBar() progress_layout.addWidget(progress_bar) # 添加取消按钮 cancel_btn = QPushButton("取消") cancel_btn.clicked.connect(self.cancel_image_export) progress_layout.addWidget(cancel_btn) progress_dialog.setLayout(progress_layout) progress_dialog.show() # 创建并启动导出线程 self.image_export_thread = ImageExportThread( self.analysis_results, export_options ) self.image_export_thread.progress_updated.connect( lambda value, text: ( progress_bar.setValue(value), progress_label.setText(text) ) ) self.image_export_thread.export_complete.connect( lambda result: ( progress_dialog.close(), QMessageBox.information(self, "导出成功", f"区块图像已成功导出到:\n{result}") ) ) self.image_export_thread.export_failed.connect( lambda error: ( progress_dialog.close(), QMessageBox.critical(self, "导出失败", f"导出区块图像时出错:\n{error}") ) ) self.image_export_thread.start() def cancel_image_export(self): """取消图像导出操作""" if hasattr(self, 'image_export_thread') and self.image_export_thread.isRunning(): self.image_export_thread.stop() self.image_export_thread.wait() # 导出标注PDF功能 def export_annotated_pdf(self): """导出标注PDF""" if not self.analysis_results: QMessageBox.warning(self, "警告", "没有可导出的内容") return # 显示PDF导出选项对话框 export_dialog = PDFExportOptionsDialog(self) if export_dialog.exec_() != QDialog.Accepted: return export_options = export_dialog.get_export_options() # 检查是否选择了输出文件 if not export_options["output_file"]: QMessageBox.warning(self, "警告", "请选择输出文件") return # 确定要处理的页面 if export_options["page_selection"] == "all": pages_to_process = range(len(self.analysis_results)) elif export_options["page_selection"] == "current": pages_to_process = [self.current_page] else: # 显示页面选择对话框 page_dialog = PageSelectionDialog(len(self.analysis_results), self.current_page, self) if page_dialog.exec_() != QDialog.Accepted: return selected_pages = page_dialog.get_selected_pages() if selected_pages is None: QMessageBox.warning(self, "警告", "请至少选择一页") return if selected_pages == "all": pages_to_process = range(len(self.analysis_results)) elif selected_pages == "current": pages_to_process = [self.current_page] else: pages_to_process = selected_pages # 添加当前页面和选择的页面到导出选项 export_options["current_page"] = self.current_page export_options["selected_pages"] = pages_to_process # 显示进度对话框 progress_dialog = QDialog(self) progress_dialog.setWindowTitle("导出进度") progress_dialog.setFixedSize(400, 150) progress_layout = QVBoxLayout() progress_label = QLabel("准备导出PDF...") progress_label.setAlignment(Qt.AlignCenter) progress_layout.addWidget(progress_label) progress_bar = QProgressBar() progress_layout.addWidget(progress_bar) # 添加取消按钮 cancel_btn = QPushButton("取消") cancel_btn.clicked.connect(self.cancel_pdf_export) progress_layout.addWidget(cancel_btn) progress_dialog.setLayout(progress_layout) progress_dialog.show() # 创建并启动导出线程 self.pdf_export_thread = PDFExportThread( self.analysis_results, export_options, file_path=export_options["output_file"] ) self.pdf_export_thread.progress_updated.connect( lambda value, text: ( progress_bar.setValue(value), progress_label.setText(text) ) ) self.pdf_export_thread.export_complete.connect( lambda result: ( progress_dialog.close(), QMessageBox.information(self, "导出成功", f"PDF已成功导出到:\n{result}") ) ) self.pdf_export_thread.export_failed.connect( lambda error: ( progress_dialog.close(), QMessageBox.critical(self, "导出失败", f"导出PDF时出错:\n{error}") ) ) self.pdf_export_thread.start() def cancel_pdf_export(self): """取消PDF导出操作""" if hasattr(self, 'pdf_export_thread') and self.pdf_export_thread.isRunning(): self.pdf_export_thread.stop() self.pdf_export_thread.wait() # 分栏裁切功能 def column_crop(self): """分栏裁切""" if not self.analysis_results: QMessageBox.warning(self, "警告", "没有可裁切的内容") return # 获取所有区块类型 block_types = set() for page in self.analysis_results: for detection in page['detections']: block_types.add(detection['label']) block_types = sorted(list(block_types)) # 显示分栏裁切选项对话框 crop_dialog = ColumnCropDialog(block_types, self) if crop_dialog.exec_() != QDialog.Accepted: return crop_options = crop_dialog.get_crop_options() # 检查是否选择了输出目录 if not crop_options["output_dir"]: QMessageBox.warning(self, "警告", "请选择输出目录") return # 创建输出目录 os.makedirs(crop_options["output_dir"], exist_ok=True) # 显示进度对话框 progress_dialog = QDialog(self) progress_dialog.setWindowTitle("分栏裁切进度") progress_dialog.setFixedSize(400, 150) progress_layout = QVBoxLayout() progress_label = QLabel("准备分栏裁切...") progress_label.setAlignment(Qt.AlignCenter) progress_layout.addWidget(progress_label) progress_bar = QProgressBar() progress_layout.addWidget(progress_bar) # 添加取消按钮 cancel_btn = QPushButton("取消") cancel_btn.clicked.connect(self.cancel_column_crop) progress_layout.addWidget(cancel_btn) progress_dialog.setLayout(progress_layout) progress_dialog.show() # 创建并启动分栏裁切线程 self.column_crop_thread = ColumnCropThread( self.analysis_results, crop_options ) self.column_crop_thread.progress_updated.connect( lambda value, text: ( progress_bar.setValue(value), progress_label.setText(text) ) ) self.column_crop_thread.crop_complete.connect( lambda result: ( progress_dialog.close(), QMessageBox.information(self, "裁切完成", f"分栏裁切已完成，结果保存在:\n{result}") ) ) self.column_crop_thread.crop_failed.connect( lambda error: ( progress_dialog.close(), QMessageBox.critical(self, "裁切失败", f"分栏裁切时出错:\n{error}") ) ) self.column_crop_thread.start() def cancel_column_crop(self): """取消分栏裁切操作""" if hasattr(self, 'column_crop_thread') and self.column_crop_thread.isRunning(): self.column_crop_thread.stop() self.column_crop_thread.wait() # 图像掩码功能 def image_mask(self): """图像掩码""" if not self.analysis_results: QMessageBox.warning(self, "警告", "没有可处理的内容") return # 获取所有区块类型 block_types = set() for page in self.analysis_results: for detection in page['detections']: block_types.add(detection['label']) block_types = sorted(list(block_types)) # 显示图像掩码选项对话框 mask_dialog = ImageMaskDialog(block_types, self) if mask_dialog.exec_() != QDialog.Accepted: return mask_options = mask_dialog.get_mask_options() # 检查是否选择了区块类型 if not mask_options["mosaic_types"]: QMessageBox.warning(self, "警告", "请至少选择一种区块类型进行马赛克处理") return # 检查是否选择了输出目录 if not mask_options["output_dir"]: QMessageBox.warning(self, "警告", "请选择输出目录") return # 确定要处理的页面 if mask_options["page_selection"] == "all": pages_to_process = range(len(self.analysis_results)) elif mask_options["page_selection"] == "current": pages_to_process = [self.current_page] else: # 显示页面选择对话框 page_dialog = PageSelectionDialog(len(self.analysis_results), self.current_page, self) if page_dialog.exec_() != QDialog.Accepted: return selected_pages = page_dialog.get_selected_pages() if selected_pages is None: QMessageBox.warning(self, "警告", "请至少选择一页") return if selected_pages == "all": pages_to_process = range(len(self.analysis_results)) elif selected_pages == "current": pages_to_process = [self.current_page] else: pages_to_process = selected_pages # 添加当前页面和选择的页面到掩码选项 mask_options["current_page"] = self.current_page mask_options["selected_pages"] = pages_to_process # 创建输出目录 os.makedirs(mask_options["output_dir"], exist_ok=True) # 显示进度对话框 progress_dialog = QDialog(self) progress_dialog.setWindowTitle("图像掩码进度") progress_dialog.setFixedSize(400, 150) progress_layout = QVBoxLayout() progress_label = QLabel("准备进行图像掩码处理...") progress_label.setAlignment(Qt.AlignCenter) progress_layout.addWidget(progress_label) progress_bar = QProgressBar() progress_layout.addWidget(progress_bar) # 添加取消按钮 cancel_btn = QPushButton("取消") cancel_btn.clicked.connect(self.cancel_image_mask) progress_layout.addWidget(cancel_btn) progress_dialog.setLayout(progress_layout) progress_dialog.show() # 创建并启动图像掩码线程 self.image_mask_thread = ImageMaskThread( self.analysis_results, mask_options ) self.image_mask_thread.progress_updated.connect( lambda value, text: ( progress_bar.setValue(value), progress_label.setText(text) ) ) self.image_mask_thread.mask_complete.connect( lambda result: ( progress_dialog.close(), QMessageBox.information(self, "处理完成", f"图像掩码处理已完成，结果保存在:\n{result}") ) ) self.image_mask_thread.mask_failed.connect( lambda error: ( progress_dialog.close(), QMessageBox.critical(self, "处理失败", f"图像掩码处理时出错:\n{error}") ) ) self.image_mask_thread.start() def cancel_image_mask(self): """取消图像掩码操作""" if hasattr(self, 'image_mask_thread') and self.image_mask_thread.isRunning(): self.image_mask_thread.stop() self.image_mask_thread.wait() if __name__ == "__main__": app = QApplication(sys.argv) # 设置应用程序样式 app.setStyle("Fusion") # 创建主窗口 window = PDFLayoutAnalyzer() window.show() sys.exit(app.exec_())