[F] Fix errors

Author: kunkunlin1221

.github/workflows/cuda-ci.yml

@@ -4,41 +4,80 @@ on:
   pull_request:
     branches: [main]
     paths-ignore:
+      - "**/version.h"
       - "doc/**"
       - "**.md"
 
+env:
+  LOCAL_REGISTRY: localhost:5000
+  DOCKERFILE: docker/pr.dockerfile
+
 jobs:
+  get_runner_and_uid:
+    name: Get Runner
+    runs-on: [self-hosted, unicorn]
+    steps:
+      - name: Get UID and GID
+        id: uid_gid
+        run: |
+          echo "uid_gid=$(id -u):$(id -g)" >> "$GITHUB_OUTPUT"
+    outputs:
+      runner: ${{ runner.name }}
+      uid: ${{ steps.uid_gid.outputs.uid_gid }}
+
+  build_docker_image:
+    name: Build Docker Image
+    needs: [get_runner_and_uid]
+    runs-on: ${{ needs.get_runner_and_uid.outputs.runner }}
+
+    steps:
+      - name: Checkout Repository
+        uses: actions/checkout@v4
+
+      - name: Build Docker Image
+        id: docker_build
+        uses: docker/build-push-action@v3
+        with:
+          file: ${{ env.DOCKERFILE }}
+          no-cache: false
+          push: true
+          tags: ${{ env.LOCAL_REGISTRY }}/capybara-container:ci
+
+    outputs:
+      image: ${{ env.LOCAL_REGISTRY }}/capybara-container:ci
+
   ci:
     name: CI
+    needs: [get_runner_and_uid, build_docker_image]
+    runs-on: ${{ needs.get_runner_and_uid.outputs.runner }}
     strategy:
-      fail-fast: false
       matrix:
-        python-version: ["3.10", "3.12"]
-
-    runs-on: ["self-hosted", "unicorn"]
+        python-version:
+          - "3.10"
+    container:
+      image: ${{ needs.build_docker_image.outputs.image }}
+      options: --user ${{ needs.get_runner_and_uid.outputs.uid }} --gpus all
 
     steps:
-      - name: Checkout repository
+      - name: Checkout Repository
         uses: actions/checkout@v4
-        with:
-          lfs: true
 
-      - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v5
-        with:
-          python-version: ${{ matrix.python-version }}
-          check-latest: true
+      - name: Install Dependencies
+        run: |
+          python3 -m pip install pytest wheel pylint pylint-flask pytest-cov typeguard
 
-      - name: Install dependencies
+      - name: Build and Install Package
         run: |
