From 12ba81cd070718aa156b8571c927a286d25b2401 Mon Sep 17 00:00:00 2001 From: SWHL Date: Tue, 18 Jun 2024 23:01:54 +0800 Subject: [PATCH] refactor: Refactor code --- demo.py | 40 +-- docs/README_Layout.md | 39 ++- layout_res.png | Bin 0 -> 260175 bytes rapid_layout/__init__.py | 3 +- rapid_layout/{rapid_layout.py => main.py} | 92 ++++-- rapid_layout/utils/__init__.py | 18 ++ rapid_layout/utils/download_model.py | 60 ++++ rapid_layout/utils/infer_engine.py | 231 ++++++++++++++ rapid_layout/utils/load_image.py | 124 ++++++++ rapid_layout/utils/logger.py | 21 ++ .../{utils.py => utils/post_prepross.py} | 291 +----------------- rapid_layout/utils/pre_procss.py | 103 +++++++ rapid_layout/utils/vis_res.py | 146 +++++++++ setup_layout.py | 5 +- tests/test_layout.py | 3 +- 15 files changed, 811 insertions(+), 365 deletions(-) create mode 100644 layout_res.png rename rapid_layout/{rapid_layout.py => main.py} (52%) create mode 100644 rapid_layout/utils/__init__.py create mode 100644 rapid_layout/utils/download_model.py create mode 100644 rapid_layout/utils/infer_engine.py create mode 100644 rapid_layout/utils/load_image.py create mode 100644 rapid_layout/utils/logger.py rename rapid_layout/{utils.py => utils/post_prepross.py} (53%) create mode 100644 rapid_layout/utils/pre_procss.py create mode 100644 rapid_layout/utils/vis_res.py diff --git a/demo.py b/demo.py index d3b5e40..e063891 100644 --- a/demo.py +++ b/demo.py @@ -7,34 +7,11 @@ import cv2 import numpy as np -from rapid_layout import RapidLayout +from rapid_layout import RapidLayout, VisLayout from rapid_orientation import RapidOrientation from rapid_table import RapidTable, VisTable -def vis_layout(img: np.ndarray, layout_res: list) -> None: - tmp_img = copy.deepcopy(img) - for v in layout_res: - bbox = np.round(v["bbox"]).astype(np.int32) - label = v["label"] - - start_point = (bbox[0], bbox[1]) - end_point = (bbox[2], bbox[3]) - - cv2.rectangle(tmp_img, start_point, end_point, (0, 255, 0), 2) - cv2.putText( - tmp_img, label, start_point, cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2 - ) - - draw_img_save = Path("./inference_results/") - if not draw_img_save.exists(): - draw_img_save.mkdir(parents=True, exist_ok=True) - - image_save = str(draw_img_save / "layout_result.jpg") - cv2.imwrite(image_save, tmp_img) - print(f"The infer result has saved in {image_save}") - - def vis_table(table_res): style_res = """