[python]智能PDF章节拆分工具——增强识别版本

   import reimport osimport fitz  # PyMuPDFimport pdfplumberfrom pathlib import Pathfrom typing import List, Dict, Tuple, Optionalimport tkinter as tkfrom tkinter import ttk, filedialog, messagebox, scrolledtextfrom dataclasses import dataclassimport threadingimport timeimport numpy as npfrom collections import Counter@dataclassclass ChapterInfo:    """章节信息数据类"""    title: str    start_page: int    end_page: int    level: int    confidence: float    source: str    font_size: float = 12.0    is_bold: bool = False    bbox: Tuple[float, float, float, float] = Noneclass ImprovedPDFChapterSplitter:    def __init__(self):        self.root = tk.Tk()        self.root.title("智能PDF章节拆分工具 v3.1 - 增强识别版")        self.root.geometry("1200x800")        self.pdf_path = None        self.chapters = []        self.setup_ui()    def setup_ui(self):        """设置用户界面"""        # 主框架        main_frame = ttk.Frame(self.root, padding="10")        main_frame.pack(fill=tk.BOTH, expand=True)        # 文件选择区域        file_frame = ttk.LabelFrame(main_frame, text="PDF文件选择", padding="10")        file_frame.pack(fill=tk.X, pady=5)        file_grid = ttk.Frame(file_frame)        file_grid.pack(fill=tk.X)        ttk.Button(file_grid, text="选择PDF文件",                   command=self.select_pdf_file).grid(row=0, column=0, padx=5)        self.file_label = ttk.Label(file_grid, text="未选择文件")        self.file_label.grid(row=0, column=1, padx=5, sticky=tk.W)        # 高级选项区域        options_frame = ttk.LabelFrame(main_frame, text="高级识别选项", padding="10")        options_frame.pack(fill=tk.X, pady=5)        options_grid = ttk.Frame(options_frame)        options_grid.pack(fill=tk.X)        # 识别模式        ttk.Label(options_grid, text="识别模式:").grid(row=0, column=0, sticky=tk.W, padx=5)        self.analysis_mode = tk.StringVar(value="advanced_hybrid")        modes = [            ("高级混合模式", "advanced_hybrid"),            ("智能视觉分析", "visual_analysis"),            ("结构分析优先", "structure_first"),            ("仅书签提取", "outline_only")        ]        for i, (text, value) in enumerate(modes):            ttk.Radiobutton(options_grid, text=text, variable=self.analysis_mode,                            value=value).grid(row=0, column=i * 2 + 1, sticky=tk.W, padx=5)        # 识别参数        param_frame = ttk.Frame(options_grid)        param_frame.grid(row=1, column=0, columnspan=8, pady=10, sticky=tk.W)        ttk.Label(param_frame, text="字体大小阈值:").grid(row=0, column=0, sticky=tk.W, padx=5)        self.font_ratio = tk.DoubleVar(value=1.2)        ttk.Scale(param_frame, from_=1.0, to=2.0, variable=self.font_ratio,                  orient=tk.HORIZONTAL, length=100).grid(row=0, column=1, padx=5)        self.font_label = ttk.Label(param_frame, text="1.2")        self.font_label.grid(row=0, column=2, padx=5)        ttk.Label(param_frame, text="最小标题长度:").grid(row=0, column=3, sticky=tk.W, padx=5)        self.min_title_len = tk.IntVar(value=2)        ttk.Entry(param_frame, textvariable=self.min_title_len, width=5).grid(row=0, column=4, padx=5)        ttk.Label(param_frame, text="置信度阈值:").grid(row=0, column=5, sticky=tk.W, padx=5)        self.confidence_thresh = tk.DoubleVar(value=0.4)        ttk.Scale(param_frame, from_=0.1, to=1.0, variable=self.confidence_thresh,                  orient=tk.HORIZONTAL, length=100).grid(row=0, column=6, padx=5)        self.conf_label = ttk.Label(param_frame, text="0.4")        self.conf_label.grid(row=0, column=7, padx=5)        # 控制按钮        control_frame = ttk.Frame(main_frame)        control_frame.pack(fill=tk.X, pady=5)        btn_style = ttk.Style()        btn_style.configure("Action.TButton", font=("Arial", 10, "bold"))        ttk.Button(control_frame, text="开始分析", style="Action.TButton",                   command=self.start_analysis).pack(side=tk.LEFT, padx=5)        ttk.Button(control_frame, text="重新分析",                   command=self.reanalyze).pack(side=tk.LEFT, padx=5)        ttk.Button(control_frame, text="手动校正",                   command=self.open_correction_tool).pack(side=tk.LEFT, padx=5)        ttk.Button(control_frame, text="导出选中",                   command=self.export_selected).pack(side=tk.LEFT, padx=5)        # 章节列表区域        chapter_frame = ttk.LabelFrame(main_frame, text="识别结果", padding="5")        chapter_frame.pack(fill=tk.BOTH, expand=True, pady=5)        # 创建树形视图        columns = ("title", "pages", "level", "font", "confidence", "source", "status")        self.chapter_tree = ttk.Treeview(chapter_frame, columns=columns, show="headings", height=15)        # 设置列标题        self.chapter_tree.heading("title", text="章节标题", anchor=tk.W)        self.chapter_tree.heading("pages", text="页码范围", anchor=tk.CENTER)        self.chapter_tree.heading("level", text="层级", anchor=tk.CENTER)        self.chapter_tree.heading("font", text="字体大小", anchor=tk.CENTER)        self.chapter_tree.heading("confidence", text="置信度", anchor=tk.CENTER)        self.chapter_tree.heading("source", text="识别来源", anchor=tk.CENTER)        self.chapter_tree.heading("status", text="状态", anchor=tk.CENTER)        # 设置列宽        self.chapter_tree.column("title", width=400, stretch=tk.YES)        self.chapter_tree.column("pages", width=100, stretch=tk.NO)        self.chapter_tree.column("level", width=60, stretch=tk.NO)        self.chapter_tree.column("font", width=80, stretch=tk.NO)        self.chapter_tree.column("confidence", width=80, stretch=tk.NO)        self.chapter_tree.column("source", width=100, stretch=tk.NO)        self.chapter_tree.column("status", width=60, stretch=tk.NO)        # 添加滚动条        scrollbar = ttk.Scrollbar(chapter_frame, orient=tk.VERTICAL,                                  command=self.chapter_tree.yview)        self.chapter_tree.configure(yscrollcommand=scrollbar.set)        self.chapter_tree.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)        # 统计信息区域        stats_frame = ttk.Frame(main_frame)        stats_frame.pack(fill=tk.X, pady=5)        self.stats_label = ttk.Label(stats_frame, text="等待分析...")        self.stats_label.pack(side=tk.LEFT)        # 日志区域        log_frame = ttk.LabelFrame(main_frame, text="处理日志", padding="5")        log_frame.pack(fill=tk.BOTH, expand=False, pady=5)        self.log_text = scrolledtext.ScrolledText(log_frame, height=8, wrap=tk.WORD)        self.log_text.pack(fill=tk.BOTH, expand=True)        # 进度条        self.progress = ttk.Progressbar(main_frame, mode='indeterminate')        self.progress.pack(fill=tk.X, pady=5)        # 绑定事件        self.font_ratio.trace_add("write", self.update_font_label)        self.confidence_thresh.trace_add("write", self.update_conf_label)    def update_font_label(self, *args):        self.font_label.config(text=f"{self.font_ratio.get():.1f}")    def update_conf_label(self, *args):        self.conf_label.config(text=f"{self.confidence_thresh.get():.2f}")    def log_message(self, message: str, level: str = "INFO"):        """添加日志消息"""        colors = {"INFO": "black", "SUCCESS": "green", "WARNING": "orange", "ERROR": "red"}        timestamp = time.strftime("%H:%M:%S")        self.log_text.insert(tk.END, f"[{timestamp}] {message}\n")        # 设置颜色        self.log_text.tag_config(level, foreground=colors.get(level, "black"))        self.log_text.tag_add(level, f"end-2l", f"end-1l")        self.log_text.see(tk.END)        self.root.update()    def select_pdf_file(self):        """选择PDF文件"""        file_path = filedialog.askopenfilename(            title="选择PDF文件",            filetypes=[("PDF文件", "*.pdf"), ("所有文件", "*.*")]        )        if file_path:            self.pdf_path = file_path            filename = Path(file_path).name            self.file_label.config(text=filename)            self.log_message(f"已选择文件: {filename}", "SUCCESS")            self.stats_label.config(text=f"已加载: {filename}")    def start_analysis(self):        """开始分析"""        if not self.pdf_path:            messagebox.showwarning("警告", "请先选择PDF文件")            return        self.progress.start()        self.log_message("开始分析PDF文档结构...", "INFO")        thread = threading.Thread(target=self._analyze_pdf)        thread.daemon = True        thread.start()    def _analyze_pdf(self):        """分析PDF文档"""        try:            mode = self.analysis_mode.get()            if mode == "outline_only":                chapters = self.extract_from_outline()            elif mode == "visual_analysis":                chapters = self.visual_analysis()            elif mode == "structure_first":                chapters = self.structure_based_analysis()            else:  # advanced_hybrid                chapters = self.hybrid_analysis()            # 过滤和排序            filtered_chapters = self.filter_chapters(chapters)            self.chapters = self.adjust_boundaries(filtered_chapters)            # 更新UI            self.root.after(0, self._update_display)        except Exception as e:            self.log_message(f"分析失败: {str(e)}", "ERROR")        finally:            self.root.after(0, self.progress.stop)    def hybrid_analysis(self) -> List[ChapterInfo]:        """混合分析方法"""        all_chapters = []        # 1. 从书签提取        outline_chapters = self.extract_from_outline()        all_chapters.extend(outline_chapters)        self.log_message(f"从书签提取到 {len(outline_chapters)} 个章节", "INFO")        # 2. 视觉分析补充        if len(outline_chapters) < 5:  # 如果书签太少            visual_chapters = self.visual_analysis()            all_chapters.extend(visual_chapters)            self.log_message(f"视觉分析补充 {len(visual_chapters)} 个章节", "INFO")        # 3. 结构分析补充        structure_chapters = self.structure_based_analysis()        all_chapters.extend(structure_chapters)        self.log_message(f"结构分析补充 {len(structure_chapters)} 个章节", "INFO")        return all_chapters    def extract_from_outline(self) -> List[ChapterInfo]:        """从书签提取章节"""        chapters = []        try:            doc = fitz.open(self.pdf_path)            outlines = doc.get_toc()            for level, title, page_num in outlines:                actual_page = max(0, page_num - 1)                # 清理标题                cleaned_title = self.clean_title(title)                chapters.append(ChapterInfo(                    title=cleaned_title,                    start_page=actual_page,                    end_page=actual_page,                    level=min(level, 4),  # 限制最大层级                    confidence=0.9,                    source="书签"                ))            doc.close()        except Exception as e:            self.log_message(f"书签提取错误: {str(e)}", "WARNING")        return chapters    def visual_analysis(self) -> List[ChapterInfo]:        """视觉特征分析"""        chapters = []        try:            doc = fitz.open(self.pdf_path)            font_ratio_thresh = self.font_ratio.get()            for page_num in range(len(doc)):                page = doc.load_page(page_num)                blocks = page.get_text("dict")                # 收集本页所有字体大小                font_sizes = []                for block in blocks["blocks"]:                    if "lines" in block:                        for line in block["lines"]:                            for span in line["spans"]:                                font_sizes.append(span["size"])                if not font_sizes:                    continue                # 计算平均字体大小                avg_font = np.mean(font_sizes)                # 查找大字体文本（可能是标题）                for block in blocks["blocks"]:                    if "lines" in block:                        for line in block["lines"]:                            for span in line["spans"]:                                text = span["text"].strip()                                font_size = span["size"]                                # 检查是否可能是标题                                if (font_size >= avg_font * font_ratio_thresh and                                        len(text) >= self.min_title_len.get() and                                        self.is_potential_title(text)):                                    is_bold = bool(span["flags"] & 2 ** 4)                                    confidence = self.calculate_confidence(text, font_size, avg_font, is_bold)                                    if confidence >= self.confidence_thresh.get():                                        chapters.append(ChapterInfo(                                            title=text,                                            start_page=page_num,                                            end_page=page_num,                                            level=self.determine_level(text, font_size, avg_font),                                            confidence=confidence,                                            source="视觉分析",                                            font_size=font_size,                                            is_bold=is_bold,                                            bbox=span["bbox"]                                        ))            doc.close()        except Exception as e:            self.log_message(f"视觉分析错误: {str(e)}", "WARNING")        return chapters    def structure_based_analysis(self) -> List[ChapterInfo]:        """基于文档结构的分析"""        chapters = []        try:            with pdfplumber.open(self.pdf_path) as pdf:                for page_num, page in enumerate(pdf.pages):                    text = page.extract_text()                    if not text:                        continue                    # 多种中文章节模式                    patterns = [                        # 中文章节                        (r'(第[一二三四五六七八九十零〇\d]+章)\s*([^\n\r]{0,50})', 1.0, 1),                        (r'(第[一二三四五六七八九十零〇\d]+节)\s*([^\n\r]{0,50})', 0.9, 2),                        (r'([一二三四五六七八九十]、)\s*([^\n\r]{0,50})', 0.8, 3),                        (r'(\d+、)\s*([^\n\r]{0,50})', 0.7, 3),                        # 数字编号                        (r'(\d+(\.\d+)+)\s+([^\n\r]{0,50})', 0.8, 2),                        # 中文项目符号                        (r'(（[一二三四五六七八九十]）)\s*([^\n\r]{0,50})', 0.7, 3),                        (r'(（\d+）)\s*([^\n\r]{0,50})', 0.6, 3),                        # 特殊章节                        (r'(前言|摘要|目录|绪论|引言|结论|参考文献|致谢|附录)\b', 0.9, 1),                        # 环境报告专用                        (r'(一、|二、|三、|四、|五、|六、|七、|八、|九、|十、)\s*([^\n\r]{0,50})', 0.9, 2),                        (r'([一二三四五六七八九十]、\s*[^\n\r]{0,30})\n([^\n\r]{0,30})', 0.8, 2),                    ]                    for pattern, base_conf, level in patterns:                        matches = re.finditer(pattern, text, re.MULTILINE)                        for match in matches:                            title = match.group(0).strip()                            if len(title) >= self.min_title_len.get():                                # 合并多行标题                                if '\n' in title:                                    lines = title.split('\n')                                    title = ' '.join([line.strip() for line in lines if line.strip()])                                if self.validate_title(title):                                    chapters.append(ChapterInfo(                                        title=title,                                        start_page=page_num,                                        end_page=page_num,                                        level=level,                                        confidence=base_conf * 0.9,                                        source="结构分析"                                    ))                                    break        except Exception as e:            self.log_message(f"结构分析错误: {str(e)}", "WARNING")        return chapters    def clean_title(self, title: str) -> str:        """清理标题"""        # 去除多余空格        title = re.sub(r'\s+', ' ', title).strip()        # 截断过长的标题        if len(title) > 100:            title = title[:97] + "..."        return title    def is_potential_title(self, text: str) -> bool:        """判断是否为潜在标题"""        if len(text) < 2 or len(text) > 100:            return False        # 过滤掉明显不是标题的文本        invalid_patterns = [            r'^\d+$',            r'^[a-zA-Z]$',            r'^[\s\.\-]+$',            r'^图\d+',            r'^表\d+',            r'^第\d+页',        ]        for pattern in invalid_patterns:            if re.match(pattern, text):                return False        # 检查是否包含中文        if not re.search('[\u4e00-\u9fff]', text):            return False        return True    def calculate_confidence(self, text: str, font_size: float, avg_font: float, is_bold: bool) -> float:        """计算置信度"""        confidence = 0.5        # 字体大小因素        if font_size > avg_font * 1.5:            confidence += 0.3        elif font_size > avg_font * 1.2:            confidence += 0.2        # 粗体因素        if is_bold:            confidence += 0.1        # 长度因素        if 5 <= len(text) <= 50:            confidence += 0.1        # 模式匹配因素        title_patterns = [            r'^第[一二三四五六七八九十\d]+[章节]',            r'^[一二三四五六七八九十]、',            r'^\d+[、\.]',        ]        for pattern in title_patterns:            if re.match(pattern, text):                confidence += 0.2                break        return min(confidence, 1.0)    def determine_level(self, text: str, font_size: float, avg_font: float) -> int:        """确定层级"""        if re.match(r'^第[一二三四五六七八九十\d]+章', text):            return 1        elif re.match(r'^第[一二三四五六七八九十\d]+节', text):            return 2        elif re.match(r'^[一二三四五六七八九十]、', text):            return 2        elif font_size > avg_font * 1.5:            return 1        elif font_size > avg_font * 1.2:            return 2        else:            return 3    def validate_title(self, title: str) -> bool:        """验证标题有效性"""        if len(title) < self.min_title_len.get():            return False        # 不能全是数字或符号        if re.match(r'^[\d\s\.\-]+$', title):            return False        return True    def filter_chapters(self, chapters: List[ChapterInfo]) -> List[ChapterInfo]:        """过滤和去重章节"""        if not chapters:            return []        # 按页码排序        chapters.sort(key=lambda x: (x.start_page, -x.confidence))        # 去重：同页相似的章节只保留置信度最高的        filtered = []        seen_pages = set()        for chapter in chapters:            page_key = f"{chapter.start_page}_{chapter.title[:20]}"            if page_key not in seen_pages and chapter.confidence >= self.confidence_thresh.get():                seen_pages.add(page_key)                filtered.append(chapter)        return filtered    def adjust_boundaries(self, chapters: List[ChapterInfo]) -> List[ChapterInfo]:        """调整章节边界"""        if not chapters:            return []        try:            with fitz.open(self.pdf_path) as doc:                total_pages = len(doc)            # 按页码排序            chapters.sort(key=lambda x: x.start_page)            # 设置结束页码            for i in range(len(chapters)):                if i < len(chapters) - 1:                    end_page = max(chapters[i].start_page, chapters[i + 1].start_page - 1)                else:                    end_page = total_pages - 1                chapters[i].end_page = end_page        except:            pass        return chapters    def _update_display(self):        """更新显示"""        # 清空现有内容        for item in self.chapter_tree.get_children():            self.chapter_tree.delete(item)        # 添加新章节        for i, chapter in enumerate(self.chapters):            page_range = f"{chapter.start_page + 1}-{chapter.end_page + 1}"            font_info = f"{chapter.font_size:.1f}" if hasattr(chapter, 'font_size') else "-"            self.chapter_tree.insert("", "end", values=(                chapter.title,                page_range,                chapter.level,                font_info,                f"{chapter.confidence:.2f}",                chapter.source,                "✓"            ))        # 更新统计信息        self.stats_label.config(text=f"共识别到 {len(self.chapters)} 个章节")        self.log_message(f"分析完成，共识别到 {len(self.chapters)} 个章节", "SUCCESS")    def reanalyze(self):        """重新分析"""        if not self.pdf_path:            return        self.log_message("重新分析文档...", "INFO")        self.start_analysis()    def open_correction_tool(self):        """打开校正工具"""        if not self.chapters:            messagebox.showinfo("提示", "请先分析文档")            return        # 创建校正窗口        dialog = tk.Toplevel(self.root)        dialog.title("章节校正工具")        dialog.geometry("800x600")        # 添加校正功能界面        ttk.Label(dialog, text="章节校正", font=("Arial", 14, "bold")).pack(pady=10)        # 这里可以添加更多校正功能        ttk.Button(dialog, text="合并相邻章节",                   command=self.merge_chapters).pack(pady=5)        ttk.Button(dialog, text="拆分章节",                   command=self.split_chapter).pack(pady=5)        ttk.Button(dialog, text="调整页码",                   command=self.adjust_page_range).pack(pady=5)    def merge_chapters(self):        """合并章节功能"""        # 实现章节合并逻辑        pass    def split_chapter(self):        """拆分章节功能"""        # 实现章节拆分逻辑        pass    def adjust_page_range(self):        """调整页码范围功能"""        # 实现页码范围调整逻辑        pass    def export_selected(self):        """导出选中章节"""        if not self.chapters:            messagebox.showwarning("警告", "没有可导出的章节")            return        output_dir = filedialog.askdirectory(title="选择输出目录")        if not output_dir:            return        self._export_chapters(output_dir)    def _export_chapters(self, output_dir: str):        """导出章节"""        try:            self.progress.start()            doc = fitz.open(self.pdf_path)            output_path = Path(output_dir)            output_path.mkdir(exist_ok=True)            for i, chapter in enumerate(self.chapters):                new_doc = fitz.open()                new_doc.insert_pdf(doc, from_page=chapter.start_page, to_page=chapter.end_page)                # 清理文件名                safe_title = re.sub(r'[\\/*?:"<>|]', "_", chapter.title)                if not safe_title.strip():                    safe_title = f"章节_{i + 1}"                filename = f"{i + 1:03d}_{safe_title}.pdf"                filepath = output_path / filename                new_doc.save(filepath)                new_doc.close()                self.log_message(f"已导出: {filename}", "SUCCESS")            doc.close()            self.log_message(f"所有章节已导出到: {output_dir}", "SUCCESS")        except Exception as e:            self.log_message(f"导出失败: {str(e)}", "ERROR")        finally:            self.root.after(0, self.progress.stop)    def run(self):        """运行应用程序"""        self.root.mainloop()def main():    """主函数"""    app = ImprovedPDFChapterSplitter()    app.run()if __name__ == "__main__":    main()