feat: test gpu exe

ziqinyeow · May 19, 2024 · d867358 · d867358
1 parent d5ab5eb
commit d867358
Show file tree

Hide file tree

Showing 5 changed files with 9 additions and 7 deletions.
diff --git a/.github/workflows/exe.yml b/.github/workflows/exe.yml
@@ -64,7 +64,7 @@ jobs:
            --add-binary="../onnxruntime-win-x64-gpu-1.17.1/lib/onnxruntime_providers_tensorrt.dll;./onnxruntime/capi/" \
             --add-binary="../onnxruntime-win-x64-gpu-1.17.1/lib/onnxruntime_providers_shared.dll;./onnxruntime/capi/" \
               --hidden-import=cv2 --hidden-import=supervision --hidden-import=addict --hidden-import=chex --hidden-import=lap --hidden-import=optax --hidden-import=einshape --hidden-import=haiku --hidden-import=mediapy \
-                --name sidecar-x86_64-pc-windows-msvc-gpu --specpath dist --distpath dist examples/fastapi-pyinstaller/server.py
+                --name sidecar-x86_64-pc-windows-msvc --specpath dist --distpath dist examples/fastapi-pyinstaller/server.py
 
       - name: Build Mac exe
         if: ${{ matrix.platform == 'macos-latest' }}

diff --git a/src/juxtapose/detectors/rtmdet/__init__.py b/src/juxtapose/detectors/rtmdet/__init__.py
@@ -45,7 +45,7 @@ def __init__(self, type: str = "m", device: str = "cpu", conf_thres: float = 0.3
         self.std = (57.375, 57.12, 58.395)
         self.device = device
         self.conf_thres = conf_thres
-        LOGGER.info(f"Loaded rtmdet-{type} onnx model.")
+        LOGGER.info(f"Loaded rtmdet-{type} onnx model into {providers}.")
 
     def inference(self, im: np.ndarray):
         im = im.transpose(2, 0, 1)

diff --git a/src/juxtapose/detectors/yolov8/__init__.py b/src/juxtapose/detectors/yolov8/__init__.py
@@ -42,10 +42,14 @@ def __init__(
         self.conf_thres = conf_thres
         self.iou_thres = iou_thres
 
+        providers = {"cpu": "CPUExecutionProvider", "cuda": "CUDAExecutionProvider"}[
+            device
+        ]
+
         # Create an inference session using the ONNX model and specify execution providers
         self.session = ort.InferenceSession(
             f"model/yolov8{type}.onnx",
-            providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
+            providers=[providers],
         )
 
         # Get the model inputs
@@ -56,7 +60,7 @@ def __init__(
         self.input_width = input_shape[2]
         self.input_height = input_shape[3]
 
-        LOGGER.info(f"Loaded yolov8{type} onnx model.")
+        LOGGER.info(f"Loaded yolov8{type} onnx model into {providers}.")
 
     def preprocess(self, img):
         img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

diff --git a/src/juxtapose/pose/rtmpose/__init__.py b/src/juxtapose/pose/rtmpose/__init__.py
@@ -41,7 +41,7 @@ def __init__(self, type: str = "m", device: str = "cpu") -> None:
         self.mean = (123.675, 116.28, 103.53)
         self.std = (58.395, 57.12, 57.375)
         self.device = device
-        LOGGER.info(f"Loaded rtmpose-{type} onnx model.")
+        LOGGER.info(f"Loaded rtmpose-{type} onnx model into {providers}.")
         # self.conf_thres = conf_thres
 
     def __call__(self, im: np.ndarray, bboxes: list = []):

diff --git a/src/juxtapose/rtm.py b/src/juxtapose/rtm.py
@@ -9,8 +9,6 @@
 from typing import List, Union, Generator, Literal
 import onnxruntime as ort
 
-import torch
-
 from juxtapose.data import load_inference_source
 from juxtapose.detectors import get_detector
 from juxtapose.pose.rtmpose import RTMPose