271 lines
19 KiB
Python
271 lines
19 KiB
Python
"""PDF -> source-traceable, reflowable EPUB. No translation calls or heavy models."""
|
||
import hashlib
|
||
import html
|
||
import json
|
||
import platform
|
||
import re
|
||
import sys
|
||
import uuid
|
||
from pathlib import Path
|
||
from ebooklib import epub
|
||
from .layout import paragraphs, repeated_furniture, join_text
|
||
|
||
SCHEMA = 1
|
||
|
||
|
||
def digest(path):
|
||
h = hashlib.sha256()
|
||
with Path(path).open('rb') as f:
|
||
for chunk in iter(lambda:f.read(1024*1024), b''): h.update(chunk)
|
||
return h.hexdigest()
|
||
|
||
|
||
def save_json(path, data):
|
||
path = Path(path); path.parent.mkdir(parents=True, exist_ok=True)
|
||
temp = path.with_suffix('.tmp'); temp.write_text(json.dumps(data, ensure_ascii=False, indent=2)+'\n'); temp.replace(path)
|
||
|
||
|
||
def page_selection(spec, count):
|
||
if not spec: return list(range(count))
|
||
selected = set()
|
||
for item in spec.split(','):
|
||
if not re.fullmatch(r'\d+(?:-\d+)?', item): raise ValueError('Pages must look like 1-20,25 (1-based inclusive)')
|
||
pair = [int(x) for x in item.split('-')]; lo, hi = pair[0], pair[-1]
|
||
if not 1 <= lo <= hi <= count: raise ValueError('Page range is outside the source PDF')
|
||
selected.update(range(lo-1, hi))
|
||
return sorted(selected)
|
||
|
||
|
||
def native_lines(page):
|
||
text = page.get_textpage()
|
||
try:
|
||
lines=[]; width,height=page.get_size()
|
||
for i in range(text.count_rects()):
|
||
left,bottom,right,top=text.get_rect(i)
|
||
value=text.get_text_bounded(left,bottom,right,top).strip()
|
||
if value:
|
||
lines.append({'text':value,'bbox':[left/width,1-top/height,right/width,1-bottom/height], 'confidence':None})
|
||
return lines
|
||
finally: text.close()
|
||
|
||
|
||
def image_boxes(page):
|
||
import pypdfium2.raw as raw
|
||
w,h=page.get_size(); boxes=[]
|
||
for obj in page.get_objects(filter=[raw.FPDF_PAGEOBJ_IMAGE]):
|
||
left,bottom,right,top=obj.get_bounds()
|
||
boxes.append([max(0,left/w), max(0,1-top/h), min(1,right/w), min(1,1-bottom/h)])
|
||
return boxes
|
||
|
||
|
||
def validate_lines(lines):
|
||
if not isinstance(lines,list): raise ValueError('OCR lines must be an array')
|
||
for line in lines:
|
||
box=line.get('bbox',[])
|
||
if not isinstance(line.get('text'),str) or len(box)!=4 or not all(isinstance(v,(int,float)) and 0<=v<=1 for v in box) or box[0]>=box[2] or box[1]>=box[3]:
|
||
raise ValueError('Invalid OCR text or normalized top-left bounding box')
|
||
|
||
|
||
def apply_corrections(pages, correction, source_sha):
|
||
"""Source-bound exact edits; never overwrite the raw OCR checkpoint."""
|
||
if correction.get('source_sha256')!=source_sha or not isinstance(correction.get('edits'),list):
|
||
raise ValueError('Corrections must be bound to this PDF SHA and contain an edits array')
|
||
lines={line['id']:line for page in pages for line in page['lines']};applied=[];seen=set()
|
||
for edit in correction['edits']:
|
||
key=edit.get('line_id');line=lines.get(key)
|
||
if key in seen or line is None or line['text']!=edit.get('expected') or not isinstance(edit.get('replacement'),str) or not edit['replacement'].strip():
|
||
raise ValueError('Stale, duplicate or invalid correction: '+str(key))
|
||
line['text']=edit['replacement'];seen.add(key);applied.append(edit)
|
||
return applied
|
||
|
||
|
||
def convert(source, output, work, language='zh-CN', title=None, author=None,
|
||
engine='auto', pages=None, dpi=180, cover_page=None, ocr_timeout=120, image_pages=None, corrections=None, ocr_url=None, ocr_model='firered-ocr', ocr_prompt=None, geometry_work=None):
|
||
import pypdfium2 as pdfium
|
||
source,output,work=Path(source),Path(output),Path(work)
|
||
if not source.is_file(): raise ValueError('PDF source does not exist')
|
||
if source.resolve()==output.resolve(): raise ValueError('Output must not overwrite source')
|
||
if output.suffix.lower()!='.epub': raise ValueError('Output must have .epub suffix')
|
||
if not 96<=dpi<=400: raise ValueError('DPI must be between 96 and 400')
|
||
work.mkdir(parents=True,exist_ok=True); source_sha=digest(source)
|
||
source_identity=work/'source.json'
|
||
if source_identity.exists() and json.loads(source_identity.read_text())['sha256']!=source_sha:
|
||
raise ValueError('Work directory belongs to another source; use a new directory')
|
||
save_json(source_identity,{'path':str(source.resolve()),'sha256':source_sha})
|
||
document=pdfium.PdfDocument(source); ocr=None
|
||
try:
|
||
selected=page_selection(pages,len(document)); partial=len(selected)!=len(document)
|
||
preserved_images=set(i+1 for i in page_selection(image_pages,len(document))) if image_pages else set()
|
||
if cover_page is not None and not 1<=cover_page<=len(document): raise ValueError('Cover page outside source PDF')
|
||
geometry_dir=Path(geometry_work) if geometry_work else None
|
||
if geometry_dir and json.loads((geometry_dir/'source.json').read_text())['sha256']!=source_sha:
|
||
raise ValueError('Reference geometry belongs to another PDF')
|
||
backend=('vision' if platform.system()=='Darwin' else 'tesseract') if engine=='auto' else engine
|
||
outlines=[]
|
||
for bookmark in document.get_toc():
|
||
dest=bookmark.get_dest()
|
||
if dest and dest.get_index() in selected:
|
||
outlines.append({'title':bookmark.get_title(),'page':dest.get_index()+1,'level':bookmark.level})
|
||
results=[]; warnings=[]
|
||
for index in selected:
|
||
page=document[index]
|
||
try:
|
||
boxes=image_boxes(page)
|
||
scan=any((b[2]-b[0])*(b[3]-b[1])>.65 for b in boxes)
|
||
native=native_lines(page) if engine=='native' or (engine=='auto' and not scan) else []
|
||
text_size=sum(len(l['text']) for l in native)
|
||
use_ocr=engine not in {'native','auto'} or (engine=='auto' and (scan or (text_size==0 and bool(boxes))))
|
||
is_cover=index+1==cover_page
|
||
is_image=is_cover or index+1 in preserved_images
|
||
if use_ocr and not is_image and ocr is None:
|
||
if backend=='firered':
|
||
from .firered import FireRedOCR
|
||
ocr=FireRedOCR(work,language,ocr_timeout,base_url=ocr_url,model=ocr_model,prompt_path=ocr_prompt)
|
||
else:
|
||
from .ocr import VisionOCR,TesseractOCR
|
||
ocr=(VisionOCR if backend=='vision' else TesseractOCR)(work,language,ocr_timeout)
|
||
used='cover-image' if is_cover else 'page-image' if is_image else backend if use_ocr else 'native'
|
||
stamp={'source':source_sha,'page':index+1,'engine':used,'ocr_identity':ocr.identity if use_ocr and not is_image else None,'language':language,'dpi':dpi,'pdfium':str(pdfium.PYPDFIUM_INFO),'schema':SCHEMA}
|
||
cache=work/'pages'/f'{index+1:04d}.json'
|
||
cached=json.loads(cache.read_text()) if cache.exists() else None
|
||
render=work/'renders'/f'{index+1:04d}.jpg';render.parent.mkdir(exist_ok=True)
|
||
if cached and cached.get('stamp')==stamp and render.exists():
|
||
result=cached
|
||
else:
|
||
bitmap=page.render(scale=dpi/72)
|
||
try: bitmap.to_pil().convert('RGB').save(render,quality=88)
|
||
finally: bitmap.close()
|
||
recognition=[] if is_image else ocr.recognize(render) if use_ocr else native
|
||
extra=recognition if isinstance(recognition,dict) else {}
|
||
lines=extra.get('lines',[]) if isinstance(recognition,dict) else recognition
|
||
validate_lines(lines)
|
||
for n,line in enumerate(lines): line['id']=f'p{index+1:04d}-l{n+1:04d}'
|
||
result={'number':index+1,'width':page.get_width(),'height':page.get_height(),'engine':used,'scan':scan,'lines':lines,'image_boxes':boxes,'render':str(render.resolve()),'stamp':stamp,**extra}
|
||
save_json(cache,result)
|
||
results.append(result)
|
||
print(f'PDF page {index+1}/{len(document)}: {used}, {len(result["lines"])} lines',file=sys.stderr,flush=True)
|
||
finally: page.close()
|
||
applied=apply_corrections(results,json.loads(Path(corrections).read_text()),source_sha) if corrections else []
|
||
repeated=repeated_furniture(results)
|
||
for page in results:
|
||
if 'markdown' in page:
|
||
from .markdown_layout import markdown_blocks
|
||
geometry=json.loads((geometry_dir/'pages'/f'{page["number"]:04d}.json').read_text()) if geometry_dir else None
|
||
if geometry and (geometry['number']!=page['number'] or geometry['stamp']['source']!=source_sha):
|
||
raise ValueError('Reference page geometry provenance mismatch')
|
||
blocks=markdown_blocks(page,language,geometry);removed=[]
|
||
else:
|
||
blocks,removed=paragraphs(page,language,repeated)
|
||
page['blocks']=blocks;page['removed_furniture']=removed
|
||
count=sum(len(b['text']) for b in blocks) if 'markdown' in page else sum(len(l['text']) for l in page['lines'])
|
||
page['image_fallback']=page['engine'] in {'cover-image','page-image'} or count<30
|
||
if page['image_fallback'] and page['engine'] not in {'cover-image','page-image'}:
|
||
warnings.append({'page':page['number'],'kind':'page_image_fallback','detail':'Blank/illustrated/short-text page retained as an image; inspect source, not claimed as reflowed text.'})
|
||
removed_ids={l['id'] for l in removed}
|
||
low=[l['id'] for l in page['lines'] if l['id'] not in removed_ids and l.get('confidence') is not None and l['confidence']<.6]
|
||
if low: warnings.append({'page':page['number'],'kind':'low_ocr_confidence','line_ids':low})
|
||
if engine=='native' and page['scan']: warnings.append({'page':page['number'],'kind':'existing_scan_text_layer','detail':'Existing OCR text layer reused without verification.'})
|
||
book={'schema':SCHEMA,'source_sha256':source_sha,'source':str(source.resolve()),'source_pages':len(document),'selected_pages':[i+1 for i in selected],'partial':partial,'language':language,'title':title or source.stem,'author':author or 'Unknown','outlines':outlines,'image_pages':sorted(preserved_images),'cover_page':cover_page,'geometry_work':str(geometry_dir.resolve()) if geometry_dir else None,'pages':results,'warnings':warnings,'applied_corrections':applied,'scope':'horizontal single-column; no automatic footnote/table reconstruction'}
|
||
save_json(work/'book.json',book)
|
||
export_epub(book,output,work,cover_page)
|
||
from src.qa.epub_audit import audit_epub
|
||
audit=audit_epub(output)
|
||
report={'source_sha256':source_sha,'output_sha256':digest(output),'output':str(output.resolve()),'partial':partial,'processed_pages':len(results),'source_pages':len(document),'image_only_pages':[p['number'] for p in results if p['image_fallback']],'removed_furniture_lines':sum(len(p['removed_furniture']) for p in results),'engines':{name:sum(p['engine']==name for p in results) for name in set(p['engine'] for p in results)},'warnings':warnings,'applied_corrections':applied,'epub_audit':audit,'visual_review':'pending','ocr_accuracy_review':'pending','translation':'not requested; English EPUB can be passed to main.py separately'}
|
||
save_json(work/'report.json',report)
|
||
if audit['issues']: raise ValueError('EPUB structural audit failed; inspect report.json')
|
||
return report
|
||
finally:
|
||
if ocr: ocr.close()
|
||
document.close()
|
||
|
||
|
||
STYLE='''body {margin: 5%; line-height:1.65;} p {margin:.5em 0;} html[lang^="zh"] p {text-indent:2em;} h1,h2 {text-align:center;page-break-before:always;} img {max-width:100%;height:auto;} figure {margin:1em 0;text-align:center;} figcaption {font-size:.8em;} .source-page {font-size:.7em;color:#777;text-indent:0;} .attribution {text-align:right !important;text-indent:0 !important;margin:.4em 0;} .image-page {text-align:center;text-indent:0;} .image-page img {max-height:90vh;width:auto;object-fit:contain;}'''
|
||
|
||
|
||
def export_epub(data, output, work, cover_page):
|
||
from PIL import Image
|
||
import io
|
||
book=epub.EpubBook(); identity={k:data[k] for k in ['source_sha256','language','title','author','selected_pages']}
|
||
identity['content_sha256']=hashlib.sha256(json.dumps([[p['number'],p['image_fallback'],[(b['kind'],b['text']) for b in p['blocks']]] for p in data['pages']],ensure_ascii=False).encode()).hexdigest()
|
||
book.set_identifier('urn:uuid:'+str(uuid.uuid5(uuid.NAMESPACE_URL,json.dumps(identity,sort_keys=True))))
|
||
label='【部分样例】' if data['partial'] else ''
|
||
book.set_title(label+data['title']+'【PDF重排初稿】');book.set_language(data['language']);book.add_author(data['author'])
|
||
book.add_metadata('DC','source','urn:sha256:'+data['source_sha256'])
|
||
book.add_metadata('DC','description','PDF结构重建初稿；OCR未逐字校对。'+('仅含选定页，非完整书。' if data['partial'] else ''))
|
||
css=epub.EpubItem(uid='style',file_name='styles/book.css',media_type='text/css',content=STYLE.encode());book.add_item(css)
|
||
chapters=[]; toc_records=[]; body=[]; current=None; previous=None; previous_page=None
|
||
def finish():
|
||
if current is not None:
|
||
current.content=('<html xmlns:epub="http://www.idpf.org/2007/ops"><head></head><body>'+''.join(body)+'</body></html>').encode()
|
||
current.add_link(href='../styles/book.css',rel='stylesheet',type='text/css')
|
||
book.add_item(current);chapters.append(current)
|
||
for page in data['pages']:
|
||
number=page['number']; outlines=[o for o in data['outlines'] if o['page']==number]
|
||
labels=[o['title'] for o in outlines]
|
||
if current is None or labels or page['image_fallback'] or previous_page is None or number!=previous_page+1:
|
||
finish();body=[];previous=None
|
||
current=epub.EpubHtml(title=' / '.join(labels) or f'PDF {number}',file_name=f'chapters/p{number:04d}.xhtml',lang=data['language'])
|
||
for outline in outlines:
|
||
toc_records.append((epub.Link(current.file_name+f'#pdf-page-{number}',outline['title'],f'nav-{len(toc_records)+1}'),outline['level']))
|
||
marker=f'<span id="pdf-page-{number}" epub:type="pagebreak" role="doc-pagebreak" aria-label="PDF {number}"></span>'
|
||
if not page['image_fallback']:
|
||
for heading in labels: body.append('<h1>'+html.escape(heading)+'</h1>')
|
||
if page['image_fallback']:
|
||
body.append(marker)
|
||
image=Path(page['render']).read_bytes();name=f'images/page-{number:04d}.jpg'
|
||
book.add_item(epub.EpubItem(uid=f'image-{number}',file_name=name,media_type='image/jpeg',content=image))
|
||
body.append(f'<div class="image-page"><img src="../{name}" alt="PDF第{number}页原貌"/></div>')
|
||
if number==cover_page: book.set_cover('images/cover.jpg',image,create_page=False)
|
||
previous=None
|
||
else:
|
||
entries=[dict(b) for b in page['blocks']]
|
||
for n,box in enumerate(page['image_boxes']):
|
||
area=(box[2]-box[0])*(box[3]-box[1])
|
||
if .02<area<.65: entries.append({'kind':'image','bbox':box,'image_index':n})
|
||
if all(b.get('bbox') for b in entries):entries.sort(key=lambda b:(b['bbox'][1],b['bbox'][0]))
|
||
marker_pending=True
|
||
for index,block in enumerate(entries):
|
||
if block['kind']=='image':
|
||
with Image.open(page['render']) as picture:
|
||
box=block['bbox'];region=picture.crop(tuple(int(v*(picture.width if i%2==0 else picture.height)) for i,v in enumerate(box)))
|
||
stream=io.BytesIO();region.save(stream,format='PNG')
|
||
n=block['image_index'];name=f'images/p{number:04d}-{n}.png'
|
||
book.add_item(epub.EpubItem(uid=f'figure-{number}-{n}',file_name=name,media_type='image/png',content=stream.getvalue()))
|
||
if marker_pending:body.append(marker);marker_pending=False
|
||
body.append(f'<figure><img src="../{name}" alt="PDF第{number}页图片"/></figure>');previous=None;continue
|
||
plain=re.sub(r'\s','',block['text'])
|
||
if any(plain==re.sub(r'\s','',label) or (index==0 and plain in {'序','序言','前言','自序','引言','后记','目录'} and re.sub(r'\s','',label).startswith(plain)) for label in labels):continue
|
||
# Merge only obvious prose continuation across consecutive source pages.
|
||
continuation=marker_pending and previous and not labels and previous_page==number-1 and previous['kind']=='paragraph' and block['kind']=='paragraph' and previous.get('role')!='attribution' and block.get('role')!='attribution' and not block['indented'] and not re.search(r'[。！？.!?:：;；”’"]$',previous['text']) and body[-1].endswith('</p>')
|
||
if continuation:
|
||
extra=join_text(previous['text'],block['text'],data['language'])[len(previous['text']):]
|
||
body[-1]=body[-1][:-4]+marker+f'<span id="{block["id"]}"></span>'+html.escape(extra)+'</p>'
|
||
marker_pending=False;previous={**block,'text':join_text(previous['text'],block['text'],data['language'])};continue
|
||
if marker_pending:body.append(marker);marker_pending=False
|
||
tag='h2' if block['kind']=='heading' else 'p'
|
||
attributes=' class="attribution"' if block.get('alignment')=='right' else ''
|
||
body.append(f'<{tag} id="{block["id"]}"{attributes}>{html.escape(block["text"])}</{tag}>')
|
||
if block['kind']=='heading' and not data['outlines']:
|
||
toc_records.append((epub.Link(current.file_name+'#'+block['id'],block['text'],f'nav-{len(toc_records)+1}'),0))
|
||
previous=block
|
||
if marker_pending:body.append(marker)
|
||
previous_page=number
|
||
if page['image_fallback']:
|
||
finish();current=None;body=[];previous=None
|
||
finish()
|
||
if not toc_records:toc_records=[(epub.Link(c.file_name,c.title,f'nav-{i+1}'),0) for i,c in enumerate(chapters)]
|
||
roots=[];stack=[]
|
||
for link,level in toc_records:
|
||
node={'link':link,'children':[]}
|
||
while stack and stack[-1][0]>=level:stack.pop()
|
||
(stack[-1][1]['children'] if stack else roots).append(node);stack.append((level,node))
|
||
def flatten(nodes):return [(node['link'],flatten(node['children'])) if node['children'] else node['link'] for node in nodes]
|
||
book.toc=flatten(roots);book.add_item(epub.EpubNcx());nav=epub.EpubNav();nav.title='目录' if data['language'].startswith('zh') else 'Contents';book.add_item(nav);book.spine=['nav']+chapters
|
||
output=Path(output);output.parent.mkdir(parents=True,exist_ok=True);temp=output.with_suffix('.tmp.epub')
|
||
try:
|
||
epub.write_epub(str(temp),book,{'epub3_pages':False})
|
||
from src.qa.epub_audit import audit_epub
|
||
if audit_epub(temp)['issues']:raise ValueError('Generated EPUB failed structural audit; existing output retained')
|
||
temp.replace(output)
|
||
finally:temp.unlink(missing_ok=True)
|