diff --git a/.github/workflows/cpu-ci.yml b/.github/workflows/cpu-ci.yml new file mode 100644 index 0000000..9cf5a9c --- /dev/null +++ b/.github/workflows/cpu-ci.yml @@ -0,0 +1,76 @@ +name: Pull Request for Multiple Platform with CPU + +on: + pull_request: + branches: [main] + paths-ignore: + - "doc/**" + - "**.md" + +jobs: + ci: + name: CI + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.12"] + os: ["ubuntu-latest", "macos-13", "macos-latest"] # "ubuntu-24.04-arm" unsupported + + runs-on: ${{ matrix.os }} + + steps: + - name: Checkout Repository + uses: actions/checkout@v4 + with: + lfs: true + + - name: Install Prerequisites for Linux + if: ${{ matrix.os == 'ubuntu-latest' || matrix.os == 'ubuntu-24.04-arm' }} + run: + sudo apt update && sudo apt install -y libturbojpeg exiftool ffmpeg libheif-dev poppler-utils + + - name: Install Prerequisites for MacOS + if: ${{ matrix.os == 'macos-13' || matrix.os == 'macos-latest' }} + run: + brew install libjpeg exiftool ffmpeg libheif poppler + + - 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: | + 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 . + + - name: Lint with Pylint + run: | + python -m pip install pylint + python -m pylint capybara --rcfile=.github/workflows/.pylintrc + + - name: Run Tests with Pytest + run: | + mkdir -p tests/coverage + python -m pip install pytest pytest-cov typeguard + python -m pytest tests --ignore tests/onnxruntime/test_engine_io_binding.py --junitxml=tests/coverage/cov-junitxml.xml --cov=capybara + + - name: Surface failing tests + uses: pmeier/pytest-results-action@main + with: + path: tests/coverage/cov-junitxml.xml + summary: true + display-options: fEX + fail-on-empty: true + title: Test results + + - name: Clean artifacts + run: + rm -rf dist wheelhouse build *.egg-info + + - name: Clean workspace when fail + if: failure() + run: + rm -rf ${{ github.workspace }} \ No newline at end of file diff --git a/.github/workflows/pull_request.yml b/.github/workflows/cuda-ci.yml similarity index 64% rename from .github/workflows/pull_request.yml rename to .github/workflows/cuda-ci.yml index a4724c7..78a8095 100644 --- a/.github/workflows/pull_request.yml +++ b/.github/workflows/cuda-ci.yml @@ -1,4 +1,4 @@ -name: Pull Request +name: Pull Request for Ubuntu with GPU on: pull_request: @@ -53,7 +53,7 @@ jobs: strategy: matrix: python-version: - - "3.10" + - "3.10.16" container: image: ${{ needs.build_docker_image.outputs.image }} options: --user ${{ needs.get_runner_and_uid.outputs.uid }} --gpus all @@ -61,6 +61,14 @@ jobs: steps: - 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: | @@ -68,7 +76,7 @@ jobs: - name: Build and Install Package run: | - python3 setup.py bdist_wheel && \ + 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 @@ -77,20 +85,28 @@ jobs: - name: Lint with Pylint run: | - python3 -m pylint capybara \ - --rcfile=.github/workflows/.pylintrc \ - --load-plugins pylint_flask + python3 -m pylint capybara --rcfile=.github/workflows/.pylintrc --load-plugins pylint_flask - - name: Run Tests with Pytest + - name: Run tests with pytest run: | - mkdir -p tests/coverage && \ - python3 -m pytest tests --junitxml=tests/coverage/cov-junitxml.xml \ - --cov=capybara | tee tests/coverage/cov.txt + mkdir -p tests/coverage + python -m pip install pytest pytest-cov typeguard + python -m pytest tests --junitxml=tests/coverage/cov-junitxml.xml --cov=capybara - - name: Pytest Coverage Comment - id: coverageComment - uses: MishaKav/pytest-coverage-comment@main + - name: Surface failing tests + uses: pmeier/pytest-results-action@main with: - github-token: ${{ secrets.GITHUB_TOKEN }} - pytest-coverage-path: tests/coverage/cov.txt - junitxml-path: tests/coverage/cov-junitxml.xml + path: tests/coverage/cov-junitxml.xml + summary: true + display-options: fEX + fail-on-empty: true + title: Test results + + - name: Clean artifacts + run: + rm -rf dist wheelhouse build *.egg-info + + - name: Clean workspace when fail + if: failure() + run: + rm -rf ${{ github.workspace }} \ No newline at end of file diff --git a/capybara/onnxengine/__init__.py b/capybara/onnxengine/__init__.py index 4d22999..18c1434 100644 --- a/capybara/onnxengine/__init__.py +++ b/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 diff --git a/capybara/onnxengine/engine.py b/capybara/onnxengine/engine.py index 46af592..bbed318 100644 --- a/capybara/onnxengine/engine.py +++ b/capybara/onnxengine/engine.py @@ -1,4 +1,3 @@ -from enum import Enum from pathlib import Path from typing import Any, Dict, Union @@ -6,17 +5,11 @@ 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, diff --git a/capybara/onnxengine/engine_io_binding.py b/capybara/onnxengine/engine_io_binding.py index ac4017d..cc8f90c 100644 --- a/capybara/onnxengine/engine_io_binding.py +++ b/capybara/onnxengine/engine_io_binding.py @@ -53,6 +53,7 @@ def __init__( providers=providers, provider_options=provider_options, ) + self.device = "cuda" if "CUDAExecutionProvider" in self.sess.get_providers() else "cpu" # setting onnxruntime session info self.model_path = model_path @@ -62,9 +63,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 +120,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=self.device, 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=self.device, device_id=self.device_id) io_binding = self.sess.io_binding() for k, v in x_ortvalues.items(): @@ -158,11 +153,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 +170,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.""" diff --git a/capybara/onnxengine/enum.py b/capybara/onnxengine/enum.py new file mode 100644 index 0000000..a697e12 --- /dev/null +++ b/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 diff --git a/capybara/onnxengine/tools.py b/capybara/onnxengine/tools.py index b0d71d7..73bffaa 100644 --- a/capybara/onnxengine/tools.py +++ b/capybara/onnxengine/tools.py @@ -2,13 +2,17 @@ from typing import Dict, List, Optional, Union import onnx -import onnxsim +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", ] @@ -75,5 +79,16 @@ def make_onnx_dynamic_axes( if x.op_type == "Reshape": raise ValueError("Reshape cannot be trasformed to dynamic axes") - new_model, _ = onnxsim.simplify(new_model) + new_model = onnxslim.slim(new_model) onnx.save(new_model, output_fpath) + + +def get_recommended_backend() -> Backend: + providers = ort.get_available_providers() + device = ort.get_device() + if "CUDAExecutionProvider" in providers and device == "GPU": + return Backend.cuda + elif "CoreMLExecutionProvider" in providers: + return Backend.coreml + else: + return Backend.cpu diff --git a/capybara/vision/videotools/__init__.py b/capybara/vision/videotools/__init__.py index 0b63f29..0872b94 100644 --- a/capybara/vision/videotools/__init__.py +++ b/capybara/vision/videotools/__init__.py @@ -1 +1,2 @@ from .video2frames import * +from .video2frames_v2 import * diff --git a/capybara/vision/videotools/video2frames_v2.py b/capybara/vision/videotools/video2frames_v2.py new file mode 100644 index 0000000..d50ae84 --- /dev/null +++ b/capybara/vision/videotools/video2frames_v2.py @@ -0,0 +1,167 @@ +from concurrent.futures import ThreadPoolExecutor, as_completed +from itertools import chain +from typing import Any, List + +import cv2 +import numpy as np + +from ..functionals import imcvtcolor, imresize +from .video2frames import is_video_file + +__all__ = ["video2frames_v2"] + + +def is_numpy_img(x: Any) -> bool: + """ + x == ndarray (H x W x C) + """ + return isinstance(x, np.ndarray) and (x.ndim == 2 or (x.ndim == 3 and x.shape[-1] in [1, 3])) + + +def flatten_list(xs: list) -> list: + """ + Function to flatten a list. + + Args: + l (List[List[...]]): + Nested lists that needs to be flattened. + + Returns: + flatten list (list): flatted list. + """ + out = list(chain(*xs)) + if len(out) and isinstance(out[0], list): + out = flatten_list(out) + return out + + +def get_step_inds(start: int, end: int, num: int): + if num > (end - start): + raise ValueError("num is larger than the number of total frames.") + return np.around(np.linspace(start=start, stop=end, num=num, endpoint=False)).astype(int).tolist() + + +def _extract_frames( + inds: List[int], video_path: str, max_size: int = 1920, color_base: str = "BGR", global_ind: int = 0 +): + # check video path + if not is_video_file(video_path): + raise TypeError(f"The video_path {video_path} is inappropriate.") + + # open cap + cap = cv2.VideoCapture(video_path) + # if start or end isn't specified lets assume 0 + start = inds[0] + end = inds[-1] + 1 + # 設定cap frame 的啟始點 + cap.set(1, start) + + def _process_frame(frame): + scale = max_size / max(frame.shape[:2]) + dst_h, dst_w = int(frame.shape[0] * scale), int(frame.shape[1] * scale) + if scale < 1: + frame = imresize(frame, (dst_h, dst_w)) + elif scale > 1: + frame = imresize(frame, (dst_h, dst_w), interpolation=cv2.INTER_AREA) + + if color_base.upper() != "BGR": + frame = imcvtcolor(frame, cvt_mode=f"BGR2{color_base}") + return frame + + def _pickup_frame(): + for idx in range(start, end): + _, frame = cap.read() + if idx == inds[0]: + inds.pop(0) + # skip error frame + if frame is None: + continue + yield _process_frame(frame) + + # extract frames + frames = list(_pickup_frame()) + + # release cap + cap.release() + return frames, global_ind + + +def video2frames_v2( + video_path: str, + frame_per_sec: int = None, + start_sec: float = 0, + end_sec: float = None, + n_threads: int = 8, + max_size: int = 1920, + color_base: str = "BGR", +) -> List[np.ndarray]: + """ + Extracts the frames from a video using ray + Inputs: + video_path (str): + path to the video. + frame_per_sec (int, Optional): + the number of extracting frames per sec. + If None, all frames will be extracted. + start_sec (int): + the start second for frame extraction. + end_sec (int): + the end second for frame extraction. + n_threads (int): + the number of threads. + max_size (int): + max resolution of extracted frames. + color_base (str): + RGB or BGR color. Defaults to 'BGR'. + Return: + frames (list) + [frame1, frame2, None, ...] or [frame1, frame2, ...,] else [] + """ + if not is_video_file(video_path): + raise TypeError(f"The video_path {video_path} is inappropriate.") + + # get total_frames frames of video + cap = cv2.VideoCapture(str(video_path)) + total_frames = round(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + fps = round(cap.get(cv2.CAP_PROP_FPS)) + cap.release() + + if total_frames == 0 or fps == 0: + return [] + + frame_per_sec = fps if frame_per_sec is None else frame_per_sec + total_sec = total_frames / fps + # get frame inds + end_sec = total_sec if end_sec is None or end_sec > total_sec else end_sec + if start_sec > end_sec: + raise ValueError(f"The start_sec should less than end_sec. {end_sec}") + + total_sec = end_sec - start_sec + start_frame = round(start_sec * fps) + end_frame = round(end_sec * fps) + num = round(total_sec * frame_per_sec) + frame_inds = get_step_inds(start_frame, end_frame, num) + + out_frames = [] + with ThreadPoolExecutor(max_workers=n_threads) as executor: + ## -----start process---- ## + # split the frames into chunk lists + divide_size = round(len(frame_inds) / n_threads) + 1 + frame_inds_list = [frame_inds[i * divide_size : (i + 1) * divide_size] for i in range(n_threads)] + future_to_frames = { + executor.submit(_extract_frames, inds, video_path, max_size, color_base, i): inds + for i, inds in enumerate(frame_inds_list) + } + out_frames = [[] for _ in range(n_threads)] + + for future in as_completed(future_to_frames): + frames = future_to_frames[future] + try: + frames, global_ind = future.result() + out_frames[global_ind] = frames + except Exception as e: + print(f"{frames} generated an exception: {e}") + + out_frames = flatten_list(out_frames) + + return out_frames diff --git a/docker/Dockerfile b/docker/Dockerfile index 43570b9..0ed3011 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,5 +1,5 @@ # syntax=docker/dockerfile:experimental -FROM nvidia/cuda:12.8.1-cudnn-runtime-ubuntu22.04 as builder +FROM nvidia/cuda:12.8.1-cudnn-runtime-ubuntu24.04 as builder ENV PYTHONDONTWRITEBYTECODE=1 \ DEBIAN_FRONTEND=noninteractive \ @@ -7,11 +7,11 @@ ENV PYTHONDONTWRITEBYTECODE=1 \ RUN apt-get update -y && apt-get upgrade -y && \ apt-get install -y --no-install-recommends \ - tzdata wget git libturbojpeg exiftool ffmpeg poppler-utils libpng-dev \ + tzdata wget git git-lfs libturbojpeg exiftool ffmpeg poppler-utils libpng-dev \ libtiff5-dev libjpeg8-dev libopenjp2-7-dev zlib1g-dev gcc \ libfreetype6-dev liblcms2-dev libwebp-dev tcl8.6-dev tk8.6-dev python3-tk \ python3-pip libharfbuzz-dev libfribidi-dev libxcb1-dev libfftw3-dev gosu \ - libpq-dev python3-dev && \ + libpq-dev python3-dev libc6 && \ ln -sf /usr/share/zoneinfo/$TZ /etc/localtime && \ dpkg-reconfigure -f noninteractive tzdata && \ apt-get clean && rm -rf /var/lib/apt/lists/* diff --git a/docker/pr.dockerfile b/docker/pr.dockerfile index 841c4cb..c74d632 100644 --- a/docker/pr.dockerfile +++ b/docker/pr.dockerfile @@ -1,5 +1,5 @@ # syntax=docker/dockerfile:experimental -FROM nvidia/cuda:12.8.1-cudnn-runtime-ubuntu22.04 as builder +FROM nvidia/cuda:12.8.1-cudnn-runtime-ubuntu24.04 as builder ENV PYTHONDONTWRITEBYTECODE=1 \ DEBIAN_FRONTEND=noninteractive \ @@ -7,11 +7,11 @@ ENV PYTHONDONTWRITEBYTECODE=1 \ RUN apt-get update -y && apt-get upgrade -y && \ apt-get install -y --no-install-recommends \ - tzdata wget git libturbojpeg exiftool ffmpeg poppler-utils libpng-dev \ + tzdata wget git git-lfs libturbojpeg exiftool ffmpeg poppler-utils libpng-dev \ libtiff5-dev libjpeg8-dev libopenjp2-7-dev zlib1g-dev gcc \ libfreetype6-dev liblcms2-dev libwebp-dev tcl8.6-dev tk8.6-dev python3-tk \ python3-pip libharfbuzz-dev libfribidi-dev libxcb1-dev libfftw3-dev gosu \ - libpq-dev python3-dev && \ + libpq-dev python3-dev libc6 && \ ln -sf /usr/share/zoneinfo/$TZ /etc/localtime && \ dpkg-reconfigure -f noninteractive tzdata && \ apt-get clean && rm -rf /var/lib/apt/lists/* diff --git a/pyproject.toml b/pyproject.toml index 7bcfd26..58ac44e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,12 +3,12 @@ requires = ["setuptools>=61", "wheel"] build-backend = "setuptools.build_meta" [project] -name = "capybara_docsaid" +name = "capybara-docsaid" dynamic = ["version"] description = "An Image Processing and Deep Learning Toolkit." readme = {file = "README.md", content-type = "text/markdown"} license = {text = "Apache License 2.0"} -requires-python = ">=3.10,<3.13" +requires-python = ">=3.10" classifiers = [ "Development Status :: 5 - Production/Stable", "License :: OSI Approved :: Apache Software License", @@ -16,7 +16,6 @@ classifiers = [ "Intended Audience :: Science/Research", "Operating System :: OS Independent", "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Topic :: Software Development :: Libraries", "Topic :: Software Development :: Libraries :: Python Modules" @@ -42,7 +41,7 @@ dependencies = [ "piexif", "matplotlib", "opencv-python>=4.12.0.88", - "onnxsim", + "onnxslim", "beautifulsoup4", "onnxruntime==1.22.0; platform_system == 'Darwin'", "onnxruntime_gpu==1.22.0; platform_system == 'Linux'", @@ -58,7 +57,8 @@ Issues = "https://github.com/DocsaidLab/Capybara/issues" include-package-data = true [tool.setuptools.packages.find] -exclude = ["docker", "demo", "tests", "benchmarks"] +include = ["capybara*"] +exclude = ["demo", "tests", "docs", ".github", "docker", "wheelhouse"] [tool.setuptools.dynamic] -version = {attr = "capybara.__version__"} \ No newline at end of file +version = { attr = "capybara.__version__" } \ No newline at end of file diff --git a/setup.cfg b/setup.cfg deleted file mode 100644 index 962694d..0000000 --- a/setup.cfg +++ /dev/null @@ -1,61 +0,0 @@ -[metadata] -name = capybara_docsaid -version = attr: capybara.__version__ -description = An Image Processing and Deep Learning Toolkit. -long_description = file: README.md -long_description_content_type = text/markdown -license = Apache License 2.0 -classifiers= - Development Status :: 5 - Production/Stable - License :: OSI Approved :: Apache Software License - Intended Audience :: Developers - Intended Audience :: Science/Research - Operating System :: OS Independent - Programming Language :: Python :: 3.10 - Programming Language :: Python :: 3.11 - Programming Language :: Python :: 3.12 - Topic :: Software Development :: Libraries - Topic :: Software Development :: Libraries :: Python Modules -python_requires = >=3.10,<3.13 -url = https://github.com/DocsaidLab/Capybara.git - -[options] -packages = find: -include_package_data = True -setup_requires= - pip - setuptools - wheel -install_requires = - dacite - psutil - requests - onnx - colored - numpy - pdf2image - ujson - pyyaml - tqdm - pybase64 - PyTurboJPEG - dill - networkx - natsort - flask - shapely - piexif - matplotlib - opencv-python>=4.12.0.88 - onnxruntime==1.22.0;platform_system=='Darwin' - onnxruntime_gpu==1.22.0;platform_system=='Linux' - onnxsim - beautifulsoup4 - pillow-heif - -[options.packages.find] -exclude = - docker - demo - tests - benchmarks \ No newline at end of file diff --git a/tests/onnxruntime/test_engine.py b/tests/onnxruntime/test_engine.py index 088dabc..67e114a 100644 --- a/tests/onnxruntime/test_engine.py +++ b/tests/onnxruntime/test_engine.py @@ -1,15 +1,12 @@ -import platform - 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 +15,10 @@ def test_ONNXEngine_CUDA(): prev_outs = outs -def test_ONNXEngine_CPU(): +@pytest.mark.skipif(get_recommended_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 +27,7 @@ def test_ONNXEngine_CPU(): prev_outs = outs -@pytest.mark.skipif(platform.system() != "Darwin", reason="Mac only") +@pytest.mark.skipif(get_recommended_backend() != "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") diff --git a/tests/onnxruntime/test_engine_io_binding.py b/tests/onnxruntime/test_engine_io_binding.py index d17a898..796d8a3 100644 --- a/tests/onnxruntime/test_engine_io_binding.py +++ b/tests/onnxruntime/test_engine_io_binding.py @@ -1,12 +1,10 @@ -import platform - import numpy as np import pytest -from capybara import ONNXEngineIOBinding, get_curdir +from capybara import Backend, ONNXEngineIOBinding, get_curdir, get_recommended_backend -@pytest.mark.skipif(platform.system() != "Linux", reason="Linux only") +@pytest.mark.skipif(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")} diff --git a/tests/vision/videotools/test_video2frames_v2.py b/tests/vision/videotools/test_video2frames_v2.py new file mode 100644 index 0000000..cbbfc58 --- /dev/null +++ b/tests/vision/videotools/test_video2frames_v2.py @@ -0,0 +1,31 @@ +import numpy as np +import pytest + +from capybara import get_curdir, video2frames_v2 + +# 測試用的影片 +video_path = get_curdir(__file__) / "video_test.mp4" + + +def test_video2frames_v2(): + # 測試從影片中提取所有幀 + frames = video2frames_v2(video_path) + assert isinstance(frames, list) + assert len(frames) > 0 + assert isinstance(frames[0], np.ndarray) + + +def test_video2frames_v2_with_fps(): + # 測試指定提取幀的速度 + frames = video2frames_v2(video_path, frame_per_sec=2, start_sec=0, end_sec=2, n_threads=2) + assert len(frames) == 4 + + +def test_video2frames_v2_invalid_input(): + # 測試不支援的影片類型 + with pytest.raises(TypeError): + video2frames_v2("invalid_video.txt") + + # 測試不存在的影片路徑 + with pytest.raises(TypeError): + video2frames_v2("non_existent_video.mp4")