Spaces:
Running
Running
| from pathlib import Path | |
| import os | |
| os.umask(0o077) | |
| import streamlit as st | |
| from display_language import english_message, diagnostic_display | |
| from core import JOBS, ROOT, new_job, launch, read, save, is_running, accept, bundle | |
| from download_names import book_details, markdown_name | |
| from chrome_job import ensure_bridge | |
| if os.environ.get('BOOK_OCR_ENABLE_BRIDGE')=='1': | |
| st.cache_resource(ensure_bridge)() | |
| st.set_page_config(page_title='Book OCR Studio',page_icon='📖',layout='wide') | |
| st.markdown('''<style>[data-testid="stStatusWidget"]{display:none!important}.block-container{max-width:1500px;padding-top:2rem}h1{letter-spacing:-.04em}textarea:disabled{-webkit-text-fill-color:#334155!important;color:#334155!important;opacity:1!important}div[data-testid="stMetric"]{background:#f1f5f9;border-radius:10px;padding:12px}</style>''',unsafe_allow_html=True) | |
| st.title('Book OCR Studio') | |
| st.subheader('Local-first book OCR') | |
| st.caption('Local processing by default. Import PDFs, images or Kindle captures; export MD, HTML, PDF and EPUB with Gemma review or an optional vision API.') | |
| st.caption('Output preserves the language of the source; it is not translated. Recognition accuracy varies by language and layout.') | |
| with st.sidebar: | |
| st.subheader('Local processing') | |
| st.write('Marker / YomiToku + Gemma 4 (12B default)') | |
| st.caption('Default Gemma review is local. The optional API connector sends images and OCR text to your selected endpoint only when enabled. Internet is needed for initial downloads and Kindle access.') | |
| st.divider() | |
| jobs=sorted(JOBS.glob('*/job.json'),reverse=True) if JOBS.exists() else [] | |
| selected=st.selectbox('Saved jobs',['New job']+[p.parent.name for p in jobs]) | |
| if selected!='New job': st.session_state['job']=str(JOBS/selected) | |
| st.caption('Original OCR is preserved. Model suggestions enter the reviewed version only after approval.') | |
| if not st.session_state.get('c1_12b_default_applied'): | |
| st.session_state['shared_model']='gemma4:12b-it-qat' | |
| st.session_state['c1_12b_default_applied']=True | |
| model=st.selectbox('Gemma review model — Möbius Custom C1 for OCR',['gemma4:12b-it-qat','gemma4:26b-a4b-it-qat'],key='shared_model') | |
| st.caption('Review profile: Möbius Custom C1 for OCR (PDF, images and Kindle). 12B is the default; 26B uses the same profile.') | |
| st.caption('12B + YomiToku supports parallel OCR and review on one GPU. A 16GB or larger card is recommended, with at least 14.5GB free VRAM. Select One GPU, parallel in GPU mode; speed gains depend on the input.') | |
| st.caption('Timing: on one 16GB GPU, 26B runs OCR and review sequentially. This may take longer than 12B in single-GPU parallel mode. Timing varies with text volume and free VRAM; the end-to-end ratio has not been measured.') | |
| from connector_ui import render as render_connector | |
| review_options=render_connector(model) | |
| engines=['Marker'] | |
| if (ROOT/'.venv-yomitoku/bin/python').exists():engines.append('YomiToku') | |
| engine=st.selectbox('OCR engine (PDF, images and Kindle)',engines,key='shared_ocr_engine',help='Marker is the default. YomiToku is optional and has non-commercial terms unless separately licensed.') | |
| if engine=='YomiToku': | |
| st.warning('YomiToku 0.15.0 code and weights: CC BY-NC-SA 4.0. Use only when your use is permitted by those terms or covered by a separate commercial license. Selecting it does not grant a license.') | |
| else: | |
| st.caption('Marker 1.10.2 / Surya 0.17.1: GPL code; model weights have separate modified AI Pubs Open RAIL-M conditions. Marker is not an unrestricted commercial-use alternative. See Licenses and source below.') | |
| with st.expander('Licenses and source'): | |
| st.markdown((ROOT/'THIRD_PARTY_NOTICES.md').read_text(encoding='utf-8')) | |
| st.download_button('Application license (AGPL-3.0-only)',(ROOT/'LICENSE').read_bytes(),file_name='LICENSE.txt',key='license-download') | |
| upload_tab, kindle_tab=st.tabs(['PDF / Images','Kindle capture']) | |
| with upload_tab: | |
| page_scope=st.radio('Pages to process',['Select a page range','All pages'],horizontal=True) | |
| all_pages=page_scope=='All pages' | |
| with st.form('upload'): | |
| uploads=st.file_uploader('One PDF or multiple page images',type=['pdf','png','jpg','jpeg','webp'],accept_multiple_files=True) | |
| a,b,c=st.columns(3) | |
| with a: spec=st.text_input('Page range','1-3',disabled=all_pages,help='Page numbers start at 1, for example 1-3,5. This field is ignored for All pages. Images are sorted by filename.') | |
| with b: mode=st.selectbox('GPU mode',['Automatic','Two GPUs in parallel','One GPU, sequential','One GPU, parallel (12B + YomiToku)']) | |
| with c: | |
| gemma=st.checkbox('Review against images with selected model',value=True) | |
| force=st.checkbox('Force OCR on text-based PDFs',value=False) | |
| st.caption('All input pages will be processed.' if all_pages else 'Only the selected pages will be processed, for example 1-3,5.') | |
| submitted=st.form_submit_button('Start conversion',type='primary',disabled=not st.session_state.get('review_ready',True)) | |
| if submitted: | |
| try: | |
| files=sorted([(u.name,u.getvalue()) for u in uploads],key=lambda x:x[0]) | |
| job=new_job(files,'' if all_pages else spec,gemma,force,{'Automatic':'auto','Two GPUs in parallel':'dual','One GPU, sequential':'sequential','One GPU, parallel (12B + YomiToku)':'shared'}[mode],ocr_engine=engine.lower(),**review_options) | |
| st.session_state['job']=str(job);launch(job) | |
| except Exception as exc:st.error(str(exc)) | |
| with kindle_tab: | |
| import kindle_ui,importlib | |
| importlib.reload(kindle_ui) | |
| kindle_ui.render_kindle() | |
| from ui_refresh import adaptive_fragment,job_pending | |
| def show_job(): | |
| if 'job' not in st.session_state:return | |
| job=Path(st.session_state['job']);cfg=read(job/'job.json');s=read(job/'status.json') | |
| st.divider();st.subheader(cfg['inputs'][0]['name']) | |
| st.caption(f'Saved to: {job}') | |
| st.caption(f"OCR engine: {cfg.get('ocr_engine','marker')}") | |
| st.caption(f"Review model: {cfg.get('model')} · Provider: {'OpenAI-compatible API' if cfg.get('review_connector') else 'Local Gemma'}") | |
| st.info(english_message(s['message'])) | |
| if (job/'fallback.json').exists(): | |
| with st.expander('VRAM fallback history'): | |
| for event in read(job/'fallback.json'): | |
| st.write(f"{event['event']} → {event['mode']}: {english_message(event['reason'])}") | |
| if (job/'execution.json').exists(): | |
| execution=read(job/'execution.json') | |
| st.caption(f"Execution mode: {execution['mode']} / OCR: GPU {execution['ocr_gpu']} / Gemma: GPU {execution['gemma_gpu']}") | |
| st.caption(english_message(execution.get('reason',''))) | |
| with st.expander('GPU assignment details'): | |
| st.json(diagnostic_display(execution)) | |
| complete=sum((job/f'page-{i+1:05d}'/'marker.json').exists() for i in cfg['selected']) | |
| checked=sum((job/f'page-{i+1:05d}'/'review.json').exists() for i in cfg['selected']) | |
| a,b,c=st.columns(3) | |
| a.metric('OCR completed',f"{complete} / {len(cfg['selected'])}");b.metric('Gemma reviewed',checked);c.metric('Total input pages',cfg['total_pages']) | |
| st.caption('Progress covers the selected input only. It does not confirm full-book capture or error-free text.') | |
| running=is_running(job) | |
| if running: | |
| if st.button('Stop after current request',key='cancel'+job.name): | |
| (job/'cancel').touch();st.warning('Processing will stop after the current Gemma response.') | |
| elif s['state'] not in {'done'}: | |
| resume_models=['gemma4:12b-it-qat','gemma4:26b-a4b-it-qat'] | |
| resume_model=st.selectbox('Review model for resume',resume_models,index=resume_models.index(cfg['model']),key='resumemodel'+job.name) | |
| shared_resume=st.checkbox('Resume in single-GPU parallel mode (12B + YomiToku only)',value=cfg.get('mode')=='shared',key='resumeshared'+job.name) | |
| st.caption('Completed reviews are preserved. Changing the model means completed and remaining pages may use different models.') | |
| if st.button('Resume from saved progress',key='resume'+job.name): | |
| if shared_resume and (resume_model!='gemma4:12b-it-qat' or cfg.get('ocr_engine')!='yomitoku'): | |
| st.error('Select YomiToku and 12B for single-GPU parallel mode.') | |
| else: | |
| cfg['model']=resume_model | |
| cfg['mode']='shared' if shared_resume else ('auto' if cfg.get('mode')=='shared' else cfg.get('mode','auto')) | |
| save(job/'job.json',cfg);launch(job) | |
| ready=[i+1 for i in cfg['selected'] if (job/f'page-{i+1:05d}'/'marker.json').exists()] | |
| from frontier_export import render_button as render_frontier | |
| from fast_review import render as render_fast_review, adopt_pending | |
| with st.expander('Quick mode — model adoption and full-context handoff'): | |
| st.caption('Skip individual approval: adopt valid undecided model candidates, then prepare the complete handoff. Original OCR and previous page decisions are preserved. Failed or stale reviews remain in the handoff but are not adopted; deferred proposals outside a candidate are not inserted.') | |
| st.warning('Not manually verified. A downstream model cannot reliably recover missing text or correct every wrong number, name or negation.') | |
| if st.button('Adopt model candidates for Quick mode',key='bulk-adopt-'+job.name,disabled=running or not ready): | |
| count=adopt_pending(job,ready) | |
| st.session_state.pop('frontier-'+job.name+'-snapshot',None) | |
| st.success(f'Adopted {count} pages without manual verification. Original OCR is preserved.') | |
| render_frontier(st,job,'frontier-'+job.name,running) | |
| if ready: | |
| with st.expander('Fast review — pending changes and exceptions',expanded=False): | |
| render_fast_review(st,job,ready,running,accept) | |
| num=st.selectbox('Page to inspect',ready,key='page'+job.name) | |
| folder=job/f'page-{num:05d}' | |
| left,right=st.columns([1,1]) | |
| with left:st.image(str(folder/'source.png'),caption=f'Source page {num}',use_container_width=True) | |
| with right: | |
| original=(folder/'original.md').read_text() | |
| t1,t2,t3=st.tabs(['Original OCR','Gemma suggestions','Changes']) | |
| with t1:st.text_area('Original text (saved)',original,height=430,disabled=True,key=f'original{job.name}{num}') | |
| with t2: | |
| if (folder/'review.json').exists(): | |
| review=read(folder/'review.json');candidate=(folder/'candidate.md').read_text() | |
| st.caption(f"Review model for this page: {review.get('model','Unknown')}") | |
| st.text_area('Suggested text from Gemma',candidate,height=330,disabled=True,key=f'candidate{job.name}{num}') | |
| for correction in review['corrections']:st.write(correction) | |
| for note in review['notes']:st.caption(note) | |
| decision=read(folder/'decision.json') if (folder/'decision.json').exists() else {} | |
| st.caption(('Quick mode: adopted without manual verification' if decision.get('method')=='bulk_model_adoption' else 'Approved') if decision.get('accepted') else 'Not approved') | |
| if not running: | |
| x,y=st.columns(2) | |
| if x.button('Approve suggestions for this page',key=f'accept{job.name}{num}'):accept(job,num,True);st.rerun(scope='fragment') | |
| if y.button('Restore original OCR',key=f'reject{job.name}{num}'):accept(job,num,False);st.rerun(scope='fragment') | |
| elif (folder/'review-error.json').exists():st.error(english_message(read(folder/'review-error.json')['error'])) | |
| else:st.caption('Waiting for review, or Gemma review is disabled.') | |
| with t3: | |
| st.code((folder/'changes.diff').read_text() if (folder/'changes.diff').exists() else 'No changes to display yet',language='diff') | |
| if not running and (job/'original.md').exists(): | |
| title,author=book_details(job,cfg) | |
| name_key=job.name+str((job/'job.json').stat().st_mtime_ns) | |
| with st.expander('Download filename settings'): | |
| title=st.text_input('Book title',value=title,key='booktitle'+name_key) | |
| author=st.text_input('Author',value=author,key='bookauthor'+name_key) | |
| override=st.text_input('Filename (optional)',value=cfg.get('download_name',''),key='downloadname'+name_key,help='Leave blank for Book title_Author.md. If the author is unknown, only the title is used.') | |
| st.caption('PDF title and author metadata are used when available. Otherwise, the input filename provides the title.') | |
| if st.button('Save title, author and filename',key='savename'+job.name): | |
| cfg.update(title=title,author=author,download_name=override);save(job/'job.json',cfg) | |
| from delivery import ensure_delivery | |
| ensure_delivery(job) | |
| st.rerun() | |
| filename=markdown_name(title,author,override) | |
| st.caption(f'Download filename: {filename}') | |
| x,y,z=st.columns(3) | |
| x.download_button('Original OCR MD',data=(job/'original.md').read_bytes(),file_name=filename,key='download1'+job.name) | |
| y.download_button('Reviewed MD',data=(job/'reviewed.md').read_bytes(),file_name=filename,key='download2'+job.name) | |
| z.download_button('Images, suggestions and records ZIP',data=bundle(job),file_name=job.name+'.zip',key='download3'+job.name) | |
| from delivery_ui import output_buttons | |
| output_buttons(job,'job-'+job.name) | |
| st.caption('The reviewed version uses suggestions only on approved pages; all other pages retain the original OCR.') | |
| with st.expander('Processing log'): | |
| st.code(english_message((job/'worker.log').read_text(errors='replace')[-6000:]) if (job/'worker.log').exists() else 'No log available yet') | |
| show_job() | |