Преглед на файлове

fix(table): add null check for OCR result in rapid table prediction

- Add a null check for OCR result in the predict method
- Return None values if OCR result is None to prevent further processing
myhloli преди 1 година
родител
ревизия
18aa1a20bf
променени са 1 файла, в които са добавени 2 реда и са изтрити 0 реда
  1. 2 0
      magic_pdf/model/sub_modules/table/rapidtable/rapid_table.py

+ 2 - 0
magic_pdf/model/sub_modules/table/rapidtable/rapid_table.py

@@ -10,5 +10,7 @@ class RapidTableModel(object):
 
     def predict(self, image):
         ocr_result, _ = self.ocr_engine(np.asarray(image))
+        if ocr_result is None:
+            return None, None, None
         html_code, table_cell_bboxes, elapse = self.table_model(np.asarray(image), ocr_result)
         return html_code, table_cell_bboxes, elapse