284 lines
10 KiB
Python
284 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import shutil
|
|
import tempfile
|
|
from collections import defaultdict
|
|
from collections.abc import Sequence
|
|
from pathlib import Path
|
|
|
|
import pymupdf
|
|
from pypdf import PdfReader, PdfWriter
|
|
|
|
from copienator import utils
|
|
from copienator import (
|
|
CliError,
|
|
EvaluationWorkspace,
|
|
ExitCode,
|
|
execute,
|
|
read_json,
|
|
target_parser,
|
|
workspace_from_target,
|
|
)
|
|
from copienator.filesystem import staged_directory
|
|
|
|
SQUARE = 1000 // 38
|
|
Coordinate = tuple[str, int, int, int, int, int]
|
|
ParsedCoordinate = tuple[str, str, int, int, int, int, int]
|
|
|
|
|
|
def decode_json(pdf_file: str | Path) -> tuple[str, list[Coordinate]]:
|
|
"""Read verified label coordinates associated with one copy PDF."""
|
|
pdf_path = Path(pdf_file)
|
|
loaded = read_json(pdf_path.with_suffix(".json"))
|
|
if not isinstance(loaded, dict):
|
|
raise TypeError(f"Expected a JSON object for {pdf_path}")
|
|
boxes = loaded.get("list")
|
|
if not isinstance(boxes, list):
|
|
raise TypeError(f"Expected a list of labels for {pdf_path}")
|
|
page_count = len(PdfReader(pdf_path).pages)
|
|
if page_count == 0:
|
|
raise ValueError(f"PDF contains no pages: {pdf_path}")
|
|
column_width = 1000 // page_count
|
|
result: list[Coordinate] = []
|
|
for entry in boxes:
|
|
if not isinstance(entry, dict):
|
|
raise TypeError(f"Malformed label entry for {pdf_path}: {entry!r}")
|
|
box = entry["box_2d"]
|
|
label = str(entry["label"])
|
|
page_number = ((box[1] + box[3]) // 2) // column_width
|
|
result.append(
|
|
(label, page_number, box[0] - SQUARE, box[2] - SQUARE, box[1], box[3])
|
|
)
|
|
result.sort(key=lambda item: (item[1], item[2]))
|
|
return str(loaded.get("name", "")), result
|
|
|
|
|
|
def _parse_coordinates(coords_list: list[Coordinate]) -> list[ParsedCoordinate]:
|
|
parsed: list[ParsedCoordinate] = []
|
|
for label, page, y0, y1, x0, x1 in coords_list:
|
|
if label.startswith("|"):
|
|
kind, clean_label = "L", label[1:]
|
|
elif label.endswith("|"):
|
|
kind, clean_label = "R", label[:-1]
|
|
else:
|
|
kind, clean_label = "N", label
|
|
parsed.append((clean_label, kind, page, y0, y1, x0, x1))
|
|
filtered: list[ParsedCoordinate] = []
|
|
for item in parsed:
|
|
if not filtered or item[0] != filtered[-1][0]:
|
|
filtered.append(item)
|
|
return filtered
|
|
|
|
|
|
def _save_cropped_page(
|
|
document: pymupdf.Document,
|
|
page_number: int,
|
|
x0: float,
|
|
y0: float,
|
|
x1: float,
|
|
y1: float,
|
|
output_path: Path,
|
|
) -> None:
|
|
# The source has been normalized by _prepare_split_pages: no rotation or
|
|
# CropBox translation remains in the coordinates passed to show_pdf_page.
|
|
visual_crop = pymupdf.Rect(x0, y0, x1, y1)
|
|
cropped = pymupdf.open()
|
|
try:
|
|
target_page = cropped.new_page(width=visual_crop.width, height=visual_crop.height)
|
|
target_page.show_pdf_page(
|
|
target_page.rect,
|
|
document,
|
|
page_number,
|
|
clip=visual_crop,
|
|
)
|
|
cropped.save(output_path)
|
|
finally:
|
|
cropped.close()
|
|
|
|
|
|
def _prepare_split_pages(document: pymupdf.Document) -> list[pymupdf.Rect]:
|
|
"""Use the label preview's full-page coordinates; retain visible bounds.
|
|
|
|
Work only on the in-memory document. Baking rotation into the content after
|
|
restoring the MediaBox avoids show_pdf_page's rotated CropBox offsets.
|
|
Intersecting with the saved bounds later preserves prior margin cropping.
|
|
"""
|
|
visible_bounds = []
|
|
for page in document:
|
|
crop = page.cropbox
|
|
media = page.mediabox
|
|
full_crop = pymupdf.Rect(media.x0, 0, media.x1, media.height)
|
|
page.set_cropbox(full_crop)
|
|
crop -= (full_crop.x0, full_crop.y0, full_crop.x0, full_crop.y0)
|
|
visible_bounds.append(crop * page.rotation_matrix)
|
|
page.remove_rotation()
|
|
return visible_bounds
|
|
|
|
|
|
def _render_split_outputs(
|
|
input_pdf: Path,
|
|
coords_list: list[Coordinate],
|
|
staging: Path,
|
|
) -> set[str]:
|
|
"""Render every current answer into an otherwise empty staging directory."""
|
|
document = pymupdf.open(input_pdf)
|
|
try:
|
|
visible_bounds = _prepare_split_pages(document)
|
|
parsed = _parse_coordinates(coords_list)
|
|
parts_by_label: defaultdict[str, list[Path]] = defaultdict(list)
|
|
with tempfile.TemporaryDirectory(prefix="copienator-split-") as temp_directory:
|
|
temporary = Path(temp_directory)
|
|
for index, item in enumerate(parsed):
|
|
clean_label, kind, start_page, y_start, _y_end, x0_raw, _x1_raw = item
|
|
if clean_label == "_":
|
|
continue
|
|
if not 0 <= start_page < document.page_count:
|
|
raise ValueError(
|
|
f"Invalid page {start_page} for {input_pdf.name}"
|
|
)
|
|
end_page = document.page_count - 1
|
|
end_y = 1000
|
|
for next_item in parsed[index + 1 :]:
|
|
_next_label, next_kind, next_page, next_y, *_rest = next_item
|
|
if (
|
|
(kind == "L" and next_kind in {"L", "N"})
|
|
or (kind == "R" and next_kind in {"R", "N"})
|
|
or kind == "N"
|
|
):
|
|
end_page = next_page
|
|
end_y = min(next_y + int(1.5 * SQUARE), 1000)
|
|
break
|
|
|
|
column_width = 1000 / document.page_count
|
|
if kind == "L":
|
|
fraction_x0 = (x0_raw % column_width) / column_width
|
|
fraction_x1 = 1.0
|
|
end_y = min(1000, end_y + 40)
|
|
elif kind == "R":
|
|
fraction_x0 = 0.0
|
|
left_labels = [entry for entry in parsed if entry[1] == "L"]
|
|
if left_labels:
|
|
closest = min(left_labels, key=lambda entry: abs(entry[3] - y_start))
|
|
center = (closest[5] + closest[6]) / 2.0
|
|
fraction_x1 = (center % column_width) / column_width
|
|
if fraction_x1 <= fraction_x0:
|
|
fraction_x1 = 1.0
|
|
else:
|
|
fraction_x1 = 1.0
|
|
else:
|
|
fraction_x0, fraction_x1 = 0.0, 1.0
|
|
|
|
for page_number in range(start_page, end_page + 1):
|
|
page = document[page_number]
|
|
y0 = (y_start / 1000) * page.rect.height if page_number == start_page else 0
|
|
y1 = (end_y / 1000) * page.rect.height if page_number == end_page else page.rect.height
|
|
clip = pymupdf.Rect(
|
|
fraction_x0 * page.rect.width, y0,
|
|
fraction_x1 * page.rect.width, y1,
|
|
) & visible_bounds[page_number] & page.rect
|
|
if clip.is_empty or clip.height <= 1 or clip.width <= 1:
|
|
continue
|
|
part_path = temporary / f"part-{index}-{page_number}.pdf"
|
|
_save_cropped_page(
|
|
document,
|
|
page_number,
|
|
clip.x0,
|
|
clip.y0,
|
|
clip.x1,
|
|
clip.y1,
|
|
part_path,
|
|
)
|
|
parts_by_label[clean_label].append(part_path)
|
|
|
|
generated: set[str] = set()
|
|
for label, parts in parts_by_label.items():
|
|
filename = f"{label}.pdf"
|
|
merger = PdfWriter()
|
|
try:
|
|
for part in parts:
|
|
merger.append(part)
|
|
merger.write(staging / filename)
|
|
finally:
|
|
merger.close()
|
|
generated.add(filename)
|
|
return generated
|
|
finally:
|
|
document.close()
|
|
|
|
|
|
def _preserve_previous_outputs(
|
|
output_dir: Path,
|
|
staging: Path,
|
|
generated_files: set[str],
|
|
) -> None:
|
|
if not output_dir.is_dir():
|
|
return
|
|
for directory in (path for path in output_dir.iterdir() if path.is_dir()):
|
|
shutil.copytree(directory, staging / directory.name, dirs_exist_ok=True)
|
|
missing_dir = staging / "Missing"
|
|
for item in (path for path in output_dir.iterdir() if path.is_file()):
|
|
if item.name in generated_files:
|
|
continue
|
|
print(f"ALERT: File '{item.name}' not generated. Moving to {missing_dir}")
|
|
missing_dir.mkdir(exist_ok=True)
|
|
shutil.copy2(item, missing_dir / item.name)
|
|
|
|
|
|
def split_an_interro(
|
|
workspace: EvaluationWorkspace,
|
|
input_pdf: Path,
|
|
coords_list: list[Coordinate],
|
|
) -> None:
|
|
"""Regenerate one copy's answers and preserve obsolete ones under Missing."""
|
|
output_dir = workspace.copies_dir / input_pdf.stem
|
|
with staged_directory(output_dir) as staging:
|
|
generated = _render_split_outputs(input_pdf, coords_list, staging)
|
|
_preserve_previous_outputs(output_dir, staging, generated)
|
|
|
|
|
|
def _selected_pdfs(workspace: EvaluationWorkspace, target: Path) -> list[Path]:
|
|
workspace.require_directories("Copies")
|
|
if target.is_file():
|
|
if target.suffix.casefold() != ".pdf":
|
|
raise CliError(f"Target is not a PDF: {target}", ExitCode.INVALID_ARGUMENTS)
|
|
return [target]
|
|
return sorted(workspace.copies_dir.glob("*.pdf"), key=lambda path: path.name.casefold())
|
|
|
|
|
|
def run(workspace: EvaluationWorkspace, target: Path) -> ExitCode:
|
|
workspace.require_files("labels")
|
|
utils.read_all_labels(workspace.root)
|
|
pdf_files = _selected_pdfs(workspace, target)
|
|
status = ExitCode.SUCCESS
|
|
for pdf_path in pdf_files:
|
|
json_path = pdf_path.with_suffix(".json")
|
|
if not json_path.is_file():
|
|
print(f"Warning: No JSON found for {pdf_path.name}")
|
|
status = ExitCode.PARTIAL
|
|
continue
|
|
name, coordinates = decode_json(pdf_path)
|
|
print(f"Decoded name: {name}")
|
|
split_an_interro(workspace, pdf_path, coordinates)
|
|
if not pdf_files:
|
|
print("No PDF copies found.")
|
|
return status
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
return target_parser("Split verified PDF copies into answers by label")
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> int:
|
|
parser = build_parser()
|
|
|
|
def handle(args: argparse.Namespace) -> ExitCode:
|
|
workspace, target = workspace_from_target(args)
|
|
return run(workspace, target)
|
|
|
|
return execute(parser, argv, handle)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|