-          python -m pip install -U pip wheel "numpy>=2" "cython>=3.0.12" "setuptools>=69"
-          python setup.py build_ext --inplace
-          python -m pip install .
+          python3 setup.py bdist_wheel --universal && \
+          wheel_file=$(ls dist/*.whl 2>/dev/null || echo '') && \
+          if [ -z "$wheel_file" ]; then
+            echo 'Error: No wheel file found in dist directory.' && exit 1
+          fi && \
+          python3 -m pip install $wheel_file --force-reinstall
 
-      - name: Lint with pylint
+      - name: Lint with Pylint
         run: |
-          python -m pip install pylint
-          python -m pylint capybara --rcfile=.github/workflows/.pylintrc
+          python3 -m pylint capybara --rcfile=.github/workflows/.pylintrc --load-plugins pylint_flask
 
       - name: Run tests with pytest
         run: |

capybara/onnxengine/__init__.py

@@ -1,11 +1,8 @@
-from .engine import Backend, ONNXEngine
+from .engine import ONNXEngine
 from .engine_io_binding import ONNXEngineIOBinding
-from .metadata import (
-    get_onnx_metadata,
-    parse_metadata_from_onnx,
-    write_metadata_into_onnx,
-)
-from .tools import get_onnx_input_infos, get_onnx_output_infos, make_onnx_dynamic_axes
+from .enum import Backend
+from .metadata import get_onnx_metadata, parse_metadata_from_onnx, write_metadata_into_onnx
+from .tools import get_onnx_input_infos, get_onnx_output_infos, get_recommended_backend, make_onnx_dynamic_axes
 
 # 暫時無法使用
 # from .quantize import quantize, quantize_static

capybara/onnxengine/engine.py

@@ -1,22 +1,15 @@
-from enum import Enum
 from pathlib import Path
 from typing import Any, Dict, Union
 
 import colored
 import numpy as np
 import onnxruntime as ort
 
-from ..enums import EnumCheckMixin
+from .enum import Backend
 from .metadata import parse_metadata_from_onnx
 from .tools import get_onnx_input_infos, get_onnx_output_infos
 
 
-class Backend(EnumCheckMixin, Enum):
-    cpu = 0
-    cuda = 1
-    coreml = 2
-
-
 class ONNXEngine:
     def __init__(
         self,

capybara/onnxengine/engine_io_binding.py

@@ -62,9 +62,7 @@ def __init__(
 
         input_infos, output_infos = self._init_io_infos(model_path, input_initializer)
 
-        io_binding, x_ortvalues, y_ortvalues = self._setup_io_binding(
-            input_infos, output_infos
-        )
+        io_binding, x_ortvalues, y_ortvalues = self._setup_io_binding(input_infos, output_infos)
         self.io_binding = io_binding
         self.x_ortvalues = x_ortvalues
         self.y_ortvalues = y_ortvalues
@@ -121,14 +119,10 @@ def _setup_io_binding(self, input_infos, output_infos):
         y_ortvalues = {}
         for k, v in input_infos.items():
             m = np.zeros(**v)
-            x_ortvalues[k] = ort.OrtValue.ortvalue_from_numpy(
-                m, device_type="cuda", device_id=self.device_id
-            )
+            x_ortvalues[k] = ort.OrtValue.ortvalue_from_numpy(m, device_type="cuda", device_id=self.device_id)
         for k, v in output_infos.items():
             m = np.zeros(**v)
-            y_ortvalues[k] = ort.OrtValue.ortvalue_from_numpy(
-                m, device_type="cuda", device_id=self.device_id
-            )
+            y_ortvalues[k] = ort.OrtValue.ortvalue_from_numpy(m, device_type="cuda", device_id=self.device_id)
 
         io_binding = self.sess.io_binding()
         for k, v in x_ortvalues.items():
@@ -158,11 +152,7 @@ def format_nested_dict(dict_data, indent=0):
                 if isinstance(value, dict):
                     info.append(f"{prefix}{key}:")
                     info.append(format_nested_dict(value, indent + 1))
-                elif (
-                    isinstance(value, str)
-                    and value.startswith("{")
-                    and value.endswith("}")
-                ):
+                elif isinstance(value, str) and value.startswith("{") and value.endswith("}"):
                     try:
                         nested_dict = eval(value)
                         if isinstance(nested_dict, dict):
@@ -179,9 +169,7 @@ def format_nested_dict(dict_data, indent=0):
         title = "DOCSAID X ONNXRUNTIME"
         divider_length = 50
         divider = f"+{'-' * divider_length}+"
-        styled_title = colored.stylize(
-            title, [colored.fg("blue"), colored.attr("bold")]
-        )
+        styled_title = colored.stylize(title, [colored.fg("blue"), colored.attr("bold")])
 
         def center_text(text, width):
             """Center text within a fixed width, handling ANSI escape codes."""

capybara/onnxengine/enum.py

@@ -0,0 +1,9 @@
+from enum import Enum
+
+from ..enums import EnumCheckMixin
+
+
+class Backend(EnumCheckMixin, Enum):
+    cpu = 0
+    cuda = 1
+    coreml = 2

capybara/onnxengine/tools.py

@@ -2,13 +2,17 @@
 from typing import Dict, List, Optional, Union
 
 import onnx
+import onnxruntime as ort
 import onnxslim
 from onnx.helper import make_graph, make_model, make_opsetid, tensor_dtype_to_np_dtype
 
+from .enum import Backend
+
 __all__ = [
     "get_onnx_input_infos",
     "get_onnx_output_infos",
     "make_onnx_dynamic_axes",
+    "get_recommended_backend",
 ]
 
 
@@ -77,3 +81,13 @@ def make_onnx_dynamic_axes(
 
     new_model = onnxslim.slim(new_model)
     onnx.save(new_model, output_fpath)
+
+
+def get_recommended_backend() -> Backend:
+    providers = ort.get_available_providers()
+    if "CUDAExecutionProvider" in providers:
+        return Backend.cuda
+    elif "CoreMLExecutionProvider" in providers:
+        return Backend.coreml
+    else:
+        return Backend.cpu

tests/onnxruntime/test_engine.py

@@ -3,13 +3,12 @@
 import numpy as np
 import pytest
 
-from capybara import ONNXEngine, get_curdir
+from capybara import Backend, ONNXEngine, get_curdir, get_recommended_backend
 
 
-@pytest.mark.skipif(platform.system() != "Linux", reason="Linux only")
-def test_ONNXEngine_CUDA():
+def test_ONNXEngine_CPU():
     model_path = get_curdir(__file__).parent / "resources/model_dynamic-axes.onnx"
-    engine = ONNXEngine(model_path, backend="cuda")
+    engine = ONNXEngine(model_path, backend="cpu")
     for i in range(5):
         xs = {"input": np.random.randn(32, 3, 224, 224).astype("float32")}
         outs = engine(**xs)
@@ -18,9 +17,13 @@ def test_ONNXEngine_CUDA():
         prev_outs = outs
 
 
-def test_ONNXEngine_CPU():
+PLATFORM = platform.system()
+
+
+@pytest.mark.skipif(PLATFORM != "Linux" and BACKEND == Backend.cuda, reason="Linux with GPU only")
+def test_ONNXEngine_CUDA():
     model_path = get_curdir(__file__).parent / "resources/model_dynamic-axes.onnx"
-    engine = ONNXEngine(model_path, backend="cpu")
+    engine = ONNXEngine(model_path, backend=get_recommended_backend())
     for i in range(5):
         xs = {"input": np.random.randn(32, 3, 224, 224).astype("float32")}
         outs = engine(**xs)
@@ -29,7 +32,7 @@ def test_ONNXEngine_CPU():
         prev_outs = outs
 
 
-@pytest.mark.skipif(platform.system() != "Darwin", reason="Mac only")
+@pytest.mark.skipif(PLATFORM != "Darwin", reason="Mac only")
 def test_ONNXEngine_COREML():
     model_path = get_curdir(__file__).parent / "resources/model_dynamic-axes.onnx"
     engine = ONNXEngine(model_path, backend="coreml")

tests/onnxruntime/test_engine_io_binding.py

@@ -3,10 +3,12 @@
 import numpy as np
 import pytest
 
-from capybara import ONNXEngineIOBinding, get_curdir
+from capybara import Backend, ONNXEngineIOBinding, get_curdir, get_recommended_backend
 
+PLATFORM = platform.system()
 
-@pytest.mark.skipif(platform.system() != "Linux", reason="Linux only")
+
+@pytest.mark.skipif(PLATFORM != "Linux" and get_recommended_backend() != Backend.cuda, reason="Linux with GPU only")
 def test_ONNXEngineIOBinding_CUDAonly():
     model_path = get_curdir(__file__).parent / "resources/model_dynamic-axes.onnx"
     input_initializer = {"input": np.random.randn(32, 3, 448, 448).astype("float32")}