205 lines
11 KiB
Python
205 lines
11 KiB
Python
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pymupdf
|
|
from PIL import Image
|
|
|
|
from copienator import atomic_write_json, configuration, read_json
|
|
from copienator.annotation_layout import (
|
|
answer_parts, copy_answer_pdf, fingerprint, layout_for, legacy_layout,
|
|
move_answer_pdf, save_parts, added_margin,
|
|
)
|
|
from copienator.commands import annotating, annotating_by_label, splitting_int
|
|
from copienator.commands.annotating_with_checks import CheckboxRenderer
|
|
from copienator.commands import reading_grouped_annotations as reader
|
|
from copienator.left_margin import analyse_margins
|
|
from copienator.pdf_cut import split_pdf
|
|
|
|
|
|
class AnnotationLayoutTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.temp.cleanup)
|
|
self.root = Path(self.temp.name)
|
|
self.pdf = self.root / 'Copies' / 'Copie01.pdf'
|
|
self.pdf.parent.mkdir()
|
|
with pymupdf.open() as document:
|
|
for _ in range(2):
|
|
page = document.new_page(width=600, height=800)
|
|
page.insert_text((100, 100), 'Answer')
|
|
document.save(self.pdf)
|
|
|
|
def test_all_pages_measured_after_early_no_margin_run(self):
|
|
images = [Image.new('RGB', (600, 100)) for _ in range(6)]
|
|
values = [10, 10, 40, None, 19.9, 20]
|
|
estimates = {id(image): value for image, value in zip(images, values)}
|
|
def estimate(image, dpi):
|
|
edge = estimates[id(image)]
|
|
return {'edge_mm': edge, 'bands': 0 if edge is None else 8,
|
|
'band_starts_mm': [] if edge is None else [edge] * 8}
|
|
with patch('copienator.left_margin.estimate_page', side_effect=estimate):
|
|
result = analyse_margins(images, 200, 600)
|
|
self.assertEqual([p['status'] for p in result['pages']],
|
|
['low', 'low', 'normal', 'unknown', 'low', 'low'])
|
|
self.assertAlmostEqual(result['pages'][3]['effective_edge_mm'], 19.98)
|
|
|
|
def test_split_records_pages_and_rejects_stale_measurements(self):
|
|
schema_path = self.root / 'Cutleft' / 'Copie01_schema.json'
|
|
atomic_write_json(schema_path, {
|
|
'source_sha256': fingerprint(self.pdf),
|
|
'auto_margin': {'pages': [{'page': 1, 'edge_mm': 40}, {'page': 2, 'edge_mm': 10}]},
|
|
})
|
|
staging = self.root / 'split'
|
|
staging.mkdir()
|
|
splitting_int._render_split_outputs(self.pdf, [('Ex 1', 0, 0, 10, 0, 10)], staging)
|
|
answer = staging / 'Ex 1.pdf'
|
|
self.assertEqual([p['page'] for p in answer_parts(answer)], [1, 2])
|
|
self.assertEqual(layout_for(answer)['page_margins'], [300, 450])
|
|
atomic_write_json(schema_path, {'source_sha256': 'stale', 'auto_margin': {
|
|
'pages': [{'page': 1, 'edge_mm': 10}]}})
|
|
splitting_int._render_split_outputs(self.pdf, [('Ex 1', 0, 0, 10, 0, 10)], staging)
|
|
self.assertEqual(layout_for(answer)['margin_left'], 300)
|
|
|
|
def test_configurable_margin_unknown_fallback_and_stale_sidecar(self):
|
|
self.assertEqual(layout_for(self.pdf)['margin_left'], 300)
|
|
save_parts(self.pdf, [{'margin_status': 'normal', 'margin_mm': 40},
|
|
{'margin_status': 'low', 'margin_mm': 20}])
|
|
with patch.object(configuration, 'ANNOTATION_MARGIN_POINTS', [(10, 500), (20, 450), (35, 300)]):
|
|
self.assertEqual(layout_for(self.pdf)['margin_left'], 450)
|
|
self.pdf.write_bytes(self.pdf.read_bytes() + b'\n')
|
|
self.assertEqual(layout_for(self.pdf)['margin_left'], 300)
|
|
|
|
def test_horizontal_answer_crop_can_remove_an_original_margin(self):
|
|
atomic_write_json(self.root / 'Cutleft' / 'Copie01_schema.json', {
|
|
'source_sha256': fingerprint(self.pdf),
|
|
'auto_margin': {'pages': [{'page': 1, 'edge_mm': 40}, {'page': 2, 'edge_mm': 40}]},
|
|
})
|
|
staging = self.root / 'horizontal'
|
|
staging.mkdir()
|
|
# 20% of a 600-point page removes 42.3 mm from the left.
|
|
splitting_int._render_split_outputs(self.pdf, [('|Ex 1', 0, 0, 10, 100, 110)], staging)
|
|
answer = staging / 'Ex 1.pdf'
|
|
self.assertEqual(layout_for(answer)['margin_left'], 450)
|
|
self.assertEqual(answer_parts(answer)[0]['margin_mm'], 0)
|
|
|
|
def test_copy_move_and_vertical_cut_preserve_page_classification(self):
|
|
save_parts(self.pdf, [{'margin_status': 'normal', 'margin_mm': 40},
|
|
{'margin_status': 'low', 'margin_mm': 20}])
|
|
copied = self.root / 'copied.pdf'
|
|
moved = self.root / 'moved.pdf'
|
|
copy_answer_pdf(self.pdf, copied)
|
|
move_answer_pdf(copied, moved)
|
|
self.assertEqual(layout_for(moved)['margin_left'], 400)
|
|
first, second = self.root / 'first.pdf', self.root / 'second.pdf'
|
|
split_pdf(moved, 50, first, second)
|
|
self.assertEqual(layout_for(first)['margin_left'], 300)
|
|
self.assertEqual(layout_for(second)['margin_left'], 400)
|
|
|
|
def render(self, layout):
|
|
callback = Mock()
|
|
result = {'score': 2, 'feedback': [{'text': 'Feedback', 'box_2d': [10, 50, 30, 100]}]}
|
|
with patch.object(annotating, 'render_score_text', return_value=Image.new('RGBA', (200, 40))):
|
|
image, header = annotating.compose_label_image(
|
|
Image.new('RGB', (1654, 100), 'blue'), 'Ex 1', result, 0,
|
|
render_fn=lambda *args, **kwargs: Image.new('RGBA', (200, 10)),
|
|
draw_callback=callback, layout=layout)
|
|
box = next(call.args[2]['box'] for call in callback.call_args_list if call.args[0] == 'local_rect')
|
|
return image, header, box
|
|
|
|
def test_interpolation_and_clamping(self):
|
|
for mm, pixels in [(None, 300), (0, 450), (10, 450), (15, 425),
|
|
(20, 400), (27.5, 350), (35, 300), (80, 300)]:
|
|
with self.subTest(mm=mm):
|
|
self.assertEqual(added_margin(mm), pixels)
|
|
|
|
def test_unknown_measurement_inherits_average_during_answer_extraction(self):
|
|
atomic_write_json(self.root / 'Cutleft' / 'Copie01_schema.json', {
|
|
'source_sha256': fingerprint(self.pdf),
|
|
'auto_margin': {'pages': [{'page': 1, 'edge_mm': 15}, {'page': 2, 'edge_mm': None}]},
|
|
})
|
|
staging = self.root / 'inferred'
|
|
staging.mkdir()
|
|
splitting_int._render_split_outputs(self.pdf, [('Ex 1', 0, 0, 10, 0, 10)], staging)
|
|
answer = staging / 'Ex 1.pdf'
|
|
self.assertEqual(answer_parts(answer)[1]['margin_status'], 'unknown')
|
|
self.assertEqual(answer_parts(answer)[1]['margin_mm'], 15)
|
|
self.assertEqual(layout_for(answer)['page_margins'], [425, 425])
|
|
|
|
def test_per_page_offsets_narrow_pages_and_saved_layout(self):
|
|
save_parts(self.pdf, [{'margin_mm': 10}, {'margin_mm': 35}])
|
|
pages = [Image.new('RGB', (1000, 100), 'blue'), Image.new('RGB', (800, 100), 'green')]
|
|
with patch.object(annotating, 'convert_from_path', return_value=pages):
|
|
base, _, _ = annotating.make_base_image(self.pdf)
|
|
result = {'score': 2, 'feedback': [{'text': 'Across pages', 'box_2d': [80, 150, 120, 200]}]}
|
|
callback = Mock()
|
|
with patch.object(annotating, 'render_score_text', return_value=Image.new('RGB', (200, 40))):
|
|
def render(**kwargs):
|
|
return annotating.compose_label_image(
|
|
base, 'Ex 1', result, 0, render_fn=lambda *a, **k: Image.new('RGBA', (10, 10)),
|
|
**kwargs)
|
|
image, header = render(source_pdf=self.pdf, draw_callback=callback)
|
|
boxes = [call.args[2]['box'] for call in callback.call_args_list if call.args[0] == 'local_rect']
|
|
self.assertEqual(boxes, [[600, 120, 650, 140], [350, 140, 400, 160]])
|
|
for x, y, color in [(450, header + 20, (0, 0, 255)), (300, header + 150, (0, 128, 0))]:
|
|
self.assertEqual(image.getpixel((x, y)), color)
|
|
self.assertEqual(image.getpixel((x - 1, y)), (255, 255, 255))
|
|
with patch.object(configuration, 'ANNOTATION_MARGIN_POINTS', [(10, 900)]):
|
|
restored, _ = render(layout=image.info['annotation_layout'])
|
|
self.assertEqual(restored.tobytes(), image.tobytes())
|
|
|
|
def test_wider_layout_translates_copy_and_feedback_without_scaling(self):
|
|
normal, height, box = self.render(legacy_layout())
|
|
low, low_height, low_box = self.render({**legacy_layout(), 'margin_left': 400})
|
|
self.assertEqual(low.width, normal.width + 100)
|
|
self.assertEqual(height, low_height)
|
|
self.assertEqual(low_box, [box[0] + 100, box[1], box[2] + 100, box[3]])
|
|
self.assertEqual(low.getpixel((400, height + 80)), (0, 0, 255))
|
|
self.assertEqual(low.getpixel((399, height + 80)), (255, 255, 255))
|
|
|
|
def test_later_rectangle_cannot_draw_inside_comment_checkbox(self):
|
|
renderer = CheckboxRenderer('Ex 1')
|
|
result = {'score': 2, 'feedback': [
|
|
{'text': 'First comment', 'box_2d': [10, 50, 50, 200]},
|
|
{'text': 'Later rectangle', 'box_2d': [10, 300, 80, 400]},
|
|
]}
|
|
with patch.object(annotating, 'render_score_text', return_value=Image.new('RGB', (800, 80), 'white')):
|
|
image, _ = annotating.compose_label_image(
|
|
Image.new('RGB', (1654, 100), 'white'), 'Ex 1', result, 0,
|
|
render_fn=lambda *args, **kwargs: Image.new('RGBA', (620, 50), (255, 200, 200, 180)),
|
|
draw_callback=renderer.callback, layout=legacy_layout())
|
|
# The later rectangle's left edge crosses the first comment's checkbox.
|
|
checkbox = next(item for item in renderer.checkboxes
|
|
if item['type'] == 'del_local' and item['index'] == 0)
|
|
x1, y1, x2, y2 = checkbox['final_box']
|
|
interior = image.crop((x1 + 5, y1 + 5, x2 - 5, y2 - 5))
|
|
self.assertEqual(interior.getextrema(), ((255, 255),) * 3)
|
|
# The rectangle still exists outside the checkbox.
|
|
self.assertEqual(image.getpixel((650, 90)), (255, 0, 0))
|
|
|
|
def test_mixed_bgnot_widths_and_saved_geometry_survive_config_change(self):
|
|
normal, height, _ = self.render(legacy_layout())
|
|
low, _, _ = self.render({**legacy_layout(), 'margin_left': 400})
|
|
annotating_by_label.save_batch([
|
|
('01', 'Ex 1', normal, height, []), ('02', 'Ex 1', low, height, []),
|
|
], 'Ex 1', 1, self.root)
|
|
output = self.root / 'Ex 1 G1'
|
|
meta = read_json(output / 'bnote.json')
|
|
self.assertEqual(meta['width'], low.width)
|
|
with Image.open(output / 'Reference.jpg') as reference:
|
|
self.assertGreater(min(reference.getpixel((low.width - 10, 80))), 245)
|
|
with patch.object(reader, 'detect_checks_and_notes', return_value=([], Image.new('RGBA', (low.width, low.height * 2)))), patch.object(configuration, 'ANNOTATION_MARGIN_POINTS', [(10, 500)]):
|
|
_, notes = reader._scan_annotation_directory(output)
|
|
self.assertEqual(notes['02']['Ex 1']['layout']['margin_left'], 400)
|
|
rerender, _, _ = self.render(notes['02']['Ex 1']['layout'])
|
|
self.assertEqual(rerender.width, low.width)
|
|
with patch.object(reader, 'detect_checks_and_notes', return_value=([], None)):
|
|
_, notes = reader._scan_annotation_directory(output)
|
|
self.assertEqual(notes['02']['Ex 1']['layout']['margin_left'], 400)
|
|
self.assertIsNone(notes['02']['Ex 1']['img'])
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|