tests/reader_v3_layout_samples.py

"""Inspectable PDF calibration pages; geometry truth is independent of OCR sorting."""
import hashlib
import os
from pathlib import Path
import sys

import pymupdf

sys.path.insert(0, str(Path(__file__).resolve().parents[2] / 'pipeline'))
from scripts import ocr_layout, pdf_reading_v3 as v3, pdf_text_layer as text


def build_samples(bundle, *, table_rows=False):
    document = pymupdf.open()
    font = pymupdf.Font('china-s')
    descriptors = []
    for kind in ('horizontal', 'rotation-90', 'rotation-270', 'vertical', 'vertical-rotation-90', 'columns', 'table'):
        page = document.new_page(width=600, height=900)
        blocks = []
        expected = []

        def line(value, x, y, size=12):
            if any(not font.has_glyph(ord(char)) for char in value):
                raise ValueError('calibration font lacks a source glyph')
            page.insert_text((x, y + size), value, fontsize=size, fontname='china-s')
            width = sum(size if ord(c) > 255 else size / 2 for c in value)
            box = [x / 600, y / 900, (x + width) / 600, (y + size) / 900]
            blocks.append({'t': value, 'b': box, 'q': [[box[0], box[1]], [box[2], box[1]],
                                                   [box[2], box[3]], [box[0], box[3]]], 'c': 1, 's': 'native'})
            expected.append(value)

        options = {'writing_mode': 'horizontal-ltr', 'join_soft_lines': False}
        if kind in ('horizontal', 'rotation-90', 'rotation-270'):
            for row in range(20):
                line(f'第{row + 1}行 校验甲 校验乙 校验丙', 55, 60 + row * 35)
            if kind != 'horizontal':
                page.set_rotation(int(kind.split('-')[1]))
        elif kind.startswith('vertical'):
            options['writing_mode'] = 'vertical-rl'
            for x, value in ((460, '竖排右列 校验甲Aé漢 校验乙丙'), (320, '竖排左列校验庚辛壬校验癸子丑')):
                if any(not font.has_glyph(ord(char)) for char in value):
                    raise ValueError('calibration font lacks a source glyph')
                for index, char in enumerate(value):
                    page.insert_text((x, 60 + (index + 1) * 28), char, fontsize=20, fontname='china-s')
                # Include real word spaces and Latin advances to exercise attached-run measurement.
                box = [x / 600, 60 / 900, (x + 20) / 600, (60 + len(value) * 28) / 900]
                blocks.append({'t': value, 'b': box, 'c': 1, 's': 'native'})
                expected.append(value)
            if kind == 'vertical-rotation-90':
                page.set_rotation(90)
        elif kind == 'columns':
            line('通栏标题', 70, 20)
            for column, x in (('左栏', 50), ('右栏', 340)):
                for row in range(20):
                    line(f'{column}{row + 1} 校验甲 校验乙', x, 100 + row * 34)
        else:
            line('合并表头与跨列说明', 55, 30)
            # Header spans all columns; the first body label spans two rows.
            xs = [45, 170, 295, 420, 555]
            ys = [100 + row * 105 for row in range(8)]
            for index, y in enumerate(ys):
                page.draw_line((xs[1] if index == 1 else xs[0], y), (xs[-1], y), color=(0, 0, 0))
            for x in xs:
                page.draw_line((x, ys[0]), (x, ys[-1]), color=(0, 0, 0))
            for row in range(7):
                for column in range(4):
                    if row == 1 and column == 0:
                        continue
                    line(f'{row + 1}项{column + 1} 校验甲', xs[column] + 5, ys[row] + 18, 9)
                    line('数值 校验乙', xs[column] + 5, ys[row] + 42, 9)
            if table_rows:
                options['regions'] = [{'box': [0, 0, 1, .09], 'writing_mode': 'horizontal-ltr'}] + [
                    {'box': [xs[column] / 600, ys[row] / 900, xs[column + 1] / 600,
                             ys[row + (2 if row == column == 0 else 1)] / 900], 'writing_mode': 'horizontal-ltr'}
                    for row in range(7) for column in range(4) if not (row == 1 and column == 0)]
        descriptors.append({'kind': kind, 'blocks': blocks, 'expected_regions': expected, 'options': options,
                            'rotation': page.rotation})
    pdf = document.tobytes(no_new_id=True)
    source = hashlib.sha256(pdf).hexdigest()
    root = Path('objects') / source[:2] / source / ('c' * 16)
    primary_path = bundle / root / 'document.pdf'
    primary_path.parent.mkdir(parents=True, exist_ok=True)
    primary_path.write_bytes(pdf)
    primary = {'bucket': v3.shared.PDF_PAGES_BUCKET, 'path': primary_path.relative_to(bundle).as_posix(),
               'sha256': source, 'bytes': len(pdf), 'role': 'runtime'}
    layers = []
    for index, descriptor in enumerate(descriptors, 1):
        arranged = ocr_layout.arrange(descriptor['blocks'], 600, 900, descriptor['options'], include_writing_modes=True)
        rotation = descriptor['rotation']
        if rotation:
            for field in ('blocks', 'raw_blocks'):
                for block in arranged[field]:
                    quad = block.get('q') or [[block['b'][0], block['b'][1]], [block['b'][2], block['b'][1]],
                                              [block['b'][2], block['b'][3]], [block['b'][0], block['b'][3]]]
                    block['q'] = [list(ocr_layout.map_point_back(x, y, 360 - rotation)) for x, y in quad]
                    block['b'] = [min(p[0] for p in block['q']), min(p[1] for p in block['q']),
                                  max(p[0] for p in block['q']), max(p[1] for p in block['q'])]
            for span in arranged['text_spans']:
                span['box'] = next(block['b'] for block in arranged['blocks'] if block['id'] == span['block'])
        page = document[index - 1]
        layer = text.from_page({'kind': 'pdf-ocr-page', 'page': index, 'source': 'native',
                               'width': page.rect.width, 'height': page.rect.height, **arranged}, source, 'ch')
        layers.append(layer)
        descriptor['page'] = index
        descriptor['expected_boxes'] = [r['box'] for r in layer['regions']]
        descriptor['actual_regions'] = [r['text'] for r in layer['regions']]
    document.close()
    layer_ref = v3.build_text_bundle(layers, [{'p': p['page']} for p in layers], root, bundle)
    def read(ref):
        return (bundle / ref['path']).read_bytes()
    spec = {'source_key': 'calibration\0layout.pdf', 'source_sha256': source, 'primary': primary,
            'text_layer': layer_ref, 'require_complete_preview': True,
            'vector_previews': os.environ.get('READER_V3_VECTOR_PREVIEWS') == '1'}
    spec = v3.generate_previews(spec, read, bundle, dpi=100)
    ref, manifest = v3.build_reading(spec, read, bundle)
    return {'reading': ref, 'manifest': manifest, 'layers': layers, 'samples': descriptors,
            'query': '校验', 'expected_total': sum(value.count('校验') for d in descriptors for value in d['expected_regions'])}


