Aman
a32b7508c7
PaddleOCR with preprocessing, scispaCy NER, figure/table detection, citation extraction, and chunked Markdown output with frontmatter. Includes watch mode and notebook reprocessing.
144 lines
4.9 KiB
Python
144 lines
4.9 KiB
Python
from __future__ import annotations
|
|
|
|
import datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import yaml
|
|
|
|
from ocr_pipeline.config import settings
|
|
from ocr_pipeline.utils.logging import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class MarkdownWriter:
|
|
def __init__(self):
|
|
self.base_dir = Path(settings.output.base_directory).expanduser().resolve()
|
|
self.organize_by = settings.output.organize_by
|
|
self.frontmatter_fields = settings.output.frontmatter.fields
|
|
|
|
def _get_output_dir(self, timestamp: str) -> Path:
|
|
dt = datetime.datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
|
|
|
|
if self.organize_by == "date_run":
|
|
date_dir = self.base_dir / dt.strftime("%Y-%m-%d")
|
|
run_dirs = sorted(date_dir.glob("run_*"))
|
|
run_num = len(run_dirs) + 1
|
|
return date_dir / f"run_{run_num:03d}"
|
|
elif self.organize_by == "source_dir":
|
|
return self.base_dir / "by_source"
|
|
else:
|
|
return self.base_dir
|
|
|
|
def write_chunk(
|
|
self,
|
|
content: str,
|
|
metadata: dict[str, Any],
|
|
figures: list[Any] = None,
|
|
tables: list[Any] = None,
|
|
entities: Any = None,
|
|
citations: list[Any] = None,
|
|
) -> Path:
|
|
output_dir = self._get_output_dir(metadata.get("timestamp", ""))
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
source_path = Path(metadata.get("source_path", "unknown"))
|
|
safe_name = "".join(c for c in source_path.stem if c.isalnum() or c in "-_")[:80]
|
|
chunk_idx = metadata.get("chunk_index", 0)
|
|
total_chunks = metadata.get("total_chunks", 1)
|
|
|
|
if total_chunks > 1:
|
|
filename = f"{safe_name}_chunk_{chunk_idx:03d}.md"
|
|
else:
|
|
filename = f"{safe_name}.md"
|
|
|
|
output_path = output_dir / filename
|
|
|
|
frontmatter = self._build_frontmatter(metadata)
|
|
body = self._build_body(content, figures, tables, entities, citations)
|
|
|
|
with open(output_path, "w") as f:
|
|
f.write("---\n")
|
|
yaml.dump(frontmatter, f, sort_keys=False, allow_unicode=True)
|
|
f.write("---\n\n")
|
|
f.write(body)
|
|
|
|
logger.debug("markdown_written", path=str(output_path))
|
|
return output_path
|
|
|
|
def _build_frontmatter(self, metadata: dict[str, Any]) -> dict[str, Any]:
|
|
fm = {}
|
|
for field in self.frontmatter_fields:
|
|
if field in metadata:
|
|
fm[field] = metadata[field]
|
|
return fm
|
|
|
|
def _build_body(
|
|
self,
|
|
content: str,
|
|
figures: list[Any] = None,
|
|
tables: list[Any] = None,
|
|
entities: Any = None,
|
|
citations: list[Any] = None,
|
|
) -> str:
|
|
parts = []
|
|
|
|
source_name = Path(content.split("\n")[0]).name if content else "Unknown"
|
|
parts.append(f"# Screenshot: {source_name}")
|
|
parts.append("")
|
|
|
|
if content:
|
|
parts.append("## Extracted Text")
|
|
parts.append("")
|
|
parts.append(content.strip())
|
|
parts.append("")
|
|
|
|
if figures:
|
|
parts.append("## Figures Detected")
|
|
parts.append("")
|
|
for i, fig in enumerate(figures):
|
|
parts.append(f"### Figure {i + 1}")
|
|
parts.append(f"- **Bounding Box**: {fig.bbox}")
|
|
parts.append(f"- **Confidence**: {fig.confidence:.2f}")
|
|
if fig.caption:
|
|
parts.append(f"- **Caption**: {fig.caption}")
|
|
parts.append("")
|
|
|
|
if tables:
|
|
parts.append("## Tables Detected")
|
|
parts.append("")
|
|
for i, table in enumerate(tables):
|
|
parts.append(f"### Table {i + 1}")
|
|
parts.append(f"- **Bounding Box**: {table.bbox}")
|
|
parts.append(f"- **Confidence**: {table.confidence:.2f}")
|
|
if table.markdown:
|
|
parts.append("")
|
|
parts.append("#### Table Content (Markdown)")
|
|
parts.append("")
|
|
parts.append(table.markdown)
|
|
parts.append("")
|
|
|
|
if entities and entities.entities:
|
|
parts.append("## Detected Entities")
|
|
parts.append("")
|
|
parts.append("| Entity | Type | Context |")
|
|
parts.append("|--------|------|---------|")
|
|
for ent in entities.entities[:50]:
|
|
context = ent.context[:80].replace("|", "\\|") + (
|
|
"..." if len(ent.context) > 80 else ""
|
|
)
|
|
parts.append(f"| {ent.text} | {ent.label} | {context} |")
|
|
parts.append("")
|
|
|
|
if citations:
|
|
parts.append("## Citations Found")
|
|
parts.append("")
|
|
parts.append("| Type | Identifier |")
|
|
parts.append("|------|------------|")
|
|
for cit in citations:
|
|
parts.append(f"| {cit.type.upper()} | `{cit.identifier}` |")
|
|
parts.append("")
|
|
|
|
return "\n".join(parts)
|