Spaces:
Sleeping
Sleeping
| # import cv2 | |
| # import matplotlib.pyplot as plt | |
| import numpy as np | |
| import streamlit as st | |
| import torch | |
| import json | |
| import base64 | |
| from doctr.io import DocumentFile | |
| from doctr.utils.visualization import visualize_page | |
| from backend.pytorch import DET_ARCHS, RECO_ARCHS, load_predictor #forward_image | |
| forward_device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
| def main(det_archs, reco_archs): | |
| """Build a streamlit layout""" | |
| # Wide mode | |
| st.set_page_config(layout="wide") | |
| st.markdown("Used Github Actions to automatically build the app on any updates on this [github repo link](https://github.com/deepanshu2207/imgtotxt_using_DocTR)") | |
| st.caption("Made with ❤️ by Deepanshu. Credits to 🤗 Spaces for Hosting this.") | |
| # Designing the interface | |
| st.title("Document Text Extraction") | |
| # For newline | |
| st.write("\n") | |
| # Instructions | |
| st.markdown("*Hint: click on the top-right corner of an image to enlarge it!*") | |
| # Set the columns | |
| # cols = st.columns((1, 1, 1, 1)) | |
| cols = st.columns((1, 1, 1)) | |
| cols[0].subheader("Input page") | |
| # cols[1].subheader("Segmentation heatmap") | |
| cols[1].subheader("OCR output") | |
| cols[2].subheader("Page reconstitution") | |
| # Sidebar | |
| # File selection | |
| st.sidebar.title("Document selection") | |
| # Choose your own image | |
| 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()) | |
| page_idx = st.sidebar.selectbox("Page selection", [idx + 1 for idx in range(len(doc))]) - 1 | |
| page = doc[page_idx] | |
| cols[0].image(page) | |
| # Model selection | |
| st.sidebar.title("Model selection") | |
| det_arch = st.sidebar.selectbox("Text detection model", det_archs) | |
| reco_arch = st.sidebar.selectbox("Text recognition model", reco_archs) | |
| # # For newline | |
| # st.sidebar.write("\n") | |
| # # Only straight pages or possible rotation | |
| # st.sidebar.title("Parameters") | |
| # assume_straight_pages = st.sidebar.checkbox("Assume straight pages", value=True) | |
| # st.sidebar.write("\n") | |
| # # Straighten pages | |
| # straighten_pages = st.sidebar.checkbox("Straighten pages", value=False) | |
| # st.sidebar.write("\n") | |
| # # Binarization threshold | |
| # bin_thresh = st.sidebar.slider("Binarization threshold", min_value=0.1, max_value=0.9, value=0.3, step=0.1) | |
| # st.sidebar.write("\n") | |
| if st.sidebar.button("Analyze page"): | |
| if uploaded_file is None: | |
| st.sidebar.write("Please upload a document") | |
| else: | |
| with st.spinner("Loading model..."): | |
| # Default Values | |
| assume_straight_pages, straighten_pages, bin_thresh = True, False, 0.3 | |
| predictor = load_predictor( | |
| det_arch, reco_arch, assume_straight_pages, straighten_pages, bin_thresh, forward_device | |
| ) | |
| with st.spinner("Analyzing..."): | |
| # # Forward the image to the model | |
| # 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) | |
| # # Plot the raw heatmap | |
| # fig, ax = plt.subplots() | |
| # ax.imshow(seg_map) | |
| # ax.axis("off") | |
| # cols[1].pyplot(fig) | |
| # Plot OCR output | |
| out = predictor([page]) | |
| fig = visualize_page(out.pages[0].export(), out.pages[0].page, interactive=False, add_labels=False) | |
| cols[1].pyplot(fig) | |
| # Page reconsitution under input page | |
| page_export = out.pages[0].export() | |
| if assume_straight_pages or (not assume_straight_pages and straighten_pages): | |
| img = out.pages[0].synthesize() | |
| cols[2].image(img, clamp=True) | |
| print('out',out) | |
| print('\n') | |
| print('page_export',page_export) | |
| print('\n') | |
| all_text = '' | |
| for i in page_export['blocks']: | |
| for line in i['lines']: | |
| for word in line['words']: | |
| all_text+=word['value'] | |
| all_text+=' ' | |
| all_text+='\n' | |
| print('all_text', all_text) | |
| print('\n') | |
| # Display Text | |
| st.markdown("\n### **Here is your text:**") | |
| st.write(all_text) | |
| # Display JSON | |
| # json_string = json.dumps(page_export) | |
| st.markdown("\n### **Here is your document structure in JSON format:**") | |
| encoded_data = base64.b64encode(json.dumps(page_export).encode("utf-8")).decode("utf-8") | |
| download_link = f"data:file/txt;base64,{encoded_data}" | |
| st.markdown(f"[Download JSON]( {download_link} )", unsafe_allow_html=True) | |
| # st.download_button(label="Download JSON", data=json_string, file_name='data.json', mime='application/json') | |
| st.json(page_export, expanded=False) | |
| st.success('Done!') | |
| st.balloons() | |
| if __name__ == "__main__": | |
| main(DET_ARCHS, RECO_ARCHS) | |