فهرست منبع

Merge pull request #1462 from myhloli/dev

Dev
Xiaomeng Zhao 10 ماه پیش
والد
کامیت
14be573d8c
2فایلهای تغییر یافته به همراه20 افزوده شده و 12 حذف شده
  1. 18 10
      magic_pdf/model/sub_modules/table/rapidtable/rapid_table.py
  2. 2 2
      projects/gradio_app/app.py

+ 18 - 10
magic_pdf/model/sub_modules/table/rapidtable/rapid_table.py

@@ -8,17 +8,25 @@ from rapid_table import RapidTable
 class RapidTableModel(object):
 class RapidTableModel(object):
     def __init__(self, ocr_engine):
     def __init__(self, ocr_engine):
         self.table_model = RapidTable()
         self.table_model = RapidTable()
-        if ocr_engine is None:
-            self.ocr_model_name = "RapidOCR"
-            if torch.cuda.is_available():
-                from rapidocr_paddle import RapidOCR
-                self.ocr_engine = RapidOCR(det_use_cuda=True, cls_use_cuda=True, rec_use_cuda=True)
-            else:
-                from rapidocr_onnxruntime import RapidOCR
-                self.ocr_engine = RapidOCR()
+        # if ocr_engine is None:
+        #     self.ocr_model_name = "RapidOCR"
+        #     if torch.cuda.is_available():
+        #         from rapidocr_paddle import RapidOCR
+        #         self.ocr_engine = RapidOCR(det_use_cuda=True, cls_use_cuda=True, rec_use_cuda=True)
+        #     else:
+        #         from rapidocr_onnxruntime import RapidOCR
+        #         self.ocr_engine = RapidOCR()
+        # else:
+        #     self.ocr_model_name = "PaddleOCR"
+        #     self.ocr_engine = ocr_engine
+
+        self.ocr_model_name = "RapidOCR"
+        if torch.cuda.is_available():
+            from rapidocr_paddle import RapidOCR
+            self.ocr_engine = RapidOCR(det_use_cuda=True, cls_use_cuda=True, rec_use_cuda=True)
         else:
         else:
-            self.ocr_model_name = "PaddleOCR"
-            self.ocr_engine = ocr_engine
+            from rapidocr_onnxruntime import RapidOCR
+            self.ocr_engine = RapidOCR()
 
 
     def predict(self, image):
     def predict(self, image):
 
 

+ 2 - 2
projects/gradio_app/app.py

@@ -193,7 +193,7 @@ if __name__ == '__main__':
                 max_pages = gr.Slider(1, 20, 10, step=1, label='Max convert pages')
                 max_pages = gr.Slider(1, 20, 10, step=1, label='Max convert pages')
                 with gr.Row():
                 with gr.Row():
                     layout_mode = gr.Dropdown(['layoutlmv3', 'doclayout_yolo'], label='Layout model', value='doclayout_yolo')
                     layout_mode = gr.Dropdown(['layoutlmv3', 'doclayout_yolo'], label='Layout model', value='doclayout_yolo')
-                    language = gr.Dropdown(all_lang, label='Language', value='')
+                    language = gr.Dropdown(all_lang, label='Language', value='auto')
                 with gr.Row():
                 with gr.Row():
                     formula_enable = gr.Checkbox(label='Enable formula recognition', value=True)
                     formula_enable = gr.Checkbox(label='Enable formula recognition', value=True)
                     is_ocr = gr.Checkbox(label='Force enable OCR', value=False)
                     is_ocr = gr.Checkbox(label='Force enable OCR', value=False)
@@ -221,6 +221,6 @@ if __name__ == '__main__':
         file.change(fn=to_pdf, inputs=file, outputs=pdf_show)
         file.change(fn=to_pdf, inputs=file, outputs=pdf_show)
         change_bu.click(fn=to_markdown, inputs=[file, max_pages, is_ocr, layout_mode, formula_enable, table_enable, language],
         change_bu.click(fn=to_markdown, inputs=[file, max_pages, is_ocr, layout_mode, formula_enable, table_enable, language],
                         outputs=[md, md_text, output_file, pdf_show])
                         outputs=[md, md_text, output_file, pdf_show])
-        clear_bu.add([file, md, pdf_show, md_text, output_file, is_ocr, language])
+        clear_bu.add([file, md, pdf_show, md_text, output_file, is_ocr])
 
 
     demo.launch(server_name='0.0.0.0')
     demo.launch(server_name='0.0.0.0')