def build_fractional_vector_sample(bundle):
    """A4-like native vector pages whose SVG intrinsic size is fractional (595.3 x 841.9)."""
    document = pymupdf.open()
    for number in range(1, 4):
        page = document.new_page(width=595.3, height=841.9)
        # Extractable text is what classifies a page as native-vector.
        page.insert_text((60, 100), f'Fractional vector page {number}', fontsize=18)
        page.draw_rect(pymupdf.Rect(40, 60, 555, 120), color=(0, 0, 0))
        page.draw_line((40, 200 + number * 10), (555, 200 + number * 10), color=(0, 0, 0))
    pdf = document.tobytes(no_new_id=True)
    document.close()
    source = hashlib.sha256(pdf).hexdigest()
    primary_path = bundle / 'objects' / source[:2] / source / ('d' * 16) / 'document.pdf'
    primary_path.parent.mkdir(parents=True, exist_ok=True)
    primary_path.write_bytes(pdf)
    primary = {'bucket': v3.shared.PDF_PAGES_BUCKET, 'path': primary_path.relative_to(bundle).as_posix(),
               'sha256': source, 'bytes': len(pdf), 'role': 'runtime'}
    def read(ref):
        return (bundle / ref['path']).read_bytes()
    spec = {'source_key': 'calibration\0fractional-vector.pdf', 'source_sha256': source, 'primary': primary,
            'require_complete_preview': True, 'vector_previews': True}
    spec = v3.generate_previews(spec, read, bundle, dpi=100)
    if any(entry['codec'] != 'svg' for entry in spec['previews']):
        raise ValueError('fractional fixture did not produce vector previews')
    ref, manifest = v3.build_reading(spec, read, bundle)
    return {'reading': ref, 'manifest': manifest, 'previews': spec['previews']}