Преглед изворни кода

Merge pull request #2088 from opendatalab/dev

fix: support non-pdf file in batch mode
Xiaomeng Zhao пре 7 месеци
родитељ
комит
dd96663c6b
1 измењених фајлова са 11 додато и 0 уклоњено
  1. 11 0
      magic_pdf/tools/cli.py

+ 11 - 0
magic_pdf/tools/cli.py

@@ -137,6 +137,17 @@ def cli(path, output_dir, method, lang, debug_able, start_page_id, end_page_id):
         doc_paths = []
         for doc_path in Path(path).glob('*'):
             if doc_path.suffix in pdf_suffixes + image_suffixes + ms_office_suffixes:
+                if doc_path.suffix in ms_office_suffixes:
+                    convert_file_to_pdf(str(doc_path), temp_dir)
+                    doc_path = Path(os.path.join(temp_dir, f'{doc_path.stem}.pdf'))
+                elif doc_path.suffix in image_suffixes:
+                    with open(str(doc_path), 'rb') as f:
+                        bits = f.read()
+                        pdf_bytes = fitz.open(stream=bits).convert_to_pdf()
+                    fn = os.path.join(temp_dir, f'{doc_path.stem}.pdf')
+                    with open(fn, 'wb') as f:
+                        f.write(pdf_bytes)
+                    doc_path = Path(fn)
                 doc_paths.append(doc_path)
         datasets = batch_build_dataset(doc_paths, 4, lang)
         batch_do_parse(output_dir, [str(doc_path.stem) for doc_path in doc_paths], datasets, method, debug_able, lang=lang)