sberbank-ai commited on
Commit
62f9f91
·
1 Parent(s): 703c829

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -3
app.py CHANGED
@@ -1,7 +1,8 @@
1
  import os
2
- os.system("pip install git+https://github.com/ai-forever/ReadingPipeline")
3
 
4
  import cv2
 
5
  import json
6
 
7
  import gradio as gr
@@ -14,9 +15,9 @@ from ocrpipeline.linefinder import get_structured_text
14
  def get_config_and_download_weights(repo_id, device='cpu'):
15
  # download weights and configs
16
  pipeline_config_path = hf_hub_download(repo_id, "pipeline_config.json")
17
- ocr_model_path = hf_hub_download(repo_id, "ocr/ocr_model.ckpt")
18
  ocr_config_path = hf_hub_download(repo_id, "ocr/ocr_config.json")
19
- segm_model_path = hf_hub_download(repo_id, "segm/segm_model.ckpt")
20
  segm_config_path = hf_hub_download(repo_id, "segm/segm_config.json")
21
 
22
  # change paths to downloaded weights and configs in main pipeline_config
@@ -25,9 +26,11 @@ def get_config_and_download_weights(repo_id, device='cpu'):
25
  pipeline_config['main_process']['SegmPrediction']['model_path'] = segm_model_path
26
  pipeline_config['main_process']['SegmPrediction']['config_path'] = segm_config_path
27
  pipeline_config['main_process']['SegmPrediction']['device'] = device
 
28
  pipeline_config['main_process']['OCRPrediction']['model_path'] = ocr_model_path
29
  pipeline_config['main_process']['OCRPrediction']['config_path'] = ocr_config_path
30
  pipeline_config['main_process']['OCRPrediction']['device'] = device
 
31
 
32
  # save pipeline_config
33
  with open(pipeline_config_path, 'w') as f:
 
1
  import os
2
+ os.system("pip install git+https://github.com/ai-forever/ReadingPipeline.git@onnx")
3
 
4
  import cv2
5
+
6
  import json
7
 
8
  import gradio as gr
 
15
  def get_config_and_download_weights(repo_id, device='cpu'):
16
  # download weights and configs
17
  pipeline_config_path = hf_hub_download(repo_id, "pipeline_config.json")
18
+ ocr_model_path = hf_hub_download(repo_id, "ocr/ocr_model.onnx")
19
  ocr_config_path = hf_hub_download(repo_id, "ocr/ocr_config.json")
20
+ segm_model_path = hf_hub_download(repo_id, "segm/segm_model.onnx")
21
  segm_config_path = hf_hub_download(repo_id, "segm/segm_config.json")
22
 
23
  # change paths to downloaded weights and configs in main pipeline_config
 
26
  pipeline_config['main_process']['SegmPrediction']['model_path'] = segm_model_path
27
  pipeline_config['main_process']['SegmPrediction']['config_path'] = segm_config_path
28
  pipeline_config['main_process']['SegmPrediction']['device'] = device
29
+ pipeline_config['main_process']['SegmPrediction']['onnx'] = True
30
  pipeline_config['main_process']['OCRPrediction']['model_path'] = ocr_model_path
31
  pipeline_config['main_process']['OCRPrediction']['config_path'] = ocr_config_path
32
  pipeline_config['main_process']['OCRPrediction']['device'] = device
33
+ pipeline_config['main_process']['OCRPrediction']['onnx'] = True
34
 
35
  # save pipeline_config
36
  with open(pipeline_config_path, 'w') as f: