feat: Added page selection & JSON display to demo

README.md

@@ -133,7 +133,7 @@ You can then easily run your app in your default browser by running:
 streamlit run demo/app.py
 ```
 
-![Demo app](https://github.com/mindee/doctr/releases/download/v0.2.1/demo.png)
+![Demo app](https://github.com/mindee/doctr/releases/download/v0.3.0/demo_update.png)
 
 ### Docker container
 

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__':