Spaces:
Sleeping
Sleeping
# Copyright (C) 2021, Mindee. | |
# This program is licensed under the Apache License version 2. | |
# See LICENSE or go to <https://www.apache.org/licenses/LICENSE-2.0.txt> for full license details. | |
import os | |
import matplotlib.pyplot as plt | |
import streamlit as st | |
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2" | |
import cv2 | |
import tensorflow as tf | |
gpu_devices = tf.config.experimental.list_physical_devices('GPU') | |
if any(gpu_devices): | |
tf.config.experimental.set_memory_growth(gpu_devices[0], True) | |
from doctr.io import DocumentFile | |
from doctr.models import ocr_predictor | |
from doctr.utils.visualization import visualize_page | |
DET_ARCHS = ["db_resnet50", "db_mobilenet_v3_large"] | |
RECO_ARCHS = ["crnn_vgg16_bn", "crnn_mobilenet_v3_small", "master", "sar_resnet31"] | |
def main(): | |
# Wide mode | |
st.set_page_config(layout="wide") | |
# Designing the interface | |
st.title("docTR: Document Text Recognition") | |
# For newline | |
st.write('\n') | |
# | |
st.write('Find more info at: https://github.com/mindee/doctr') | |
# 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.beta_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") | |
# Sidebar | |
# File selection | |
st.sidebar.title("Document selection") | |
# Disabling warning | |
st.set_option('deprecation.showfileUploaderEncoding', False) | |
# 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 | |
cols[0].image(doc[page_idx]) | |
# 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') | |
if st.sidebar.button("Analyze page"): | |
if uploaded_file is None: | |
st.sidebar.write("Please upload a document") | |
else: | |
with st.spinner('Loading model...'): | |
predictor = ocr_predictor(det_arch, reco_arch, pretrained=True) | |
with st.spinner('Analyzing...'): | |
# Forward the image to the model | |
processed_batches = predictor.det_predictor.pre_processor([doc[page_idx]]) | |
out = predictor.det_predictor.model(processed_batches[0], return_model_output=True) | |
seg_map = out["out_map"] | |
seg_map = tf.squeeze(seg_map[0, ...], axis=[2]) | |
seg_map = cv2.resize(seg_map.numpy(), (doc[page_idx].shape[1], doc[page_idx].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([doc[page_idx]]) | |
fig = visualize_page(out.pages[0].export(), doc[page_idx], interactive=False) | |
cols[2].pyplot(fig) | |
# Page reconsitution under input page | |
page_export = out.pages[0].export() | |
img = out.pages[0].synthesize() | |
cols[3].image(img, clamp=True) | |
# Display JSON | |
st.markdown("\nHere are your analysis results in JSON format:") | |
st.json(page_export) | |
if __name__ == '__main__': | |
main() |