| 1 | import base64 |
| 2 | import io |
| 3 | from types import SimpleNamespace |
| 4 | import unittest |
| 5 | from unittest.mock import patch |
| 6 | import xml.etree.ElementTree as ET |
| 7 | |
| 8 | from PIL import Image |
| 9 | |
| 10 | from baseline_fixture import build |
| 11 | from compare_baselines import baseline_report, pdf_rgb, register_images |
| 12 | from native_fixture import NS |
| 13 | |
| 14 | |
| 15 | class BaselineAnchors(unittest.TestCase): |
| 16 | def setUp(self): |
| 17 | self.xml = ET.fromstring(build()['pages'][0]['xml']) |
| 18 | self.images = [] |
| 19 | for node in self.xml.findall(f'{{{NS}}}Image'): |
| 20 | with Image.open(io.BytesIO(base64.b64decode(node.findtext(f'{{{NS}}}Data')))) as image: |
| 21 | pixels, dimensions = image.convert('RGB').tobytes(), image.size |
| 22 | position, size = node.find(f'{{{NS}}}Position'), node.find(f'{{{NS}}}Size') |
| 23 | x, y = float(position.get('x')) + 36, float(position.get('y')) + 21.6 |
| 24 | self.images.append({ |
| 25 | 'srcsize': dimensions, 'bits': 8, 'colorspace': [SimpleNamespace(name='DeviceRGB')], |
| 26 | 'x0': x, 'x1': x + float(size.get('width')), 'top': y, 'bottom': y + float(size.get('height')), |
| 27 | 'stream': SimpleNamespace(attrs={}, get_data=lambda data=pixels: data), |
| 28 | }) |
| 29 | |
| 30 | def test_exact_pixels_establish_translation(self): |
| 31 | anchors = register_images(self.xml, [SimpleNamespace(images=self.images)])[0] |
| 32 | self.assertEqual(len(anchors), 3) |
| 33 | for anchor in anchors: |
| 34 | self.assertAlmostEqual(anchor['translation'][0], 36) |
| 35 | self.assertAlmostEqual(anchor['translation'][1], 21.6) |
| 36 | self.assertEqual(anchor['extent_scale'], [1, 1]) |
| 37 | |
| 38 | def test_same_dimensions_with_changed_pixels_are_not_an_anchor(self): |
| 39 | self.images[0]['stream'].get_data = lambda: b'\0' * (22 * 26 * 3) |
| 40 | anchors = register_images(self.xml, [SimpleNamespace(images=self.images)])[0] |
| 41 | self.assertEqual(len(anchors), 2) |
| 42 | self.assertTrue(all(a['source_xywh'][0] != 20.125 for a in anchors)) |
| 43 | |
| 44 | def test_duplicates_and_masks_are_not_usable(self): |
| 45 | self.images.append(self.images[0]) |
| 46 | self.images[1]['stream'].attrs['SMask'] = object() |
| 47 | anchors = register_images(self.xml, [SimpleNamespace(images=self.images)])[0] |
| 48 | self.assertEqual(len(anchors), 1) |
| 49 | self.assertEqual(anchors[0]['source_xywh'][0], 25.75) |
| 50 | |
| 51 | def test_indexed_four_bit_pixels_respect_row_padding(self): |
| 52 | colors = b'\xff\0\0\0\xff\0\0\0\xff' |
| 53 | image = {'srcsize': (3, 2), 'bits': 4, |
| 54 | 'colorspace': [SimpleNamespace(name='Indexed'), SimpleNamespace(name='DeviceRGB'), 2, colors], |
| 55 | 'stream': SimpleNamespace(attrs={}, get_data=lambda: b'\x01\x20\x01\x20')} |
| 56 | self.assertEqual(pdf_rgb(image), colors * 2) |
| 57 | |
| 58 | def test_text_geometry_cannot_change_image_registration(self): |
| 59 | probe = {'objects': [{'kind': 'outline', 'id': 'o', 'is_title': False, 'layout': {'y': 400}}]} |
| 60 | wraps = {'counts': {'matched': 1}, 'paragraphs': [{ |
| 61 | 'status': 'matched', 'outline_id': 'o', 'paragraph_id': 'p', 'canvas_baselines_in_outline': [10], |
| 62 | 'native_instances': [{'pdf_page': 1, 'baseline_ranges_from_pdf_top': [[430, 430]]}], |
| 63 | }]} |
| 64 | with patch('compare_baselines.compare', return_value=wraps): |
| 65 | before = baseline_report(probe, self.xml, [SimpleNamespace(images=self.images)]) |
| 66 | probe['objects'][0]['layout']['y'] += 20 |
| 67 | after = baseline_report(probe, self.xml, [SimpleNamespace(images=self.images)]) |
| 68 | self.assertEqual(before['anchors_by_pdf_page'], after['anchors_by_pdf_page']) |
| 69 | self.assertAlmostEqual(before['paragraphs'][0]['lines'][0]['residual_using_median_anchor_translation'][0], -1.6) |
| 70 | self.assertAlmostEqual(after['paragraphs'][0]['lines'][0]['residual_using_median_anchor_translation'][0], -21.6) |
| 71 | |
| 72 | def test_wrap_mismatch_remains_visible_without_baseline_measurements(self): |
| 73 | wraps = {'counts': {'different_wraps': 1}, 'paragraphs': [{'status': 'different_wraps', 'paragraph_id': 'p'}]} |
| 74 | with patch('compare_baselines.compare', return_value=wraps): |
| 75 | report = baseline_report({'objects': []}, self.xml, [SimpleNamespace(images=self.images)]) |
| 76 | self.assertEqual(report['paragraphs'], [{'status': 'different_wraps', 'paragraph_id': 'p', 'lines': []}]) |
| 77 | |
| 78 | |
| 79 | if __name__ == '__main__': |
| 80 | unittest.main() |