prompt: 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()
  
  ```
  
golden_diff: 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
  
verification_info: {"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"}]}