(args)
| 302 | |
| 303 | |
| 304 | def main(args): |
| 305 | image_file_list = get_image_file_list(args.image_dir) |
| 306 | image_file_list = image_file_list |
| 307 | image_file_list = image_file_list[args.process_id :: args.total_process_num] |
| 308 | |
| 309 | if not args.use_pdf2docx_api: |
| 310 | structure_sys = StructureSystem(args) |
| 311 | save_folder = os.path.join(args.output, structure_sys.mode) |
| 312 | os.makedirs(save_folder, exist_ok=True) |
| 313 | img_num = len(image_file_list) |
| 314 | |
| 315 | for i, image_file in enumerate(image_file_list): |
| 316 | logger.info("[{}/{}] {}".format(i, img_num, image_file)) |
| 317 | img, flag_gif, flag_pdf = check_and_read(image_file) |
| 318 | img_name = os.path.basename(image_file).split(".")[0] |
| 319 | |
| 320 | if args.recovery and args.use_pdf2docx_api and flag_pdf: |
| 321 | try_import("pdf2docx") |
| 322 | from pdf2docx.converter import Converter |
| 323 | |
| 324 | os.makedirs(args.output, exist_ok=True) |
| 325 | docx_file = os.path.join(args.output, "{}_api.docx".format(img_name)) |
| 326 | cv = Converter(image_file) |
| 327 | cv.convert(docx_file) |
| 328 | cv.close() |
| 329 | logger.info("docx save to {}".format(docx_file)) |
| 330 | continue |
| 331 | |
| 332 | if not flag_gif and not flag_pdf: |
| 333 | img = cv2.imread(image_file) |
| 334 | |
| 335 | if not flag_pdf: |
| 336 | if img is None: |
| 337 | logger.error("error in loading image:{}".format(image_file)) |
| 338 | continue |
| 339 | imgs = [img] |
| 340 | else: |
| 341 | imgs = img |
| 342 | |
| 343 | all_res = [] |
| 344 | for index, img in enumerate(imgs): |
| 345 | res, time_dict = structure_sys(img, img_idx=index) |
| 346 | img_save_path = os.path.join( |
| 347 | save_folder, img_name, "show_{}.jpg".format(index) |
| 348 | ) |
| 349 | os.makedirs(os.path.join(save_folder, img_name), exist_ok=True) |
| 350 | if structure_sys.mode == "structure" and res != []: |
| 351 | draw_img = draw_structure_result(img, res, args.vis_font_path) |
| 352 | save_structure_res(res, save_folder, img_name, index) |
| 353 | elif structure_sys.mode == "kie": |
| 354 | if structure_sys.kie_predictor.predictor is not None: |
| 355 | draw_img = draw_re_results(img, res, font_path=args.vis_font_path) |
| 356 | else: |
| 357 | draw_img = draw_ser_results(img, res, font_path=args.vis_font_path) |
| 358 | |
| 359 | with open( |
| 360 | os.path.join(save_folder, img_name, "res_{}_kie.txt".format(index)), |
| 361 | "w", |
no test coverage detected
searching dependent graphs…