| import cv2 |
| import matplotlib.pyplot as plt |
| import numpy as np |
| import streamlit as st |
| import json |
|
|
| from doctr.file_utils import is_tf_available |
| from doctr.io import DocumentFile |
| from doctr.utils.visualization import visualize_page |
| from doctr.models import ocr_predictor |
|
|
| import torch |
| from backend.pytorch import DET_ARCHS, RECO_ARCHS, forward_image, load_predictor |
|
|
| forward_device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") |
|
|
| def extract_text_data(ocr_result, page_idx): |
| """Extracts and formats text data from OCR results, with the correct page index.""" |
| page_data = {"page_idx": page_idx, "text": ""} |
| for block in ocr_result.pages[0].blocks: |
| for line in block.lines: |
| for word in line.words: |
| page_data["text"] += word.value + " " |
| page_data["text"] += "\n " |
| return page_data |
|
|
| def main(det_archs, reco_archs): |
| """Build a Streamlit layout""" |
| st.set_page_config(layout="wide") |
|
|
| st.title("docTR: Document Text Recognition") |
| st.write("\n") |
| st.markdown("Hint: click on the top-right corner of an image to enlarge it!") |
|
|
| |
| cols = st.columns((1, 1, 1, 1)) |
| cols[0].subheader("Input page") |
| cols[1].subheader("Segmentation heatmap") |
| cols[2].subheader("OCR output") |
| cols[3].subheader("Page reconstitution") |
|
|
| st.sidebar.title("Document selection") |
| uploaded_file = st.sidebar.file_uploader("Upload files", type=["pdf", "png", "jpeg", "jpg"]) |
| if uploaded_file is not None: |
| if uploaded_file.name.endswith(".pdf"): |
| doc = DocumentFile.from_pdf(uploaded_file.read()) |
| else: |
| doc = DocumentFile.from_images(uploaded_file.read()) |
| cols[0].image(doc) |
|
|
| st.sidebar.title("Model selection") |
| st.sidebar.markdown("*Backend*: " + ("TensorFlow" if is_tf_available() else "PyTorch")) |
| det_arch = st.sidebar.selectbox("Text detection model", det_archs) |
| reco_arch = st.sidebar.selectbox("Text recognition model", reco_archs) |
| |
| assume_straight_pages = True |
| straighten_pages = False |
| bin_thresh = 0.3 |
| box_thresh = 0.1 |
|
|
| if st.sidebar.button("Analyze document"): |
| with st.spinner("Loading model..."): |
| predictor = load_predictor( |
| det_arch, reco_arch, assume_straight_pages, straighten_pages, bin_thresh, box_thresh, forward_device |
| ) |
|
|
| with st.spinner("Analyzing..."): |
| all_pages_export = [] |
| all_pages_text = [] |
| for page_idx, page in enumerate(doc): |
| st.write(f"Processing page {page_idx + 1}/{len(doc)}...") |
| |
| seg_map = forward_image(predictor, page, forward_device) |
| seg_map = np.squeeze(seg_map) |
| seg_map = cv2.resize(seg_map, (page.shape[1], page.shape[0]), interpolation=cv2.INTER_LINEAR) |
|
|
| fig, ax = plt.subplots() |
| ax.imshow(seg_map) |
| ax.axis("off") |
| cols[1].pyplot(fig) |
| |
| out = predictor([page]) |
| fig = visualize_page(out.pages[0].export(), out.pages[0].page, interactive=False, add_labels=False) |
| cols[2].pyplot(fig) |
|
|
| page_export = out.pages[0].export() |
| all_pages_export.append(page_export) |
| img = out.pages[0].synthesize() |
| cols[3].image(img, clamp=True) |
|
|
| page_text = extract_text_data(out, page_idx) |
| all_pages_text.append(page_text) |
|
|
| st.markdown("\nHere are your analysis results in JSON format for all pages:") |
| st.json(all_pages_text, expanded=False) |
|
|
| st.markdown("\nDownload here:") |
| st.download_button(label="Download JSON", data=json.dumps(all_pages_text), file_name="OCR.json", mime="application/json") |
| |
| |
| |
| |
| |
| |
|
|
| if __name__ == "__main__": |
| main(DET_ARCHS, RECO_ARCHS) |
|
|
