Repository: mindee/doctr

Issue:
[demo] Improve UI for OCR result display
For very dense documents, since the predicted text value is plotted statically, there can be some readability issues. We should try to improve this


Discussion:
I would argue that this was solved by #320
But perhaps we should wait for additional feedback on the readability of this new layout, what do you think @charlesmindee?

Resolving patch:
diff --git a/demo/app.py b/demo/app.py
--- a/demo/app.py
+++ b/demo/app.py
@@ -33,10 +33,14 @@ def main():
     st.title("DocTR: Document Text Recognition")
     # 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))
-    cols[0].subheader("Input document (first page)")
-    cols[1].subheader("Raw heatmap (segmentation task)")
+    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
@@ -50,7 +54,8 @@ def main():
             doc = DocumentFile.from_pdf(uploaded_file.read()).as_images(output_size=(1024, 1024))
         else:
             doc = DocumentFile.from_images(uploaded_file.read())
-        cols[0].image(doc[0], width=640)
+        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")
@@ -60,7 +65,7 @@ def main():
     # For newline
     st.sidebar.write('\n')
 
-    if st.sidebar.button("Analyze document"):
+    if st.sidebar.button("Analyze page"):
 
         if uploaded_file is None:
             st.sidebar.write("Please upload a document")
@@ -72,11 +77,11 @@ def main():
             with st.spinner('Analyzing...'):
 
                 # Forward the image to the model
-                processed_batches = predictor.det_predictor.pre_processor(doc)
+                processed_batches = predictor.det_predictor.pre_processor([doc[page_idx]])
                 out = predictor.det_predictor.model(processed_batches[0], return_model_output=True, training=False)
                 seg_map = out["out_map"]
                 seg_map = tf.squeeze(seg_map[0, ...], axis=[2])
-                seg_map = cv2.resize(seg_map.numpy(), (doc[0].shape[1], doc[0].shape[0]),
+                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()
@@ -85,15 +90,18 @@ def main():
                 cols[1].pyplot(fig)
 
                 # Plot OCR output
-                out = predictor(doc, training=False)
-                cols[1].subheader("OCR output")
-                fig = visualize_page(out.pages[0].export(), doc[0], interactive=False)
-                cols[1].pyplot(fig)
+                out = predictor([doc[page_idx]], training=False)
+                fig = visualize_page(out.pages[0].export(), doc[page_idx], interactive=False)
+                cols[2].pyplot(fig)
 
                 # Page reconsitution under input page
-                cols[0].subheader("Page reconstitution from OCR output")
-                img = synthetize_page(out.pages[0].export())
-                cols[0].image(img, clamp=True, width=640)
+                page_export = out.pages[0].export()
+                img = synthetize_page(page_export)
+                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__':


Test patch: