Skip to content

Commit aacfe58

Browse files
kunkunlin1221kunkunlin1221
authored
Merge pull request #27 from DocsaidLab/feat/update_onnxsim_to_onnxslim
[C] Replace `onnxsim` to `onnxslim` and rm `setup.cfg` [A] Add `video2frame_v2`
2 parents 0eaf29c + 8a30568 commit aacfe58

16 files changed

Lines changed: 365 additions & 136 deletions

File tree

‎.github/workflows/cpu-ci.yml‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
name: Pull Request for Multiple Platform with CPU
2+
3+
on:
4+
pull_request:
5+
branches: [main]
6+
paths-ignore:
7+
- "doc/**"
8+
- "**.md"
9+
10+
jobs:
11+
ci:
12+
name: CI
13+
strategy:
14+
fail-fast: false
15+
matrix:
16+
python-version: ["3.10", "3.12"]
17+
os: ["ubuntu-latest", "macos-13", "macos-latest"] # "ubuntu-24.04-arm" unsupported
18+
19+
runs-on: ${{ matrix.os }}
20+
21+
steps:
22+
- name: Checkout Repository
23+
uses: actions/checkout@v4
24+
with:
25+
lfs: true
26+
27+
- name: Install Prerequisites for Linux
28+
if: ${{ matrix.os == 'ubuntu-latest' || matrix.os == 'ubuntu-24.04-arm' }}
29+
run:
30+
sudo apt update && sudo apt install -y libturbojpeg exiftool ffmpeg libheif-dev poppler-utils
31+
32+
- name: Install Prerequisites for MacOS
33+
if: ${{ matrix.os == 'macos-13' || matrix.os == 'macos-latest' }}
34+
run:
35+
brew install libjpeg exiftool ffmpeg libheif poppler
36+
37+
- name: Set up Python ${{ matrix.python-version }}
38+
uses: actions/setup-python@v5
39+
with:
40+
python-version: ${{ matrix.python-version }}
41+
check-latest: true
42+
43+
- name: Install Dependencies
44+
run: |
45+
python -m pip install -U pip wheel "numpy>=2" "cython>=3.0.12" "setuptools>=69"
46+
python setup.py build_ext --inplace
47+
python -m pip install .
48+
49+
- name: Lint with Pylint
50+
run: |
51+
python -m pip install pylint
52+
python -m pylint capybara --rcfile=.github/workflows/.pylintrc
53+
54+
- name: Run Tests with Pytest
55+
run: |
56+
mkdir -p tests/coverage
57+
python -m pip install pytest pytest-cov typeguard
58+
python -m pytest tests --ignore tests/onnxruntime/test_engine_io_binding.py --junitxml=tests/coverage/cov-junitxml.xml --cov=capybara
59+
60+
- name: Surface failing tests
61+
uses: pmeier/pytest-results-action@main
62+
with:
63+
path: tests/coverage/cov-junitxml.xml
64+
summary: true
65+
display-options: fEX
66+
fail-on-empty: true
67+
title: Test results
68+
69+
- name: Clean artifacts
70+
run:
71+
rm -rf dist wheelhouse build *.egg-info
72+
73+
- name: Clean workspace when fail
74+
if: failure()
75+
run:
76+
rm -rf ${{ github.workspace }}

.github/workflows/pull_request.yml renamed to .github/workflows/cuda-ci.yml

Lines changed: 32 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
name: Pull Request
1+
name: Pull Request for Ubuntu with GPU
22

33
on:
44
pull_request:
@@ -53,22 +53,30 @@ jobs:
5353
strategy:
5454
matrix:
5555
python-version:
56-
- "3.10"
56+
- "3.10.16"
5757
container:
5858
image: ${{ needs.build_docker_image.outputs.image }}
5959
options: --user ${{ needs.get_runner_and_uid.outputs.uid }} --gpus all
6060

6161
steps:
6262
- name: Checkout Repository
6363
uses: actions/checkout@v4
64+
with:
65+
lfs: true
66+
67+
- name: Set up Python ${{ matrix.python-version }}
68+
uses: actions/setup-python@v5
69+
with:
70+
python-version: ${{ matrix.python-version }}
71+
check-latest: true
6472

