216 lines
13 KiB
Python
216 lines
13 KiB
Python
"""Lightweight import: selection, provenance, OCR routing, reflow and EPUB links."""
|
|
import json
|
|
import zipfile
|
|
from pathlib import Path
|
|
import pytest
|
|
from src.pdf_import.converter import convert, page_selection, validate_lines
|
|
from src.pdf_import.layout import paragraphs, repeated_furniture, join_text
|
|
from src.qa.epub_audit import audit_epub
|
|
|
|
|
|
@pytest.fixture
|
|
def native_pdf(tmp_path):
|
|
pytest.importorskip('pypdfium2')
|
|
# Minimal native-text PDF fixture; no external PDF authoring dependency.
|
|
def write(name='input.pdf', image=False):
|
|
streams=[b'BT /F1 18 Tf 60 740 Td (Chapter 1: Start) Tj 0 -40 Td /F1 12 Tf (A paragraph with enough native English text to use without OCR.) Tj ET',b'BT /F1 12 Tf 60 740 Td (This is the second page with more native English prose.) Tj ET']
|
|
if image:streams=[b'q 600 0 0 800 0 0 cm /Im1 Do Q '+s for s in streams]
|
|
resources=b'<< /Font << /F1 7 0 R >>'+(b' /XObject << /Im1 8 0 R >>' if image else b'')+b' >>'
|
|
objects=[b'<< /Type /Catalog /Pages 2 0 R >>',b'<< /Type /Pages /Kids [3 0 R 4 0 R] /Count 2 >>']
|
|
objects.extend(b'<< /Type /Page /Parent 2 0 R /MediaBox [0 0 600 800] /Resources '+resources+b' /Contents '+str(n).encode()+b' 0 R >>' for n in [5,6])
|
|
objects.extend(b'<< /Length '+str(len(s)).encode()+b' >>\nstream\n'+s+b'\nendstream' for s in streams)
|
|
objects.append(b'<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>')
|
|
if image:objects.append(b'<< /Type /XObject /Subtype /Image /Width 1 /Height 1 /ColorSpace /DeviceRGB /BitsPerComponent 8 /Length 3 >>\nstream\n\xff\xff\xff\nendstream')
|
|
raw=b'%PDF-1.4\n';offsets=[0]
|
|
for i,obj in enumerate(objects):offsets.append(len(raw));raw+=f'{i+1} 0 obj\n'.encode()+obj+b'\nendobj\n'
|
|
xref=len(raw);raw+=f'xref\n0 {len(objects)+1}\n0000000000 65535 f \n'.encode()
|
|
raw+=b''.join(f'{n:010d} 00000 n \n'.encode() for n in offsets[1:])
|
|
raw+=f'trailer\n<< /Size {len(objects)+1} /Root 1 0 R >>\nstartxref\n{xref}\n%%EOF'.encode()
|
|
path=tmp_path/name;path.write_bytes(raw);return path
|
|
return write
|
|
|
|
|
|
def test_page_ranges_do_not_silently_drop_pages():
|
|
assert page_selection(None,3)==[0,1,2]
|
|
assert page_selection('3,1-2,2',3)==[0,1,2]
|
|
for value in ['0','3-2','1-4','-1','1,','abc']:
|
|
with pytest.raises(ValueError):page_selection(value,3)
|
|
|
|
|
|
def test_folio_ocr_variants_are_removed_but_body_years_survive():
|
|
lines=[{'id':'header','text':'002 | 中国的世界系','bbox':[.15,.04,.4,.06]}, {'id':'body','text':'2019年发生的事。','bbox':[.15,.12,.8,.14]}]
|
|
blocks,removed=paragraphs({'number':5,'lines':lines},'zh-CN',set())
|
|
assert [b['text'] for b in blocks]==['2019年发生的事。']
|
|
assert removed[0]['id']=='header'
|
|
assert join_text('跨页','继续。','zh-CN')=='跨页继续。'
|
|
assert join_text('Some','words','en')=='Some words'
|
|
|
|
|
|
def test_native_english_import_is_independent_and_complete(native_pdf,tmp_path):
|
|
source=native_pdf();out=tmp_path/'book.epub';work=tmp_path/'work'
|
|
report=convert(source,out,work,language='en',engine='native')
|
|
assert report['engines']=={'native':2} and not report['partial']
|
|
assert audit_epub(out)['issues']==[]
|
|
with zipfile.ZipFile(out) as z:
|
|
chapters=[n for n in z.namelist() if '/chapters/' in n]
|
|
assert len(chapters)==1 # Consecutive pages flow inside one reading document.
|
|
text=z.read(chapters[0]).decode()
|
|
assert 'native English' in text and 'second page' in text
|
|
assert 'pdf-page-1' in text and 'pdf-page-2' in text
|
|
assert '../styles/book.css' in text and 'lang="en"' in text
|
|
from src.preprocessing.epub_cleaner import EpubCleaner
|
|
from src.preprocessing.text_extractor import FineGrainedExtractor
|
|
from src.common.data_model import BookStructure
|
|
structure=BookStructure.load(EpubCleaner(out,tmp_path/'translation-input').clean())
|
|
units=FineGrainedExtractor(complete_book=True).extract(structure)
|
|
assert any('native English' in unit.original_text for unit in units)
|
|
ir=json.loads((work/'book.json').read_text())
|
|
assert len(ir['pages'])==2 and ir['pages'][0]['lines'][0]['bbox']
|
|
partial=convert(source,tmp_path/'partial.epub',work,language='en',engine='native',pages='2')
|
|
assert partial['partial'] and partial['processed_pages']==1
|
|
with zipfile.ZipFile(tmp_path/'partial.epub') as z:
|
|
opf=z.read('EPUB/content.opf').decode();assert '部分样例' in opf
|
|
with pytest.raises(ValueError,match='another source'):
|
|
convert(native_pdf('another.pdf',image=True),tmp_path/'other.epub',work,engine='native')
|
|
|
|
|
|
def test_scanned_text_layer_is_replaced_and_ocr_cache_is_reused(native_pdf,tmp_path,monkeypatch):
|
|
from src.pdf_import import ocr
|
|
import src.pdf_import.converter as converter
|
|
called=[]
|
|
class FakeOCR:
|
|
identity='fixture-v1'
|
|
def __init__(self,*args):pass
|
|
def recognize(self,image):
|
|
called.append(image)
|
|
return [{'text':'这是来自扫描图像的新识别结果，必须保留完整正文内容，不能使用已有错误文本层。','bbox':[.1,.2,.9,.25],'confidence':.9}]
|
|
def close(self):pass
|
|
monkeypatch.setattr(ocr,'VisionOCR',FakeOCR)
|
|
monkeypatch.setattr(ocr,'TesseractOCR',FakeOCR)
|
|
source=native_pdf(image=True);out=tmp_path/'scan.epub';work=tmp_path/'work'
|
|
convert(source,out,work,language='zh-CN',engine='auto')
|
|
assert len(called)==2
|
|
convert(source,out,work,language='zh-CN',engine='auto')
|
|
assert len(called)==2
|
|
with zipfile.ZipFile(out) as z:
|
|
body=b''.join(z.read(n) for n in z.namelist() if '/chapters/' in n).decode()
|
|
assert '新识别结果' in body and 'native English' not in body
|
|
assert 'lang="zh-CN"' in body
|
|
FakeOCR.identity='fixture-v2'
|
|
convert(source,out,work,language='zh-CN',engine='auto')
|
|
assert len(called)==4
|
|
|
|
|
|
def test_invalid_ocr_geometry_fails_closed():
|
|
with pytest.raises(ValueError):validate_lines([{'text':'bad','bbox':[.8,.2,.1,.4]}])
|
|
with pytest.raises(ValueError):validate_lines([{'text':'bad','bbox':[0,0,float('nan'),1]}])
|
|
|
|
|
|
def test_ocr_failure_leaves_page_checkpoint_and_no_partial_final(native_pdf,tmp_path,monkeypatch):
|
|
from src.pdf_import import ocr
|
|
called=[]
|
|
class FakeOCR:
|
|
identity='resume-fixture'
|
|
fail=True
|
|
def __init__(self,*args):pass
|
|
def recognize(self,image):
|
|
called.append(Path(image).name)
|
|
if Path(image).stem=='0002' and self.fail:raise RuntimeError('fixture OCR failure')
|
|
return [{'text':'A sufficiently long recognized paragraph to remain searchable in the EPUB.', 'bbox':[.1,.2,.9,.3], 'confidence':1}]
|
|
def close(self):pass
|
|
monkeypatch.setattr(ocr,'VisionOCR',FakeOCR);monkeypatch.setattr(ocr,'TesseractOCR',FakeOCR)
|
|
source=native_pdf(image=True);out=tmp_path/'output.epub';work=tmp_path/'work'
|
|
with pytest.raises(RuntimeError,match='fixture OCR failure'):convert(source,out,work,language='en')
|
|
assert not out.exists() and (work/'pages/0001.json').exists() and not (work/'pages/0002.json').exists()
|
|
FakeOCR.fail=False
|
|
convert(source,out,work,language='en')
|
|
assert called==['0001.jpg','0002.jpg','0002.jpg']
|
|
assert audit_epub(out)['issues']==[]
|
|
|
|
|
|
def test_consistently_indented_quote_lines_stay_one_paragraph():
|
|
lines=[{'id':'a','text':'这是正文段落。','bbox':[.1,.11,.9,.13]}, {'id':'b','text':'一段尚未结束的引文','bbox':[.2,.16,.8,.18]}, {'id':'c','text':'继续在相同缩进位置。','bbox':[.2,.19,.8,.21]}]
|
|
blocks,_=paragraphs({'number':1,'lines':lines},'zh-CN',set())
|
|
assert len(blocks)==2
|
|
assert blocks[1]['text']=='一段尚未结束的引文继续在相同缩进位置。'
|
|
|
|
|
|
def test_corrections_require_source_and_exact_expected_text():
|
|
from src.pdf_import.converter import apply_corrections
|
|
pages=[{'lines':[{'id':'p0001-l0001','text':'牲富'}]}]
|
|
patch={'source_sha256':'source','edits':[{'line_id':'p0001-l0001','expected':'牲富','replacement':'牲畜'}]}
|
|
with pytest.raises(ValueError):apply_corrections(pages,patch,'different-source')
|
|
applied=apply_corrections(pages,patch,'source')
|
|
assert pages[0]['lines'][0]['text']=='牲畜' and len(applied)==1
|
|
with pytest.raises(ValueError,match='Stale'):apply_corrections(pages,patch,'source')
|
|
|
|
|
|
def test_signature_and_date_keep_separate_right_aligned_paragraphs():
|
|
page={'number':1,'lines':[{'id':'body','text':'正文最后一段。','bbox':[.1,.4,.85,.44]}, {'id':'name','text':'作者甲','bbox':[.72,.50,.84,.52]}, {'id':'date','text':'2020.2.16','bbox':[.70,.54,.84,.56]}]}
|
|
blocks,_=paragraphs(page,'zh-CN',set())
|
|
assert [b['text'] for b in blocks][-2:]==['作者甲','2020.2.16']
|
|
assert all(b.get('alignment')=='right' and b.get('role')=='attribution' for b in blocks[-2:])
|
|
from src.pdf_import.markdown_layout import markdown_blocks
|
|
md=markdown_blocks({'number':1,'markdown':'识别后的正文最后一段。\n\n作者甲\n\n2020.2.16'},'zh-CN',page)
|
|
assert md[0]['text']=='识别后的正文最后一段。'
|
|
assert all(b.get('alignment')=='right' for b in md[-2:])
|
|
assert md[-1]['geometry_source']=='reference-vision-text-alignment'
|
|
assert md[-1]['line_ids']==[] # Never pretend FireRed emitted coordinates or Vision text.
|
|
|
|
|
|
def test_firered_http_cache_and_truncation_rejection(tmp_path,monkeypatch):
|
|
from PIL import Image
|
|
import io
|
|
import src.pdf_import.firered as firered
|
|
calls=[]
|
|
def response(request,timeout):
|
|
payload=json.loads(request.data);calls.append(payload)
|
|
return io.BytesIO(json.dumps({'model':'firered-ocr','choices':[{'finish_reason':'stop','message':{'content':'原文段落。\n\n作者甲'}}],'usage':{'completion_tokens':10}}).encode())
|
|
monkeypatch.setattr(firered,'urlopen',response)
|
|
image=tmp_path/'page.png';Image.new('RGB',(40,60),'white').save(image)
|
|
engine=firered.FireRedOCR(tmp_path,'zh-CN',base_url='http://ocr.test:8001')
|
|
result=engine.recognize(image);assert result['markdown'].startswith('原文段落') and result['lines']==[]
|
|
assert engine.recognize(image)==result and len(calls)==1
|
|
assert calls[0]['model']=='firered-ocr' and calls[0]['temperature']==0
|
|
def truncated(request,timeout):
|
|
return io.BytesIO(json.dumps({'choices':[{'finish_reason':'length','message':{'content':'仅一部分'}}]}).encode())
|
|
monkeypatch.setattr(firered,'urlopen',truncated)
|
|
rejected=firered.FireRedOCR(tmp_path/'rejected','zh-CN',base_url='http://ocr.test:8001')
|
|
with pytest.raises(RuntimeError,match='incomplete'):rejected.recognize(image)
|
|
assert list((tmp_path/'rejected/firered-raw').glob('*.rejected.json'))
|
|
|
|
|
|
def test_snippet_comparison_excludes_page_context_and_counts_missing_characters():
|
|
from scripts.compare_pdf_ocr import snippet_errors
|
|
assert snippet_errors('原始正文','页眉原始正文页脚')['errors']==0
|
|
result=snippet_errors('拥有识字能力','页眉有识字能力页脚')
|
|
assert result['errors']==1
|
|
assert sum(result[k] for k in ['substitutions','deletions','insertions'])==1
|
|
|
|
|
|
def test_markdown_ocr_is_exported_as_text_not_an_image(native_pdf,tmp_path,monkeypatch):
|
|
import src.pdf_import.firered as firered
|
|
from src.pdf_import.converter import digest,save_json
|
|
class FakeOCR:
|
|
identity='markdown-fixture'
|
|
def __init__(self,*args,**kwargs):pass
|
|
def recognize(self,image):return {'lines':[],'markdown':'这是重新识别后的完整正文段落，保留新模型文字而不是旧模型的文字内容。\n\n作者甲\n\n2020.2.16'}
|
|
def close(self):pass
|
|
monkeypatch.setattr(firered,'FireRedOCR',FakeOCR)
|
|
source=native_pdf(image=True);work=tmp_path/'work';geom=tmp_path/'geometry';save_json(geom/'source.json',{'sha256':digest(source)})
|
|
for n in [1,2]:
|
|
save_json(geom/'pages'/f'{n:04d}.json',{'number':n,'stamp':{'source':digest(source)},'lines':[{'id':f'p{n}-a','text':'这是旧模型的正文段落。','bbox':[.1,.2,.9,.3]},{'id':f'p{n}-b','text':'作者甲','bbox':[.72,.5,.84,.52]},{'id':f'p{n}-c','text':'2020.2.16','bbox':[.7,.54,.84,.56]}]})
|
|
out=tmp_path/'md.epub';convert(source,out,work,engine='firered',ocr_url='http://ocr.test',geometry_work=geom)
|
|
with zipfile.ZipFile(out) as z:
|
|
text=''.join(z.read(n).decode() for n in z.namelist() if '/chapters/' in n)
|
|
assert '重新识别后的完整正文' in text and 'class="attribution"' in text
|
|
assert '作者甲2020' not in text
|
|
assert audit_epub(out)['issues']==[]
|
|
|
|
|
|
def test_markdown_does_not_erase_uncertain_sentence_spacing():
|
|
from src.pdf_import.markdown_layout import markdown_blocks
|
|
blocks=markdown_blocks({'number':1,'markdown':'前一句 后一句'},'zh-CN')
|
|
assert blocks[0]['text']=='前一句 后一句'
|
|
assert blocks[0]['bbox'] is None
|