moebiusT7/book-ocr-studio

1from pathlib import Path2import os3os.umask(0o077)4import streamlit as st5from display_language import english_message, diagnostic_display6from core import JOBS, ROOT, new_job, launch, read, save, is_running, accept, bundle7from download_names import book_details, markdown_name8 9from chrome_job import ensure_bridge10if os.environ.get('BOOK_OCR_ENABLE_BRIDGE')=='1':11    st.cache_resource(ensure_bridge)()12 13st.set_page_config(page_title='Book OCR Studio',page_icon='📖',layout='wide')14st.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)15st.title('Book OCR Studio')16st.subheader('Local-first book OCR')17st.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.')18st.caption('Output preserves the language of the source; it is not translated. Recognition accuracy varies by language and layout.')19with st.sidebar:20    st.subheader('Local processing')21    st.write('Marker / YomiToku + Gemma 4 (12B default)')22    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.')23    st.divider()24    jobs=sorted(JOBS.glob('*/job.json'),reverse=True) if JOBS.exists() else []25    selected=st.selectbox('Saved jobs',['New job']+[p.parent.name for p in jobs])26    if selected!='New job': st.session_state['job']=str(JOBS/selected)27    st.caption('Original OCR is preserved. Model suggestions enter the reviewed version only after approval.')28 29if not st.session_state.get('c1_12b_default_applied'):30    st.session_state['shared_model']='gemma4:12b-it-qat'31    st.session_state['c1_12b_default_applied']=True32model=st.selectbox('Gemma review model — Möbius Custom C1 for OCR',['gemma4:12b-it-qat','gemma4:26b-a4b-it-qat'],key='shared_model')33st.caption('Review profile: Möbius Custom C1 for OCR (PDF, images and Kindle). 12B is the default; 26B uses the same profile.')34st.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.')35st.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.')36from connector_ui import render as render_connector37review_options=render_connector(model)38engines=['Marker']39if (ROOT/'.venv-yomitoku/bin/python').exists():engines.append('YomiToku')40engine=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.')41if engine=='YomiToku':42    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.')43else:44    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.')45with st.expander('Licenses and source'):46    st.markdown((ROOT/'THIRD_PARTY_NOTICES.md').read_text(encoding='utf-8'))47    st.download_button('Application license (AGPL-3.0-only)',(ROOT/'LICENSE').read_bytes(),file_name='LICENSE.txt',key='license-download')48upload_tab, kindle_tab=st.tabs(['PDF / Images','Kindle capture'])49with upload_tab:50    page_scope=st.radio('Pages to process',['Select a page range','All pages'],horizontal=True)51    all_pages=page_scope=='All pages'52    with st.form('upload'):53        uploads=st.file_uploader('One PDF or multiple page images',type=['pdf','png','jpg','jpeg','webp'],accept_multiple_files=True)54        a,b,c=st.columns(3)55        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.')56        with b: mode=st.selectbox('GPU mode',['Automatic','Two GPUs in parallel','One GPU, sequential','One GPU, parallel (12B + YomiToku)'])57        with c:58            gemma=st.checkbox('Review against images with selected model',value=True)59            force=st.checkbox('Force OCR on text-based PDFs',value=False)60        st.caption('All input pages will be processed.' if all_pages else 'Only the selected pages will be processed, for example 1-3,5.')61        submitted=st.form_submit_button('Start conversion',type='primary',disabled=not st.session_state.get('review_ready',True))62    if submitted:63        try:64            files=sorted([(u.name,u.getvalue()) for u in uploads],key=lambda x:x[0])65            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)66            st.session_state['job']=str(job);launch(job)67        except Exception as exc:st.error(str(exc))68with kindle_tab:69    import kindle_ui,importlib70    importlib.reload(kindle_ui)71    kindle_ui.render_kindle()72 73from ui_refresh import adaptive_fragment,job_pending74 75@adaptive_fragment(3,lambda: job_pending(st.session_state.get('job')))76def show_job():77    if 'job' not in st.session_state:return78    job=Path(st.session_state['job']);cfg=read(job/'job.json');s=read(job/'status.json')79    st.divider();st.subheader(cfg['inputs'][0]['name'])80    st.caption(f'Saved to: {job}')81    st.caption(f"OCR engine: {cfg.get('ocr_engine','marker')}")82    st.caption(f"Review model: {cfg.get('model')} · Provider: {'OpenAI-compatible API' if cfg.get('review_connector') else 'Local Gemma'}")83    st.info(english_message(s['message']))84    if (job/'fallback.json').exists():85        with st.expander('VRAM fallback history'):86            for event in read(job/'fallback.json'):87                st.write(f"{event['event']} → {event['mode']}: {english_message(event['reason'])}")88    if (job/'execution.json').exists():89        execution=read(job/'execution.json')90        st.caption(f"Execution mode: {execution['mode']} / OCR: GPU {execution['ocr_gpu']} / Gemma: GPU {execution['gemma_gpu']}")91        st.caption(english_message(execution.get('reason','')))92        with st.expander('GPU assignment details'):93            st.json(diagnostic_display(execution))94    complete=sum((job/f'page-{i+1:05d}'/'marker.json').exists() for i in cfg['selected'])95    checked=sum((job/f'page-{i+1:05d}'/'review.json').exists() for i in cfg['selected'])96    a,b,c=st.columns(3)97    a.metric('OCR completed',f"{complete} / {len(cfg['selected'])}");b.metric('Gemma reviewed',checked);c.metric('Total input pages',cfg['total_pages'])98    st.caption('Progress covers the selected input only. It does not confirm full-book capture or error-free text.')99    running=is_running(job)100    if running:101        if st.button('Stop after current request',key='cancel'+job.name):102            (job/'cancel').touch();st.warning('Processing will stop after the current Gemma response.')103    elif s['state'] not in {'done'}:104        resume_models=['gemma4:12b-it-qat','gemma4:26b-a4b-it-qat']105        resume_model=st.selectbox('Review model for resume',resume_models,index=resume_models.index(cfg['model']),key='resumemodel'+job.name)106        shared_resume=st.checkbox('Resume in single-GPU parallel mode (12B + YomiToku only)',value=cfg.get('mode')=='shared',key='resumeshared'+job.name)107        st.caption('Completed reviews are preserved. Changing the model means completed and remaining pages may use different models.')108        if st.button('Resume from saved progress',key='resume'+job.name):109            if shared_resume and (resume_model!='gemma4:12b-it-qat' or cfg.get('ocr_engine')!='yomitoku'):110                st.error('Select YomiToku and 12B for single-GPU parallel mode.')111            else:112                cfg['model']=resume_model113                cfg['mode']='shared' if shared_resume else ('auto' if cfg.get('mode')=='shared' else cfg.get('mode','auto'))114                save(job/'job.json',cfg);launch(job)115    ready=[i+1 for i in cfg['selected'] if (job/f'page-{i+1:05d}'/'marker.json').exists()]116    from frontier_export import render_button as render_frontier117    from fast_review import render as render_fast_review, adopt_pending118    with st.expander('Quick mode — model adoption and full-context handoff'):119        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.')120        st.warning('Not manually verified. A downstream model cannot reliably recover missing text or correct every wrong number, name or negation.')121        if st.button('Adopt model candidates for Quick mode',key='bulk-adopt-'+job.name,disabled=running or not ready):122            count=adopt_pending(job,ready)123            st.session_state.pop('frontier-'+job.name+'-snapshot',None)124            st.success(f'Adopted {count} pages without manual verification. Original OCR is preserved.')125        render_frontier(st,job,'frontier-'+job.name,running)126    if ready:127        with st.expander('Fast review — pending changes and exceptions',expanded=False):128            render_fast_review(st,job,ready,running,accept)129        num=st.selectbox('Page to inspect',ready,key='page'+job.name)130        folder=job/f'page-{num:05d}'131        left,right=st.columns([1,1])132        with left:st.image(str(folder/'source.png'),caption=f'Source page {num}',use_container_width=True)133        with right:134            original=(folder/'original.md').read_text()135            t1,t2,t3=st.tabs(['Original OCR','Gemma suggestions','Changes'])136            with t1:st.text_area('Original text (saved)',original,height=430,disabled=True,key=f'original{job.name}{num}')137            with t2:138                if (folder/'review.json').exists():139                    review=read(folder/'review.json');candidate=(folder/'candidate.md').read_text()140                    st.caption(f"Review model for this page: {review.get('model','Unknown')}")141                    st.text_area('Suggested text from Gemma',candidate,height=330,disabled=True,key=f'candidate{job.name}{num}')142                    for correction in review['corrections']:st.write(correction)143                    for note in review['notes']:st.caption(note)144                    decision=read(folder/'decision.json') if (folder/'decision.json').exists() else {}145                    st.caption(('Quick mode: adopted without manual verification' if decision.get('method')=='bulk_model_adoption' else 'Approved') if decision.get('accepted') else 'Not approved')146                    if not running:147                        x,y=st.columns(2)148                        if x.button('Approve suggestions for this page',key=f'accept{job.name}{num}'):accept(job,num,True);st.rerun(scope='fragment')149                        if y.button('Restore original OCR',key=f'reject{job.name}{num}'):accept(job,num,False);st.rerun(scope='fragment')150                elif (folder/'review-error.json').exists():st.error(english_message(read(folder/'review-error.json')['error']))151                else:st.caption('Waiting for review, or Gemma review is disabled.')152            with t3:153                st.code((folder/'changes.diff').read_text() if (folder/'changes.diff').exists() else 'No changes to display yet',language='diff')154        if not running and (job/'original.md').exists():155            title,author=book_details(job,cfg)156            name_key=job.name+str((job/'job.json').stat().st_mtime_ns)157            with st.expander('Download filename settings'):158                title=st.text_input('Book title',value=title,key='booktitle'+name_key)159                author=st.text_input('Author',value=author,key='bookauthor'+name_key)160                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.')161                st.caption('PDF title and author metadata are used when available. Otherwise, the input filename provides the title.')162                if st.button('Save title, author and filename',key='savename'+job.name):163                    cfg.update(title=title,author=author,download_name=override);save(job/'job.json',cfg)164                    from delivery import ensure_delivery165                    ensure_delivery(job)166                    st.rerun()167            filename=markdown_name(title,author,override)168            st.caption(f'Download filename: {filename}')169            x,y,z=st.columns(3)170            x.download_button('Original OCR MD',data=(job/'original.md').read_bytes(),file_name=filename,key='download1'+job.name)171            y.download_button('Reviewed MD',data=(job/'reviewed.md').read_bytes(),file_name=filename,key='download2'+job.name)172            z.download_button('Images, suggestions and records ZIP',data=bundle(job),file_name=job.name+'.zip',key='download3'+job.name)173            from delivery_ui import output_buttons174            output_buttons(job,'job-'+job.name)175            st.caption('The reviewed version uses suggestions only on approved pages; all other pages retain the original OCR.')176    with st.expander('Processing log'):177        st.code(english_message((job/'worker.log').read_text(errors='replace')[-6000:]) if (job/'worker.log').exists() else 'No log available yet')178show_job()179