fg-mindee commited on
Commit
dd06629
1 Parent(s): 7459353

feat: Updated architecture selection

Browse files
Files changed (1) hide show
  1. app.py +20 -7
app.py CHANGED
@@ -1,16 +1,29 @@
1
- import os
2
- import cv2
 
 
 
 
 
 
3
  import streamlit as st
 
 
 
 
4
  import tensorflow as tf
5
- import matplotlib.pyplot as plt
 
 
 
6
 
7
  from doctr.io import DocumentFile
8
  from doctr.models import ocr_predictor
9
  from doctr.utils.visualization import visualize_page
10
 
 
 
11
 
12
- DET_ARCHS = ['db_resnet50', 'db_mobilenet_v3_large', 'linknet16']
13
- RECO_ARCHS = ["crnn_vgg16_bn", "master", "sar_resnet31"]
14
 
15
  def main():
16
 
@@ -18,10 +31,10 @@ def main():
18
  st.set_page_config(layout="wide")
19
 
20
  # Designing the interface
21
- st.title("DocTR: Document Text Recognition")
22
  # For newline
23
  st.write('\n')
24
- #
25
  st.write('Find more info at: https://github.com/mindee/doctr')
26
  # For newline
27
  st.write('\n')
 
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
+
8
+ import matplotlib.pyplot as plt
9
  import streamlit as st
10
+
11
+ os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
12
+
13
+ import cv2
14
  import tensorflow as tf
15
+
16
+ gpu_devices = tf.config.experimental.list_physical_devices('GPU')
17
+ if any(gpu_devices):
18
+ tf.config.experimental.set_memory_growth(gpu_devices[0], True)
19
 
20
  from doctr.io import DocumentFile
21
  from doctr.models import ocr_predictor
22
  from doctr.utils.visualization import visualize_page
23
 
24
+ DET_ARCHS = ["db_resnet50", "db_mobilenet_v3_large"]
25
+ RECO_ARCHS = ["crnn_vgg16_bn", "crnn_mobilenet_v3_small", "master", "sar_resnet31"]
26
 
 
 
27
 
28
  def main():
29
 
 
31
  st.set_page_config(layout="wide")
32
 
33
  # Designing the interface
34
+ st.title("docTR: Document Text Recognition")
35
  # For newline
36
  st.write('\n')
37
+ #
38
  st.write('Find more info at: https://github.com/mindee/doctr')
39
  # For newline
40
  st.write('\n')