We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Demo app error when analyzing my first document
## 🐛 Bug

I tried to analyze a PNG and a PDF, got the same error. I try to change the model, didn't change anything.

## To Reproduce

Steps to reproduce the behavior:

1. Upload a PNG
2. Click on analyze document


<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->
```
KeyError: 0
Traceback:
File "/Users/thibautmorla/opt/anaconda3/lib/python3.8/site-packages/streamlit/script_runner.py", line 337, in _run_script
    exec(code, module.__dict__)
File "/Users/thibautmorla/Downloads/doctr/demo/app.py", line 93, in <module>
    main()
File "/Users/thibautmorla/Downloads/doctr/demo/app.py", line 77, in main
    seg_map = predictor.det_predictor.model(processed_batches[0])[0]
```


## Additional context

First image upload

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `demo/app.py`
Content:
```
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 
6 import os
7 import streamlit as st
8 import matplotlib.pyplot as plt
9 
10 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
11 
12 import tensorflow as tf
13 import cv2
14 
15 gpu_devices = tf.config.experimental.list_physical_devices('GPU')
16 if any(gpu_devices):
17     tf.config.experimental.set_memory_growth(gpu_devices[0], True)
18 
19 from doctr.documents import DocumentFile
20 from doctr.models import ocr_predictor
21 from doctr.utils.visualization import visualize_page
22 
23 DET_ARCHS = ["db_resnet50"]
24 RECO_ARCHS = ["crnn_vgg16_bn", "crnn_resnet31", "sar_vgg16_bn", "sar_resnet31"]
25 
26 
27 def main():
28 
29     # Wide mode
30     st.set_page_config(layout="wide")
31 
32     # Designing the interface
33     st.title("DocTR: Document Text Recognition")
34     # For newline
35     st.write('\n')
36     # Set the columns
37     cols = st.beta_columns((1, 1, 1))
38     cols[0].header("Input document")
39     cols[1].header("Text segmentation")
40     cols[-1].header("OCR output")
41 
42     # Sidebar
43     # File selection
44     st.sidebar.title("Document selection")
45     # Disabling warning
46     st.set_option('deprecation.showfileUploaderEncoding', False)
47     # Choose your own image
48     uploaded_file = st.sidebar.file_uploader("Upload files", type=['pdf', 'png', 'jpeg', 'jpg'])
49     if uploaded_file is not None:
50         if uploaded_file.name.endswith('.pdf'):
51             doc = DocumentFile.from_pdf(uploaded_file.read())
52         else:
53             doc = DocumentFile.from_images(uploaded_file.read())
54         cols[0].image(doc[0], "First page", use_column_width=True)
55 
56     # Model selection
57     st.sidebar.title("Model selection")
58     det_arch = st.sidebar.selectbox("Text detection model", DET_ARCHS)
59     reco_arch = st.sidebar.selectbox("Text recognition model", RECO_ARCHS)
60 
61     # For newline
62     st.sidebar.write('\n')
63 
64     if st.sidebar.button("Analyze document"):
65 
66         if uploaded_file is None:
67             st.sidebar.write("Please upload a document")
68 
69         else:
70             with st.spinner('Loading model...'):
71                 predictor = ocr_predictor(det_arch, reco_arch, pretrained=True)
72 
73             with st.spinner('Analyzing...'):
74 
75                 # Forward the image to the model
76                 processed_batches = predictor.det_predictor.pre_processor(doc)
77                 seg_map = predictor.det_predictor.model(processed_batches[0])[0]
78                 seg_map = cv2.resize(seg_map.numpy(), (doc[0].shape[1], doc[0].shape[0]),
79                                      interpolation=cv2.INTER_LINEAR)
80                 # Plot the raw heatmap
81                 fig, ax = plt.subplots()
82                 ax.imshow(seg_map)
83                 ax.axis('off')
84                 cols[1].pyplot(fig)
85 
86                 # OCR
87                 out = predictor(doc)
88                 fig = visualize_page(out.pages[0].export(), doc[0], interactive=False)
89                 cols[-1].pyplot(fig)
90 
91 
92 if __name__ == '__main__':
93     main()
94
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/demo/app.py b/demo/app.py
--- a/demo/app.py
+++ b/demo/app.py
@@ -74,7 +74,8 @@
 
                 # Forward the image to the model
                 processed_batches = predictor.det_predictor.pre_processor(doc)
-                seg_map = predictor.det_predictor.model(processed_batches[0])[0]
+                seg_map = predictor.det_predictor.model(processed_batches[0])["proba_map"]
+                seg_map = tf.squeeze(seg_map, axis=[0, 3])
                 seg_map = cv2.resize(seg_map.numpy(), (doc[0].shape[1], doc[0].shape[0]),
                                      interpolation=cv2.INTER_LINEAR)
                 # Plot the raw heatmap

{"golden_diff": "diff --git a/demo/app.py b/demo/app.py\n--- a/demo/app.py\n+++ b/demo/app.py\n@@ -74,7 +74,8 @@\n \n                 # Forward the image to the model\n                 processed_batches = predictor.det_predictor.pre_processor(doc)\n-                seg_map = predictor.det_predictor.model(processed_batches[0])[0]\n+                seg_map = predictor.det_predictor.model(processed_batches[0])[\"proba_map\"]\n+                seg_map = tf.squeeze(seg_map, axis=[0, 3])\n                 seg_map = cv2.resize(seg_map.numpy(), (doc[0].shape[1], doc[0].shape[0]),\n                                      interpolation=cv2.INTER_LINEAR)\n                 # Plot the raw heatmap\n", "issue": "Demo app error when analyzing my first document\n## \ud83d\udc1b Bug\r\n\r\nI tried to analyze a PNG and a PDF, got the same error. I try to change the model, didn't change anything.\r\n\r\n## To Reproduce\r\n\r\nSteps to reproduce the behavior:\r\n\r\n1. Upload a PNG\r\n2. Click on analyze document\r\n\r\n\r\n<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->\r\n```\r\nKeyError: 0\r\nTraceback:\r\nFile \"/Users/thibautmorla/opt/anaconda3/lib/python3.8/site-packages/streamlit/script_runner.py\", line 337, in _run_script\r\n    exec(code, module.__dict__)\r\nFile \"/Users/thibautmorla/Downloads/doctr/demo/app.py\", line 93, in <module>\r\n    main()\r\nFile \"/Users/thibautmorla/Downloads/doctr/demo/app.py\", line 77, in main\r\n    seg_map = predictor.det_predictor.model(processed_batches[0])[0]\r\n```\r\n\r\n\r\n## Additional context\r\n\r\nFirst image upload\n", "before_files": [{"content": "# Copyright (C) 2021, Mindee.\n\n# This program is licensed under the Apache License version 2.\n# See LICENSE or go to <https://www.apache.org/licenses/LICENSE-2.0.txt> for full license details.\n\nimport os\nimport streamlit as st\nimport matplotlib.pyplot as plt\n\nos.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"2\"\n\nimport tensorflow as tf\nimport cv2\n\ngpu_devices = tf.config.experimental.list_physical_devices('GPU')\nif any(gpu_devices):\n    tf.config.experimental.set_memory_growth(gpu_devices[0], True)\n\nfrom doctr.documents import DocumentFile\nfrom doctr.models import ocr_predictor\nfrom doctr.utils.visualization import visualize_page\n\nDET_ARCHS = [\"db_resnet50\"]\nRECO_ARCHS = [\"crnn_vgg16_bn\", \"crnn_resnet31\", \"sar_vgg16_bn\", \"sar_resnet31\"]\n\n\ndef main():\n\n    # Wide mode\n    st.set_page_config(layout=\"wide\")\n\n    # Designing the interface\n    st.title(\"DocTR: Document Text Recognition\")\n    # For newline\n    st.write('\\n')\n    # Set the columns\n    cols = st.beta_columns((1, 1, 1))\n    cols[0].header(\"Input document\")\n    cols[1].header(\"Text segmentation\")\n    cols[-1].header(\"OCR output\")\n\n    # Sidebar\n    # File selection\n    st.sidebar.title(\"Document selection\")\n    # Disabling warning\n    st.set_option('deprecation.showfileUploaderEncoding', False)\n    # Choose your own image\n    uploaded_file = st.sidebar.file_uploader(\"Upload files\", type=['pdf', 'png', 'jpeg', 'jpg'])\n    if uploaded_file is not None:\n        if uploaded_file.name.endswith('.pdf'):\n            doc = DocumentFile.from_pdf(uploaded_file.read())\n        else:\n            doc = DocumentFile.from_images(uploaded_file.read())\n        cols[0].image(doc[0], \"First page\", use_column_width=True)\n\n    # Model selection\n    st.sidebar.title(\"Model selection\")\n    det_arch = st.sidebar.selectbox(\"Text detection model\", DET_ARCHS)\n    reco_arch = st.sidebar.selectbox(\"Text recognition model\", RECO_ARCHS)\n\n    # For newline\n    st.sidebar.write('\\n')\n\n    if st.sidebar.button(\"Analyze document\"):\n\n        if uploaded_file is None:\n            st.sidebar.write(\"Please upload a document\")\n\n        else:\n            with st.spinner('Loading model...'):\n                predictor = ocr_predictor(det_arch, reco_arch, pretrained=True)\n\n            with st.spinner('Analyzing...'):\n\n                # Forward the image to the model\n                processed_batches = predictor.det_predictor.pre_processor(doc)\n                seg_map = predictor.det_predictor.model(processed_batches[0])[0]\n                seg_map = cv2.resize(seg_map.numpy(), (doc[0].shape[1], doc[0].shape[0]),\n                                     interpolation=cv2.INTER_LINEAR)\n                # Plot the raw heatmap\n                fig, ax = plt.subplots()\n                ax.imshow(seg_map)\n                ax.axis('off')\n                cols[1].pyplot(fig)\n\n                # OCR\n                out = predictor(doc)\n                fig = visualize_page(out.pages[0].export(), doc[0], interactive=False)\n                cols[-1].pyplot(fig)\n\n\nif __name__ == '__main__':\n    main()\n", "path": "demo/app.py"}], "after_files": [{"content": "# Copyright (C) 2021, Mindee.\n\n# This program is licensed under the Apache License version 2.\n# See LICENSE or go to <https://www.apache.org/licenses/LICENSE-2.0.txt> for full license details.\n\nimport os\nimport streamlit as st\nimport matplotlib.pyplot as plt\n\nos.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"2\"\n\nimport tensorflow as tf\nimport cv2\n\ngpu_devices = tf.config.experimental.list_physical_devices('GPU')\nif any(gpu_devices):\n    tf.config.experimental.set_memory_growth(gpu_devices[0], True)\n\nfrom doctr.documents import DocumentFile\nfrom doctr.models import ocr_predictor\nfrom doctr.utils.visualization import visualize_page\n\nDET_ARCHS = [\"db_resnet50\"]\nRECO_ARCHS = [\"crnn_vgg16_bn\", \"crnn_resnet31\", \"sar_vgg16_bn\", \"sar_resnet31\"]\n\n\ndef main():\n\n    # Wide mode\n    st.set_page_config(layout=\"wide\")\n\n    # Designing the interface\n    st.title(\"DocTR: Document Text Recognition\")\n    # For newline\n    st.write('\\n')\n    # Set the columns\n    cols = st.beta_columns((1, 1, 1))\n    cols[0].header(\"Input document\")\n    cols[1].header(\"Text segmentation\")\n    cols[-1].header(\"OCR output\")\n\n    # Sidebar\n    # File selection\n    st.sidebar.title(\"Document selection\")\n    # Disabling warning\n    st.set_option('deprecation.showfileUploaderEncoding', False)\n    # Choose your own image\n    uploaded_file = st.sidebar.file_uploader(\"Upload files\", type=['pdf', 'png', 'jpeg', 'jpg'])\n    if uploaded_file is not None:\n        if uploaded_file.name.endswith('.pdf'):\n            doc = DocumentFile.from_pdf(uploaded_file.read())\n        else:\n            doc = DocumentFile.from_images(uploaded_file.read())\n        cols[0].image(doc[0], \"First page\", use_column_width=True)\n\n    # Model selection\n    st.sidebar.title(\"Model selection\")\n    det_arch = st.sidebar.selectbox(\"Text detection model\", DET_ARCHS)\n    reco_arch = st.sidebar.selectbox(\"Text recognition model\", RECO_ARCHS)\n\n    # For newline\n    st.sidebar.write('\\n')\n\n    if st.sidebar.button(\"Analyze document\"):\n\n        if uploaded_file is None:\n            st.sidebar.write(\"Please upload a document\")\n\n        else:\n            with st.spinner('Loading model...'):\n                predictor = ocr_predictor(det_arch, reco_arch, pretrained=True)\n\n            with st.spinner('Analyzing...'):\n\n                # Forward the image to the model\n                processed_batches = predictor.det_predictor.pre_processor(doc)\n                seg_map = predictor.det_predictor.model(processed_batches[0])[\"proba_map\"]\n                seg_map = tf.squeeze(seg_map, axis=[0, 3])\n                seg_map = cv2.resize(seg_map.numpy(), (doc[0].shape[1], doc[0].shape[0]),\n                                     interpolation=cv2.INTER_LINEAR)\n                # Plot the raw heatmap\n                fig, ax = plt.subplots()\n                ax.imshow(seg_map)\n                ax.axis('off')\n                cols[1].pyplot(fig)\n\n                # OCR\n                out = predictor(doc)\n                fig = visualize_page(out.pages[0].export(), doc[0], interactive=False)\n                cols[-1].pyplot(fig)\n\n\nif __name__ == '__main__':\n    main()\n", "path": "demo/app.py"}]}