MarketOne/OCRMO

1# Copyright (C) 2021, Mindee.2 3# This program is licensed under the Apache License version 2.4# See LICENSE or go to <https://www.apache.org/licenses/LICENSE-2.0.txt> for full license details.5 6import os7 8import matplotlib.pyplot as plt9import streamlit as st10 11os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"12 13import cv214import tensorflow as tf15 16gpu_devices = tf.config.experimental.list_physical_devices('GPU')17if any(gpu_devices):18    tf.config.experimental.set_memory_growth(gpu_devices[0], True)19 20from doctr.io import DocumentFile21from doctr.models import ocr_predictor22from doctr.utils.visualization import visualize_page23 24DET_ARCHS = ["db_resnet50", "db_mobilenet_v3_large"]25RECO_ARCHS = ["crnn_vgg16_bn", "crnn_mobilenet_v3_small", "master", "sar_resnet31"]26 27 28def main():29 30    # Wide mode31    st.set_page_config(layout="wide")32 33    # Designing the interface34    st.title("docTR: Document Text Recognition")35    # For newline36    st.write('\n')37    #38    st.write('Find more info at: https://marketone.co/')39    # For newline40    st.write('\n')41    # Instructions42    st.markdown("*Hint: click on the top-right corner of an image to enlarge it!*")43    # Set the columns44    cols = st.beta_columns((1, 1, 1, 1))45    cols[0].subheader("Input page")46    cols[1].subheader("Segmentation heatmap")47    cols[2].subheader("OCR output")48    cols[3].subheader("Page reconstitution")49 50    # Sidebar51    # File selection52    st.sidebar.title("Document selection")53    # Disabling warning54    st.set_option('deprecation.showfileUploaderEncoding', False)55    # Choose your own image56    uploaded_file = st.sidebar.file_uploader("Upload files", type=['pdf', 'png', 'jpeg', 'jpg'])57    if uploaded_file is not None:58        if uploaded_file.name.endswith('.pdf'):59            doc = DocumentFile.from_pdf(uploaded_file.read())60        else:61            doc = DocumentFile.from_images(uploaded_file.read())62        page_idx = st.sidebar.selectbox("Page selection", [idx + 1 for idx in range(len(doc))]) - 163        cols[0].image(doc[page_idx])64 65    # Model selection66    st.sidebar.title("Model selection")67    det_arch = st.sidebar.selectbox("Text detection model", DET_ARCHS)68    reco_arch = st.sidebar.selectbox("Text recognition model", RECO_ARCHS)69 70    # For newline71    st.sidebar.write('\n')72 73    if st.sidebar.button("Analyze page"):74 75        if uploaded_file is None:76            st.sidebar.write("Please upload a document")77 78        else:79            with st.spinner('Loading model...'):80                predictor = ocr_predictor(det_arch, reco_arch, pretrained=True)81 82            with st.spinner('Analyzing...'):83 84                # Forward the image to the model85                processed_batches = predictor.det_predictor.pre_processor([doc[page_idx]])86                out = predictor.det_predictor.model(processed_batches[0], return_model_output=True)87                seg_map = out["out_map"]88                seg_map = tf.squeeze(seg_map[0, ...], axis=[2])89                seg_map = cv2.resize(seg_map.numpy(), (doc[page_idx].shape[1], doc[page_idx].shape[0]),90                                     interpolation=cv2.INTER_LINEAR)91                # Plot the raw heatmap92                fig, ax = plt.subplots()93                ax.imshow(seg_map)94                ax.axis('off')95                cols[1].pyplot(fig)96 97                # Plot OCR output98                out = predictor([doc[page_idx]])99                fig = visualize_page(out.pages[0].export(), doc[page_idx], interactive=False)100                cols[2].pyplot(fig)101 102                # Page reconsitution under input page103                page_export = out.pages[0].export()104                img = out.pages[0].synthesize()105                cols[3].image(img, clamp=True)106 107                # Display JSON108                st.markdown("\nHere are your analysis results in JSON format:")109                st.json(page_export)110 111 112if __name__ == '__main__':113    main()