obsidian-plugin/lllin000_PaperForge
1
0
Fork 0
mirror of https://github.com/lllin000/PaperForge.git synced 2026-07-22 06:50:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat: add paperforge ocr redo command for re-running OCR on marked papers

Browse source
This commit is contained in:
Research Assistant 2026-06-01 11:29:06 +08:00
parent 6acfa8b5c2
commit eeb4e8f85d
5 changed files with 380 additions and 49 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
paperforge/cli.py
View file

@ -231,6 +231,8 @@ def build_parser() -> argparse.ArgumentParser:
ocr_sub.add_parser("run", help="Run OCR queue")
doctor_parser = ocr_sub.add_parser("doctor", help="Diagnose OCR configuration and connectivity")
doctor_parser.add_argument("--live", action="store_true", help="Run live PDF test (L4)")
redo_parser = ocr_sub.add_parser("redo", help="Re-run OCR for papers marked ocr_redo: true")
redo_parser.set_defaults(ocr_action="redo")
# context (Phase 26: traceable AI context packs)
p_context = sub.add_parser("context", help="Generate traceable AI context pack for paper(s)")

62
paperforge/commands/ocr.py
View file

@ -2,6 +2,8 @@
import argparse
import logging
import re as _re
import shutil
from pathlib import Path
from paperforge import __version__
@ -126,6 +128,57 @@ def _get_run_ocr():
return run_ocr
def _run_ocr_redo(vault: Path, verbose: bool = False, no_progress: bool = False) -> int:
"""Scan for papers with ocr_redo: true, reset their OCR state, then run OCR queue."""
from paperforge.adapters.obsidian_frontmatter import extract_preserved_ocr_redo
from paperforge.worker._utils import pipeline_paths
from paperforge.worker.ocr import run_ocr
paths = pipeline_paths(vault)
ocr_root = paths.get("ocr")
lit_root = paths.get("literature")
if not lit_root or not lit_root.exists():
logger.info("No literature directory found, nothing to redo")
return 0
redo_entries = []
for note_file in sorted(lit_root.rglob("*.md")):
if note_file.name in ("fulltext.md", "deep-reading.md", "discussion.md"):
continue
try:
text = note_file.read_text(encoding="utf-8")
except Exception:
continue
if not extract_preserved_ocr_redo(text):
continue
key_match = _re.search(r"^zotero_key:\s*(.+)$", text, _re.MULTILINE)
if not key_match:
continue
zotero_key = key_match.group(1).strip().strip('"').strip("'")
redo_entries.append((zotero_key, note_file, text))
if not redo_entries:
logger.info("No papers with ocr_redo: true found")
return 0
logger.info("Found %d paper(s) with ocr_redo: true", len(redo_entries))
for zotero_key, note_file, text in redo_entries:
# Delete OCR output directory
ocr_dir = ocr_root / zotero_key if ocr_root else None
if ocr_dir and ocr_dir.exists():
shutil.rmtree(ocr_dir)
logger.info("Deleted OCR directory for %s", zotero_key)
# Update library note frontmatter: ocr_status -> pending, ocr_redo -> false
text = _re.sub(r"^ocr_status:\s*.+$", "ocr_status: pending", text, flags=_re.MULTILINE)
text = _re.sub(r"^ocr_redo:\s*.+$", "ocr_redo: false", text, flags=_re.MULTILINE)
note_file.write_text(text, encoding="utf-8")
logger.info("Reset ocr_status to pending and ocr_redo to false for %s", zotero_key)
return run_ocr(vault, verbose=verbose, no_progress=no_progress)
def run(args: argparse.Namespace) -> int:
"""Run OCR command.
@ -150,6 +203,15 @@ def run(args: argparse.Namespace) -> int:
if ocr_action == "doctor" or diagnose_only:
return _diagnose(vault, live=live, json_output=json_output)
if ocr_action == "redo":
logger.info("OCR redo: scanning for ocr_redo: true papers...")
rc = _run_ocr_redo(
vault,
verbose=getattr(args, "verbose", False),
no_progress=getattr(args, "no_progress", False),
)
return rc
if key:
logger.info("Processing specific key: %s", key)

185
paperforge/worker/ocr.py
View file

@ -532,6 +532,26 @@ def parse_reference_number(text: str) -> int | None:
return None
def is_reference_lead_block(text: str) -> bool:
text = clean_block_text(text)
if not text:
return False
if parse_reference_number(text) is not None:
return True
if re.match(r"^\s*[^\n]{1,160}\(\d{4}[a-z]?\)\.", text):
return True
if re.match(r"^\s*[A-Z][A-Za-z'’\-]+(?:\s+[A-Z][A-Za-z'’\-]+)*(?:,\s|\s+et al\.)", text):
return True
return False
def _bbox_column_bucket_xc(block: dict) -> float | None:
bbox = block.get("block_bbox", [0, 0, 0, 0])
if bbox[2] <= bbox[0]:
return None
return (bbox[0] + bbox[2]) / 2
def sort_reference_blocks(blocks: list[dict]) -> list[dict]:
numbered_count = sum(parse_reference_number(block.get("block_content", "")) is not None for block in blocks)
if numbered_count >= max(3, len(blocks) // 2):
@ -543,9 +563,23 @@ def sort_reference_blocks(blocks: list[dict]) -> list[dict]:
block_sort_key(block),
),
)
return sorted(
blocks, key=lambda block: (block.get("block_bbox", [0, 0, 0, 0])[0], block.get("block_bbox", [0, 0, 0, 0])[1])
)
xcs = [xc for block in blocks if (xc := _bbox_column_bucket_xc(block)) is not None]
if not xcs:
mid = 600
else:
mid = sum(xcs) / len(xcs)
left: list[dict] = []
right: list[dict] = []
for block in blocks:
xc = _bbox_column_bucket_xc(block)
if xc is None or xc < mid - 60:
left.append(block)
else:
right.append(block)
y_sort = lambda block: int(block.get("block_bbox", [0, 0, 0, 0])[1])
left.sort(key=y_sort)
right.sort(key=y_sort)
return left + right
def assign_reference_continuation(continuation_block: dict, reference_blocks: list[dict]) -> int | None:
@ -566,6 +600,53 @@ def assign_reference_continuation(continuation_block: dict, reference_blocks: li
return candidates[0][0]
def build_reference_section_lines(reference_blocks: list[dict], reference_continuations: list[dict]) -> list[str]:
if not reference_blocks:
return []
lines = [""]
ordered_reference_blocks = sort_reference_blocks(reference_blocks)
continuation_map: dict[int, list[str]] = {}
sorted_continuations = sorted(reference_continuations, key=block_sort_key)
assigned_continuation_indexes: set[int] = set()
incomplete_reference_indexes = [
index
for index, block in enumerate(ordered_reference_blocks)
if clean_block_text(block.get("block_content", "")).rstrip().endswith(":")
]
for index, continuation in zip(incomplete_reference_indexes, sorted_continuations, strict=False):
continuation_text = clean_block_text(continuation.get("block_content", ""))
if continuation_text:
continuation_map.setdefault(index, []).append(continuation_text)
assigned_continuation_indexes.add(id(continuation))
for continuation in sorted_continuations:
if id(continuation) in assigned_continuation_indexes:
continue
target_index = assign_reference_continuation(continuation, ordered_reference_blocks)
if target_index is None:
continue
continuation_map.setdefault(target_index, []).append(clean_block_text(continuation.get("block_content", "")))
assigned_continuation_indexes.add(id(continuation))
unassigned_continuations = [
clean_block_text(continuation.get("block_content", ""))
for continuation in sorted_continuations
if id(continuation) not in assigned_continuation_indexes
]
for index, block in enumerate(ordered_reference_blocks):
text = clean_block_text(block.get("block_content", ""))
if text:
lines.append(text)
for continuation_text in continuation_map.get(index, []):
if continuation_text:
lines.append(continuation_text)
if text.rstrip().endswith(":") and unassigned_continuations:
lines.append(unassigned_continuations.pop(0))
return lines
def append_reference_section(rendered: list[str], reference_blocks: list[dict], reference_continuations: list[dict]) -> None:
rendered.extend(build_reference_section_lines(reference_blocks, reference_continuations))
def clean_author_line(text: str) -> str:
text = clean_block_text(text)
text = text.replace("$^{ID}$", "")
@ -1324,6 +1405,10 @@ def render_page_blocks(
reference_continuations: list[dict] = []
footnotes: list[str] = []
footnote_counter = 0
references_heading_seen = False
references_section_active = False
references_emitted = False
references_insert_index: int | None = None
rendered_cluster_ids: set[int] = set()
rendered_caption_media_ids: set[int] = set()
first_page_meta_done = page_index != 1
@ -1367,7 +1452,16 @@ def render_page_blocks(
rendered.extend(deferred_meta)
deferred_meta.clear()
first_page_meta_done = True
if references_heading_seen and (not references_emitted) and title.lower() != "references":
if reference_blocks:
append_reference_section(rendered, reference_blocks, reference_continuations)
references_emitted = True
references_section_active = False
rendered.append(f"### {title}")
if title.lower() == "references":
references_heading_seen = True
references_section_active = True
references_insert_index = len(rendered)
continue
if label == "abstract":
if page_index == 1 and affiliation_buffer:
@ -1382,10 +1476,10 @@ def render_page_blocks(
if label == "reference_content":
text = clean_block_text(content)
if text and (not is_reference_tail_noise_line(text)):
if parse_reference_number(text) is None:
reference_continuations.append(block)
else:
if is_reference_lead_block(text):
reference_blocks.append(block)
else:
reference_continuations.append(block)
continue
if label == "text":
text = clean_block_text(content)
@ -1397,7 +1491,7 @@ def render_page_blocks(
or is_embedded_figure_text_block(block, blocks, page_width=ocr_width, page_height=ocr_height)
):
continue
if raw_reference_blocks and bbox[1] >= first_reference_y - 10:
if references_section_active and raw_reference_blocks and bbox[1] >= first_reference_y - 10:
reference_continuations.append(block)
continue
if is_numbered_figure_caption(text):
@ -1525,46 +1619,12 @@ def render_page_blocks(
rendered.append("")
rendered.extend(footnotes)
rendered = dedupe_page_media_lines(rendered)
if reference_blocks:
rendered.append("")
ordered_reference_blocks = sort_reference_blocks(reference_blocks)
continuation_map: dict[int, list[str]] = {}
sorted_continuations = sorted(reference_continuations, key=block_sort_key)
assigned_continuation_indexes: set[int] = set()
incomplete_reference_indexes = [
index
for index, block in enumerate(ordered_reference_blocks)
if clean_block_text(block.get("block_content", "")).rstrip().endswith(":")
]
for index, continuation in zip(incomplete_reference_indexes, sorted_continuations, strict=False):
continuation_text = clean_block_text(continuation.get("block_content", ""))
if continuation_text:
continuation_map.setdefault(index, []).append(continuation_text)
assigned_continuation_indexes.add(id(continuation))
for continuation in sorted_continuations:
if id(continuation) in assigned_continuation_indexes:
continue
target_index = assign_reference_continuation(continuation, ordered_reference_blocks)
if target_index is None:
continue
continuation_map.setdefault(target_index, []).append(
clean_block_text(continuation.get("block_content", ""))
)
assigned_continuation_indexes.add(id(continuation))
unassigned_continuations = [
clean_block_text(continuation.get("block_content", ""))
for continuation in sorted_continuations
if id(continuation) not in assigned_continuation_indexes
]
for index, block in enumerate(ordered_reference_blocks):
text = clean_block_text(block.get("block_content", ""))
if text:
rendered.append(text)
for continuation_text in continuation_map.get(index, []):
if continuation_text:
rendered.append(continuation_text)
if text.rstrip().endswith(":") and unassigned_continuations:
rendered.append(unassigned_continuations.pop(0))
if reference_blocks and not references_emitted:
reference_lines = build_reference_section_lines(reference_blocks, reference_continuations)
if references_insert_index is not None:
rendered[references_insert_index:references_insert_index] = reference_lines
else:
rendered.extend(reference_lines)
return [part for part in rendered if part]
@ -1607,7 +1667,7 @@ def postprocess_ocr_result(vault: Path, key: str, all_results: list[dict]) -> tu
return (page_num, markdown_path, json_path, fulltext_md_path)
def run_ocr(vault: Path, verbose: bool = False, no_progress: bool = False) -> int:
def run_ocr(vault: Path, verbose: bool = False, no_progress: bool = False, redo_mode: bool = False) -> int:
from paperforge.pdf_resolver import resolve_pdf_path
paths = pipeline_paths(vault)
@ -1644,6 +1704,35 @@ def run_ocr(vault: Path, verbose: bool = False, no_progress: bool = False) -> in
control_actions = load_control_actions(paths)
target_keys = {key for key, action in control_actions.items() if action.get("do_ocr", False)}
if redo_mode:
redo_keys = [key for key, action in control_actions.items() if action.get("ocr_redo", False)]
if redo_keys:
logger.info("OCR redo mode: processing %d paper(s) with ocr_redo: true", len(redo_keys))
ocr_root = paths.get("ocr")
lit_root = paths.get("literature")
for key in redo_keys:
ocr_dir = ocr_root / key if ocr_root else None
if ocr_dir and ocr_dir.exists():
shutil.rmtree(ocr_dir)
logger.info("Deleted OCR directory for %s", key)
target_keys.add(key)
if lit_root and lit_root.exists():
for note_file in lit_root.rglob("*.md"):
if note_file.name in ("fulltext.md", "deep-reading.md", "discussion.md"):
continue
try:
text = note_file.read_text(encoding="utf-8")
except Exception:
continue
for key in redo_keys:
if key in text:
text = re.sub(r"^ocr_status:\s*.+$", "ocr_status: pending", text, flags=re.MULTILINE)
text = re.sub(r"^ocr_redo:\s*.+$", "ocr_redo: false", text, flags=re.MULTILINE)
note_file.write_text(text, encoding="utf-8")
logger.info("Reset frontmatter for %s", key)
break
target_rows = []
for export_path in sorted(paths["exports"].glob("*.json")):
for item in load_export_rows(export_path):

3
paperforge/worker/sync.py
View file

@ -198,7 +198,8 @@ def load_control_actions(paths: dict[str, Path]) -> dict[str, dict]:
)
if analyze_match:
analyze = analyze_match.group(1).lower() == "true"
actions[zotero_key] = {"analyze": analyze, "do_ocr": do_ocr}
ocr_redo = bool(extract_preserved_ocr_redo(text))
actions[zotero_key] = {"analyze": analyze, "do_ocr": do_ocr, "ocr_redo": ocr_redo}
return actions

177
tests/test_ocr_redo.py Normal file
View file

@ -0,0 +1,177 @@
"""Tests for paperforge ocr redo workflow."""
import json
import re
import shutil
from pathlib import Path
import pytest
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_vault(tmp_path: Path) -> tuple[Path, Path, Path, Path]:
vault = tmp_path / "vault"
vault.mkdir()
(vault / "paperforge.json").write_text("{}", encoding="utf-8")
ocr_root = vault / "PaperForge" / "ocr"
ocr_root.mkdir(parents=True)
exports = vault / "PaperForge" / "exports"
exports.mkdir(parents=True)
literature = vault / "Resources" / "Literature"
literature.mkdir(parents=True)
return vault, ocr_root, exports, literature
def _make_ocr_meta(ocr_root: Path, key: str, status: str = "done") -> dict:
meta_dir = ocr_root / key
meta_dir.mkdir(parents=True, exist_ok=True)
(meta_dir / "images").mkdir(exist_ok=True)
meta = {
"zotero_key": key,
"ocr_status": status,
"ocr_provider": "PaddleOCR-VL-1.6",
"source_pdf": f"some/path/{key}.pdf",
"ocr_job_id": "job-123",
"ocr_started_at": "2025-01-01T00:00:00",
"ocr_finished_at": "2025-01-01T01:00:00",
"page_count": 5,
"markdown_path": f"PaperForge/ocr/{key}/fulltext.md",
}
(meta_dir / "meta.json").write_text(json.dumps(meta, indent=2), encoding="utf-8")
return meta
def _make_library_note(lit_dir: Path, key: str, ocr_redo: bool = True, ocr_status: str = "done") -> Path:
domain_dir = lit_dir / "test_domain"
domain_dir.mkdir(parents=True, exist_ok=True)
note_path = domain_dir / f"{key}.md"
note_text = f"""---
title: "Test Paper"
zotero_key: {key}
do_ocr: true
analyze: true
ocr_status: {ocr_status}
ocr_redo: {"true" if ocr_redo else "false"}
tags:
- test
---
# Test Paper
Some content
"""
note_path.write_text(note_text, encoding="utf-8")
return note_path
# ---------------------------------------------------------------------------
# Test: OCR redo resets ocr_status in meta.json
# ---------------------------------------------------------------------------
def test_ocr_redo_resets_ocr_status():
"""Setting ocr_status to pending in meta.json simulates redo reset."""
meta = {"zotero_key": "KEY001", "ocr_status": "done", "ocr_job_id": "job-xyz"}
meta["ocr_status"] = "pending"
meta["ocr_job_id"] = ""
assert meta["ocr_status"] == "pending"
assert meta["ocr_job_id"] == ""
assert meta["zotero_key"] == "KEY001"
# ---------------------------------------------------------------------------
# Test: OCR redo clears OCR directory
# ---------------------------------------------------------------------------
def test_ocr_redo_clears_ocr_dir(tmp_path):
"""Verify OCR output directory is removed on redo."""
vault, ocr_root, exports, literature = _make_vault(tmp_path)
key = "KEY002"
_make_ocr_meta(ocr_root, key, status="done")
ocr_dir = ocr_root / key
assert ocr_dir.exists()
assert (ocr_dir / "meta.json").exists()
assert (ocr_dir / "images").exists()
# Simulate redo: delete the OCR directory
shutil.rmtree(ocr_dir)
assert not ocr_dir.exists()
# ---------------------------------------------------------------------------
# Test: OCR redo updates library note frontmatter
# ---------------------------------------------------------------------------
def test_ocr_redo_updates_frontmatter(tmp_path):
"""Verify ocr_status and ocr_redo are updated in the library note."""
vault, ocr_root, exports, literature = _make_vault(tmp_path)
key = "KEY003"
_make_ocr_meta(ocr_root, key, status="done")
note_path = _make_library_note(literature, key, ocr_redo=True, ocr_status="done")
# Simulate redo: update frontmatter
text = note_path.read_text(encoding="utf-8")
text = re.sub(r"^ocr_status:\s*.+$", "ocr_status: pending", text, flags=re.MULTILINE)
text = re.sub(r"^ocr_redo:\s*.+$", "ocr_redo: false", text, flags=re.MULTILINE)
note_path.write_text(text, encoding="utf-8")
updated = note_path.read_text(encoding="utf-8")
assert "ocr_status: pending" in updated
assert "ocr_redo: false" in updated
# ocr_redo: false and ocr_redo: true must not both appear
assert len(re.findall(r"^ocr_redo:", updated, re.MULTILINE)) == 1
assert len(re.findall(r"^ocr_status:", updated, re.MULTILINE)) == 1
# ---------------------------------------------------------------------------
# Test: redo subcommand is registered
# ---------------------------------------------------------------------------
def test_redo_subcommand_registered():
"""Verify 'paperforge ocr redo' is a registered subcommand."""
from paperforge.cli import build_parser
parser = build_parser()
# Parse ocr redo --help should not fail
with pytest.raises(SystemExit) as exc:
parser.parse_args(["ocr", "redo", "--help"])
assert exc.value.code == 0
# Parse ocr redo should set ocr_action="redo"
args = parser.parse_args(["ocr", "redo"])
assert args.command == "ocr"
assert args.ocr_action == "redo"
# ---------------------------------------------------------------------------
# Test: _run_ocr_redo scan logic
# ---------------------------------------------------------------------------
def test_ocr_redo_scan_finds_marked_papers(tmp_path):
"""Verify scan finds papers with ocr_redo: true."""
vault, ocr_root, exports, literature = _make_vault(tmp_path)
_make_library_note(literature, "KEY_A", ocr_redo=True)
_make_library_note(literature, "KEY_B", ocr_redo=False)
_make_library_note(literature, "KEY_C", ocr_redo=True)
from paperforge.adapters.obsidian_frontmatter import extract_preserved_ocr_redo
found = []
for note_file in sorted(literature.rglob("*.md")):
if note_file.name in ("fulltext.md", "deep-reading.md", "discussion.md"):
continue
text = note_file.read_text(encoding="utf-8")
if not extract_preserved_ocr_redo(text):
continue
key_match = re.search(r"^zotero_key:\s*(.+)$", text, re.MULTILINE)
assert key_match is not None
zkey = key_match.group(1).strip().strip('"').strip("'")
found.append(zkey)
assert "KEY_A" in found
assert "KEY_B" not in found
assert "KEY_C" in found
assert len(found) == 2