Create pymupdf_test_WIP_2.py

pymupdf_test_WIP_2.py

@@ -0,0 +1,318 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import os
+import sys
+import json
+import math
+import logging
+import queue
+import argparse
+import threading
+import multiprocessing as mp
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import fitz
+import tkinter as tk
+from tkinter import filedialog, messagebox
+from joblib import Parallel, cpu_count, delayed
+
+logging.basicConfig(format="%(asctime)s [%(levelname)s] %(message)s", level=logging.INFO)
+_LOGGER = logging.getLogger(__name__)
+
+@dataclass
+class Config:
+    parallel_threshold: int = 16
+    line_tolerance: float = 1.0
+    min_rect_area: float = 1e4
+    cores: int = max(1, cpu_count() - 2)
+
+_CFG = Config()
+
+def _cluster(vals, tol):
+    if not vals:
+        return [], {}
+    arr = np.sort(np.unique(np.array(vals, dtype=float)))
+    if arr.size == 1:
+        return [arr], {round(float(arr[0]), 2): 0}
+    breaks = np.where(np.diff(arr) > tol)[0] + 1
+    clusters = np.split(arr, breaks)
+    mapping = {round(float(v), 2): idx for idx, cl in enumerate(clusters) for v in cl}
+    return clusters, mapping
+
+def _detect_table_bboxes(page):
+    horiz, vert = [], []
+    for d in page.get_drawings():
+        if d["type"] != 1 or d["linewidth"] > 2:
+            continue
+        x0, y0, x1, y1 = d["bbox"]
+        if abs(y1 - y0) < 2:
+            horiz.append(((y0 + y1) / 2, x0, x1))
+        elif abs(x1 - x0) < 2:
+            vert.append(((x0 + x1) / 2, y0, y1))
+    if not horiz or not vert:
+        return []
+    _, row_map = _cluster([y for y, _, _ in horiz], _CFG.line_tolerance)
+    _, col_map = _cluster([x for x, _, _ in vert], _CFG.line_tolerance)
+    rows = {}
+    for y, x0, x1 in horiz:
+        rows.setdefault(row_map[round(y, 2)], []).append(y)
+    cols = {}
+    for x, y0, y1 in vert:
+        cols.setdefault(col_map[round(x, 2)], []).append(x)
+    row_pos = sorted(sum(v) / len(v) for v in rows.values())
+    col_pos = sorted(sum(v) / len(v) for v in cols.values())
+    rects = [
+        fitz.Rect(c0, r0, c1, r1)
+        for r0, r1 in zip(row_pos, row_pos[1:])
+        for c0, c1 in zip(col_pos, col_pos[1:])
+        if fitz.Rect(c0, r0, c1, r1).get_area() >= _CFG.min_rect_area
+    ]
+    uniq = []
+    for r in rects:
+        if not any(u.contains(r) or r.contains(u) for u in uniq):
+            uniq.append(r)
+    return uniq
+
+def _clean(text):
+    return " ".join(text.replace("-\n", "").replace("\n", " ").split()) if isinstance(text, str) else ""
+
+def _extract_table(words, rect):
+    inside = [w for w in words if rect.x0 <= w[0] <= rect.x1 and rect.y0 <= w[1] <= rect.y1]
+    if not inside:
+        return []
+    inside.sort(key=lambda w: (w[1], w[0]))
+    x_centers = [(w[0] + w[2]) / 2 for w in inside]
+    _, col_map = _cluster(x_centers, 10)
+    columns = sorted(set(col_map.values()))
+    rows, cur_y, cur = [], None, []
+    for w in inside:
+        if cur_y is None or abs(w[1] - cur_y) > 5:
+            if cur:
+                rows.append(cur)
+            cur = [w]
+            cur_y = w[1]
+        else:
+            cur.append(w)
+    if cur:
+        rows.append(cur)
+    def _row_to_cells(r):
+        cells = [""] * len(columns)
+        for w in r:
+            col = col_map[round((w[0] + w[2]) / 2, 2)]
+            cells[col] += (" " if cells[col] else "") + w[4]
+        return [_clean(c) for c in cells]
+    table_matrix = [_row_to_cells(r) for r in rows]
+    if len(table_matrix) < 2:
+        return []
+    headers = table_matrix[0]
+    return [dict(zip(headers, row)) for row in table_matrix[1:]]
+
+def _process_page(page, page_number):
+    words = page.get_text("words")
+    tables = _detect_table_bboxes(page)
+    tbl_json = [json.dumps(_extract_table(words, r), ensure_ascii=False, indent=1) for r in tables if _extract_table(words, r)]
+    tbl_boxes = [(r.x0, r.y0, r.x1, r.y1) for r in tables]
+    outside = [
+        w for w in words
+        if not any(bx0 <= w[0] <= bx2 and by0 <= w[1] <= by3 for bx0, by0, bx2, by3 in tbl_boxes)
+    ]
+    outside.sort(key=lambda w: (w[1], w[0]))
+    text_lines, cur_y, cur = [], None, []
+    for w in outside:
+        if cur_y is None or abs(w[1] - cur_y) > 10:
+            if cur:
+                text_lines.append(" ".join(cur))
+            cur, cur_y = [w[4]], w[1]
+        else:
+            cur.append(w[4])
+    if cur:
+        text_lines.append(" ".join(cur))
+    result = f"Page {page_number + 1}\n" + "\n".join(text_lines) + "\n"
+    for idx, tbl in enumerate(tbl_json, 1):
+        result += f'"table {idx}":\n{tbl}\n'
+    return result
+
+def _write_output(base_path, content):
+    out_path = Path(base_path).with_suffix(".txt")
+    out_path.write_text(content, encoding="utf-8")
+    _LOGGER.info("Wrote %s", out_path)
+
+def _process_document(path):
+    try:
+        with fitz.open(path) as doc:
+            results = []
+            for page_number in range(doc.page_count):
+                results.append(_process_page(doc.load_page(page_number), page_number))
+        _write_output(path, "\n".join(results))
+    except fitz.FileDataError as e:
+        _LOGGER.error("Cannot open %s: %s", path, e)
+    except Exception as e:
+        _LOGGER.error("Error processing %s: %s", path, e)
+
+def _run_cli(files):
+    pdfs = []
+    for p in files:
+        try:
+            _ = fitz.open(p)
+            pdfs.append(p)
+        except fitz.FileDataError:
+            _LOGGER.error("Password-protected PDF skipped: %s", p)
+        except Exception as e:
+            _LOGGER.error("Error opening %s: %s", p, e)
+    if not pdfs:
+        return
+    if len(pdfs) == 1 or _CFG.cores == 1:
+        for p in pdfs:
+            _process_document(p)
+    else:
+        Parallel(n_jobs=min(_CFG.cores, len(pdfs)))(delayed(_process_document)(p) for p in pdfs)
+
+class GUI:
+    def __init__(self):
+        self.root = tk.Tk()
+        self.root.title("Parser-Sevenof9 — PyMuPDF")
+        self.queue = queue.Queue()
+        self.stop_event = threading.Event()
+        self.files = []
+        self._build_widgets()
+        self.thread = None
+        self.root.protocol("WM_DELETE_WINDOW", self._on_close)
+        self._poll()
+        self.root.mainloop()
+
+    def _build_widgets(self):
+        tk.Label(self.root, text="Selected PDF files:").pack(pady=5)
+        list_frame = tk.Frame(self.root)
+        list_frame.pack(pady=5)
+        sb = tk.Scrollbar(list_frame)
+        self.listbox = tk.Listbox(list_frame, selectmode=tk.MULTIPLE, width=80, height=6, yscrollcommand=sb.set)
+        sb.config(command=self.listbox.yview)
+        self.listbox.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
+        sb.pack(side=tk.RIGHT, fill=tk.Y)
+        btn_frame = tk.Frame(self.root)
+        btn_frame.pack(pady=10)
+        tk.Button(btn_frame, text="Add Folder", command=self._add_folder).pack(side=tk.LEFT, padx=5)
+        tk.Button(btn_frame, text="Select Files", command=self._add_files).pack(side=tk.LEFT, padx=5)
+        tk.Button(btn_frame, text="Remove Selected", command=self._remove_sel).pack(side=tk.LEFT, padx=5)
+        tk.Button(btn_frame, text="Remove All", command=self._remove_all).pack(side=tk.LEFT, padx=5)
+        tk.Button(self.root, text="Start Parser", command=self._start_parser).pack(pady=5)
+        tk.Button(self.root, text="Stop", command=self._stop_parser).pack(pady=5)
+        prog_frame = tk.Frame(self.root)
+        prog_frame.pack(padx=10, pady=5)
+        sbp = tk.Scrollbar(prog_frame)
+        self.prog = tk.Text(prog_frame, height=15, width=100, state=tk.DISABLED, yscrollcommand=sbp.set)
+        sbp.config(command=self.prog.yview)
+        self.prog.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
+        sbp.pack(side=tk.RIGHT, fill=tk.Y)
+
+    def _add_folder(self):
+        folder = filedialog.askdirectory(title="Select Folder")
+        if not folder:
+            return
+        for root, _, fs in os.walk(folder):
+            for f in fs:
+                if f.lower().endswith(".pdf"):
+                    p = os.path.join(root, f)
+                    if p not in self.files:
+                        self.files.append(p)
+                        self.listbox.insert(tk.END, p)
+
+    def _add_files(self):
+        for p in filedialog.askopenfilenames(title="Select PDF Files", filetypes=[("PDF Files", "*.pdf")]):
+            if p not in self.files:
+                self.files.append(p)
+                self.listbox.insert(tk.END, p)
+
+    def _remove_sel(self):
+        sel = self.listbox.curselection()
+        for idx in reversed(sel):
+            self.listbox.delete(idx)
+            del self.files[idx]
+
+    def _remove_all(self):
+        self.listbox.delete(0, tk.END)
+        self.files.clear()
+
+    def _start_parser(self):
+        if not self.files:
+            messagebox.showinfo("No Files", "Please select at least one file.")
+            return
+        if self.thread and self.thread.is_alive():
+            messagebox.showinfo("Running", "Parser is already running.")
+            return
+        self.stop_event.clear()
+        self.thread = threading.Thread(target=self._worker, daemon=True)
+        self.thread.start()
+
+    def _stop_parser(self):
+        if self.thread and self.thread.is_alive():
+            self.stop_event.set()
+
+    def _worker(self):
+        for f in self.files:
+            if self.stop_event.is_set():
+                self.queue.put("[INFO] Parser stopped.\n")
+                return
+            self.queue.put(f"[INFO] Processing {Path(f).name}\n")
+            _process_document(f)
+        self.queue.put("[INFO] Parser finished.\n")
+
+    def _poll(self):
+        try:
+            while True:
+                msg = self.queue.get_nowait()
+                self._append(msg)
+        except queue.Empty:
+            pass
+        self.root.after(100, self._poll)
+
+    def _append(self, txt):
+        self.prog.config(state=tk.NORMAL)
+        self.prog.insert(tk.END, txt)
+        self.prog.see(tk.END)
+        self.prog.config(state=tk.DISABLED)
+
+    def _on_close(self):
+        self._stop_parser()
+        self.root.destroy()
+
+def _positive_int(v):
+    iv = int(v)
+    if iv <= 0:
+        raise argparse.ArgumentTypeError("value must be positive")
+    return iv
+
+def _positive_float(v):
+    fv = float(v)
+    if fv <= 0:
+        raise argparse.ArgumentTypeError("value must be positive")
+    return fv
+
+def _parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("pdfs", nargs="*", help="PDF files")
+    p.add_argument("--gui", action="store_true", help="Launch GUI")
+    p.add_argument("--threshold", type=_positive_int, default=_CFG.parallel_threshold)
+    p.add_argument("--tol", type=_positive_float, default=_CFG.line_tolerance)
+    p.add_argument("--cores", type=_positive_int, default=_CFG.cores)
+    p.add_argument("--log", default="INFO", help="Logging level")
+    return p.parse_args()
+
+def main():
+    args = _parse_args()
+    _CFG.parallel_threshold = args.threshold
+    _CFG.line_tolerance = args.tol
+    _CFG.cores = args.cores
+    _LOGGER.setLevel(args.log.upper())
+    if args.gui or (not args.pdfs and not args.gui):
+        GUI()
+    else:
+        _run_cli(args.pdfs)
+
+if __name__ == "__main__":
+    mp.set_start_method("spawn", force=True)
+    mp.freeze_support()
+    main()