chore: optimize code

README.md

@@ -10,6 +10,9 @@
 <a href=""><img src="https://img.shields.io/badge/OS-Linux%2C%20Win%2C%20Mac-pink.svg"></a>
 <a href="https://pypi.org/project/rapid-orientation/"><img alt="PyPI" src="https://img.shields.io/pypi/v/rapid-orientation"></a>
 <a href="https://pepy.tech/project/rapid-orientation"><img src="https://static.pepy.tech/personalized-badge/rapid-orientation?period=total&units=abbreviation&left_color=grey&right_color=blue&left_text=Downloads"></a>
+<a href="https://github.com/RapidAI/RapidOrientation/stargazers"><img src="https://img.shields.io/github/stars/RapidAI/RapidOrientation?color=ccf"></a>
+<a href="https://semver.org/"><img alt="SemVer2.0" src="https://img.shields.io/badge/SemVer-2.0-brightgreen"></a>
+<a href="https://github.com/psf/black"><img src="https://img.shields.io/badge/code%20style-black-000000.svg"></a>
 
 </div>
 

demo.py

@@ -5,7 +5,36 @@
 
 from rapid_orientation import RapidOrientation
 
+
+def scale_resize(img, resize_value=(280, 32)):
+    """
+    @params:
+    img: ndarray
+    resize_value: (width, height)
+    """
+    # padding
+    ratio = resize_value[0] / resize_value[1]  # w / h
+    h, w = img.shape[:2]
+    if w / h < ratio:
+        # 补宽
+        t = int(h * ratio)
+        w_padding = (t - w) // 2
+        img = cv2.copyMakeBorder(
+            img, 0, 0, w_padding, w_padding, cv2.BORDER_CONSTANT, value=(0, 0, 0)
+        )
+    else:
+        # 补高  (top, bottom, left, right)
+        t = int(w / ratio)
+        h_padding = (t - h) // 2
+        color = tuple([int(i) for i in img[0, 0, :]])
+        img = cv2.copyMakeBorder(
+            img, h_padding, h_padding, 0, 0, cv2.BORDER_CONSTANT, value=(0, 0, 0)
+        )
+    img = cv2.resize(img, resize_value, interpolation=cv2.INTER_LANCZOS4)
+    return img
+
+
 orientation_engine = RapidOrientation()
-img = cv2.imread("tests/test_files/img_rot180_demo.jpg")
+img = cv2.imread("tests/test_files/1.png")
 cls_result, _ = orientation_engine(img)
 print(cls_result)

rapid_orientation/config.yaml

@@ -6,11 +6,3 @@ CUDAExecutionProvider:
     arena_extend_strategy: kNextPowerOfTwo
     cudnn_conv_algo_search: EXHAUSTIVE
     do_copy_in_default_stream: true
-
-PreProcess:
-    - ResizeImage:
-        resize_short: 256
-    - CropImage:
-        size: 224
-    - NormalizeImage:
-    - ToCHWImage:

rapid_orientation/main.py

@@ -17,53 +17,52 @@
 import argparse
 import time
 from pathlib import Path
-from typing import Optional, Union
+from typing import Union
 
 import cv2
 import numpy as np
-import yaml
 
-from .utils import LoadImage, OrtInferSession, create_operators
+from .utils.infer_engine import OrtInferSession
+from .utils.load_image import LoadImage
+from .utils.preprocess import Preprocess
+from .utils.utils import read_yaml
 
 root_dir = Path(__file__).resolve().parent
+DEFAULT_PATH = root_dir / "models" / "rapid_orientation.onnx"
+DEFAULT_CFG = root_dir / "config.yaml"
 
 
 class RapidOrientation:
-    def __init__(self, model_path: Optional[str] = None):
-        config_path = str(root_dir / "config.yaml")
-        config = self.read_yaml(config_path)
-        if model_path is None:
-            model_path = str(root_dir / "models" / "rapid_orientation.onnx")
+    def __init__(
+        self,
+        model_path: Union[str, Path] = DEFAULT_PATH,
+        cfg_path: Union[str, Path] = DEFAULT_CFG,
+    ):
+        config = read_yaml(cfg_path)
         config["model_path"] = model_path
 
         self.session = OrtInferSession(config)
-        self.labels = self.session.get_metadata()["character"].splitlines()
-
-        self.preprocess_ops = create_operators(config["PreProcess"])
+        self.labels = self.session.get_character_list()
 
+        self.preprocess = Preprocess()
         self.load_img = LoadImage()
 
     def __call__(self, img_content: Union[str, np.ndarray, bytes, Path]):
-        images = self.load_img(img_content)
+        image = self.load_img(img_content)
+
+        s = time.perf_counter()
 
-        s = time.time()
-        for ops in self.preprocess_ops:
-            images = ops(images)
-        image = np.array(images)[None, ...]
+        image = self.preprocess(image)
+        image = image[None, ...]
 
         pred_output = self.session(image)[0]
 
         pred_output = pred_output.squeeze()
         pred_idx = np.argmax(pred_output)
         pred_txt = self.labels[pred_idx]
-        elapse = time.time() - s
-        return pred_txt, elapse
 
-    @staticmethod
-    def read_yaml(yaml_path):
-        with open(yaml_path, "rb") as f:
-            data = yaml.load(f, Loader=yaml.Loader)
-        return data
+        elapse = time.perf_counter() - s
+        return pred_txt, elapse
 
 
 def main():

rapid_orientation/utils/__init__.py

@@ -0,0 +1,3 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com