97 lines
4.6 KiB
Python
97 lines
4.6 KiB
Python
"""Compare extracted answers against the same Poppler view used for labels."""
|
|
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pymupdf
|
|
from pdf2image import convert_from_path
|
|
|
|
from copienator.commands.splitting_int import (
|
|
ANSWER_TOP_PADDING_POINTS,
|
|
_render_split_outputs,
|
|
)
|
|
|
|
|
|
class SplittingGeometryTests(unittest.TestCase):
|
|
def make_source(self, path, rotation, cropped=True):
|
|
with pymupdf.open() as document:
|
|
page = document.new_page(width=800, height=1000)
|
|
# Asymmetric colours exercise both translation and orientation.
|
|
for row in range(20):
|
|
for column in range(8):
|
|
colour = (row / 20, column / 8, (row + column) % 7 / 7)
|
|
page.draw_rect(
|
|
pymupdf.Rect(column * 100, row * 50,
|
|
(column + 1) * 100, (row + 1) * 50),
|
|
color=None, fill=colour,
|
|
)
|
|
page.insert_text((200, 500), "Middle of the answer")
|
|
if cropped:
|
|
page.set_cropbox(pymupdf.Rect(20, 80, 760, 920))
|
|
page.set_rotation(rotation)
|
|
document.save(path)
|
|
|
|
def assert_matches_preview(self, answer, preview, bounds):
|
|
actual = np.asarray(convert_from_path(answer, dpi=72)[0]).astype(float)
|
|
expected = np.asarray(preview.crop(bounds)).astype(float)
|
|
self.assertEqual(actual.shape, expected.shape)
|
|
self.assertLess(np.abs(actual - expected).mean(), 0.5)
|
|
|
|
def test_cropped_and_uncropped_pages_at_every_rotation(self):
|
|
visible = {
|
|
0: (20, 80, 760, 920),
|
|
90: (80, 20, 920, 760),
|
|
180: (40, 80, 780, 920),
|
|
270: (80, 40, 920, 780),
|
|
}
|
|
for rotation in (0, 90, 180, 270):
|
|
for cropped in (False, True):
|
|
for full_answer in (False, True):
|
|
with self.subTest(rotation=rotation, cropped=cropped, full=full_answer):
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
root = Path(temporary)
|
|
source = root / "copy.pdf"
|
|
self.make_source(source, rotation, cropped)
|
|
before = source.read_bytes()
|
|
preview = convert_from_path(source, dpi=72, use_cropbox=False)[0]
|
|
width, height = preview.size
|
|
bounds = visible[rotation] if cropped else (0, 0, width, height)
|
|
if full_answer:
|
|
coordinates = [("A", 0, 0, 10, 0, 100)]
|
|
else:
|
|
coordinates = [("A", 0, 250, 260, 0, 100),
|
|
("_", 0, 711, 721, 0, 100)]
|
|
answer_top = int(height // 4 - ANSWER_TOP_PADDING_POINTS)
|
|
bounds = (bounds[0], max(bounds[1], answer_top),
|
|
bounds[2], min(bounds[3], height * 3 // 4))
|
|
_render_split_outputs(source, coordinates, root)
|
|
self.assert_matches_preview(root / "A.pdf", preview, bounds)
|
|
self.assertEqual(source.read_bytes(), before)
|
|
|
|
def test_answer_continues_to_bottom_then_next_page(self):
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
root = Path(temporary)
|
|
source = root / "copy.pdf"
|
|
first = root / "first.pdf"
|
|
second = root / "second.pdf"
|
|
self.make_source(first, 180)
|
|
self.make_source(second, 0)
|
|
with pymupdf.open() as document:
|
|
for path in (first, second):
|
|
with pymupdf.open(path) as part:
|
|
document.insert_pdf(part)
|
|
document.save(source)
|
|
previews = convert_from_path(source, dpi=72, use_cropbox=False)
|
|
_render_split_outputs(source, [("A", 0, 700, 720, 0, 100),
|
|
("_", 1, 211, 230, 500, 600)], root)
|
|
rendered = convert_from_path(root / "A.pdf", dpi=72)
|
|
self.assertEqual(len(rendered), 2)
|
|
for actual, preview, bounds in zip(rendered, previews,
|
|
[(40, 688, 780, 920), (20, 80, 760, 250)]):
|
|
expected = np.asarray(preview.crop(bounds)).astype(float)
|
|
actual = np.asarray(actual).astype(float)
|
|
self.assertEqual(actual.shape, expected.shape)
|
|
self.assertLess(np.abs(actual - expected).mean(), 0.5)
|