6573
- name: Install Dependencies
6674
run: |
6775
python3 -m pip install pytest wheel pylint pylint-flask pytest-cov typeguard
6876
6977
- name: Build and Install Package
7078
run: |
71-
python3 setup.py bdist_wheel && \
79+
python3 setup.py bdist_wheel --universal && \
7280
wheel_file=$(ls dist/*.whl 2>/dev/null || echo '') && \
7381
if [ -z "$wheel_file" ]; then
7482
echo 'Error: No wheel file found in dist directory.' && exit 1
@@ -77,20 +85,28 @@ jobs:
7785
7886
- name: Lint with Pylint
7987
run: |
80-
python3 -m pylint capybara \
81-
--rcfile=.github/workflows/.pylintrc \
82-
--load-plugins pylint_flask
88+
python3 -m pylint capybara --rcfile=.github/workflows/.pylintrc --load-plugins pylint_flask
8389
84-
- name: Run Tests with Pytest
90+
- name: Run tests with pytest
8591
run: |
86-
mkdir -p tests/coverage && \
87-
python3 -m pytest tests --junitxml=tests/coverage/cov-junitxml.xml \
88-
--cov=capybara | tee tests/coverage/cov.txt
92+
mkdir -p tests/coverage
93+
python -m pip install pytest pytest-cov typeguard
94+
python -m pytest tests --junitxml=tests/coverage/cov-junitxml.xml --cov=capybara
8995
90-
- name: Pytest Coverage Comment
91-
id: coverageComment
92-
uses: MishaKav/pytest-coverage-comment@main
96+
- name: Surface failing tests
97+
uses: pmeier/pytest-results-action@main
9398
with:
94-
github-token: ${{ secrets.GITHUB_TOKEN }}
95-
pytest-coverage-path: tests/coverage/cov.txt
96-
junitxml-path: tests/coverage/cov-junitxml.xml
99+
path: tests/coverage/cov-junitxml.xml
100+
summary: true
101+
display-options: fEX
102+
fail-on-empty: true
103+
title: Test results
104+
105+
- name: Clean artifacts
106+
run:
107+
rm -rf dist wheelhouse build *.egg-info
108+
109+
- name: Clean workspace when fail
110+
if: failure()
111+
run:
112+
rm -rf ${{ github.workspace }}

‎capybara/onnxengine/__init__.py‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,8 @@
1-
from .engine import Backend, ONNXEngine
1+
from .engine import ONNXEngine
22
from .engine_io_binding import ONNXEngineIOBinding
3-
from .metadata import (
4-
get_onnx_metadata,
5-
parse_metadata_from_onnx,
6-
write_metadata_into_onnx,
7-
)
8-
from .tools import get_onnx_input_infos, get_onnx_output_infos, make_onnx_dynamic_axes
3+
from .enum import Backend
4+
from .metadata import get_onnx_metadata, parse_metadata_from_onnx, write_metadata_into_onnx
5+
from .tools import get_onnx_input_infos, get_onnx_output_infos, get_recommended_backend, make_onnx_dynamic_axes
96

107
# 暫時無法使用
118
# from .quantize import quantize, quantize_static

‎capybara/onnxengine/engine.py‎

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,15 @@
1-
from enum import Enum
21
from pathlib import Path
32
from typing import Any, Dict, Union
43

54
import colored
65
import numpy as np
76
import onnxruntime as ort
87

9-
from ..enums import EnumCheckMixin
8+
from .enum import Backend
109
from .metadata import parse_metadata_from_onnx
1110
from .tools import get_onnx_input_infos, get_onnx_output_infos
1211

1312

14-
class Backend(EnumCheckMixin, Enum):
15-
cpu = 0
16-
cuda = 1
17-
coreml = 2
18-
19-
2013
class ONNXEngine:
2114
def __init__(
2215
self,

‎capybara/onnxengine/engine_io_binding.py‎

Lines changed: 6 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ def __init__(
5353
providers=providers,
5454
provider_options=provider_options,
5555
)
56+
self.device = "cuda" if "CUDAExecutionProvider" in self.sess.get_providers() else "cpu"
5657

5758
# setting onnxruntime session info
5859
self.model_path = model_path
@@ -62,9 +63,7 @@ def __init__(
6263

6364
input_infos, output_infos = self._init_io_infos(model_path, input_initializer)
6465

65-
io_binding, x_ortvalues, y_ortvalues = self._setup_io_binding(
66-
input_infos, output_infos
67-
)
66+
io_binding, x_ortvalues, y_ortvalues = self._setup_io_binding(input_infos, output_infos)
6867
self.io_binding = io_binding
6968
self.x_ortvalues = x_ortvalues
7069
self.y_ortvalues = y_ortvalues
@@ -121,14 +120,10 @@ def _setup_io_binding(self, input_infos, output_infos):
121120
y_ortvalues = {}
122121
for k, v in input_infos.items():
123122
m = np.zeros(**v)
124-
x_ortvalues[k] = ort.OrtValue.ortvalue_from_numpy(
125-
m, device_type="cuda", device_id=self.device_id
126-
)
123+
x_ortvalues[k] = ort.OrtValue.ortvalue_from_numpy(m, device_type=self.device, device_id=self.device_id)
127124
for k, v in output_infos.items():
128125
m = np.zeros(**v)
129-
y_ortvalues[k] = ort.OrtValue.ortvalue_from_numpy(
130-
m, device_type="cuda", device_id=self.device_id
131-
)
126+
y_ortvalues[k] = ort.OrtValue.ortvalue_from_numpy(m, device_type=self.device, device_id=self.device_id)
132127

133128
io_binding = self.sess.io_binding()
134129
for k, v in x_ortvalues.items():
@@ -158,11 +153,7 @@ def format_nested_dict(dict_data, indent=0):
158153
if isinstance(value, dict):
159154
info.append(f"{prefix}{key}:")
160155
info.append(format_nested_dict(value, indent + 1))
161-
elif (
162-
isinstance(value, str)
163-
and value.startswith("{")
164-
and value.endswith("}")
165-
):
156+
elif isinstance(value, str) and value.startswith("{") and value.endswith("}"):
166157
try:
167158
nested_dict = eval(value)
168159
if isinstance(nested_dict, dict):
@@ -179,9 +170,7 @@ def format_nested_dict(dict_data, indent=0):
179170
title = "DOCSAID X ONNXRUNTIME"
180171
divider_length = 50
181172
divider = f"+{'-' * divider_length}+"
182-
styled_title = colored.stylize(
183-
title, [colored.fg("blue"), colored.attr("bold")]
184-
)
173+
styled_title = colored.stylize(title, [colored.fg("blue"), colored.attr("bold")])
185174

186175
def center_text(text, width):
187176
"""Center text within a fixed width, handling ANSI escape codes."""

‎capybara/onnxengine/enum.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
from enum import Enum
2+
3+
from ..enums import EnumCheckMixin
4+
5+
6+
class Backend(EnumCheckMixin, Enum):
7+
cpu = 0
8+
cuda = 1
9+
coreml = 2

‎capybara/onnxengine/tools.py‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,17 @@
22
from typing import Dict, List, Optional, Union
33

44
import onnx
5-
import onnxsim
5+
import onnxruntime as ort
6+
import onnxslim
67
from onnx.helper import make_graph, make_model, make_opsetid, tensor_dtype_to_np_dtype
78

9+
from .enum import Backend
10+
811
__all__ = [
912
"get_onnx_input_infos",
1013
"get_onnx_output_infos",
1114
"make_onnx_dynamic_axes",
15+
"get_recommended_backend",
1216
]
1317

1418

@@ -75,5 +79,16 @@ def make_onnx_dynamic_axes(
7579
if x.op_type == "Reshape":
7680
raise ValueError("Reshape cannot be trasformed to dynamic axes")
7781

78-
new_model, _ = onnxsim.simplify(new_model)
82+
new_model = onnxslim.slim(new_model)
7983
onnx.save(new_model, output_fpath)
84+
85+
86+
def get_recommended_backend() -> Backend:
87+
providers = ort.get_available_providers()
88+
device = ort.get_device()
89+
if "CUDAExecutionProvider" in providers and device == "GPU":
90+
return Backend.cuda
91+
elif "CoreMLExecutionProvider" in providers:
92+
return Backend.coreml
93+
else:
94+
return Backend.cpu

‎capybara/vision/videotools/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,2 @@
11
from .video2frames import *
2+
from .video2frames_v2 import *

0 commit comments

Comments
 (0)