From 01430ed41e8f94d7ae572c424e0fb4bf62c948a3 Mon Sep 17 00:00:00 2001
From: Ning hao <135016091+ccddos@users.noreply.github.com>
Date: Sun, 12 May 2024 10:53:19 +0800
Subject: [PATCH] feat(cli): Support GPU acceleration via ONNX Runtime DirectML
 on Windows (#176)

* feat(cli): Support GPU acceleration via ONNX Runtime DirectML on Windows

This commit adds support for GPU acceleration in the CLI using ONNX Runtime DirectML. This allows users with compatible GPUs on Windows to leverage their hardware for faster inference.

* fix small bug
---
 python/rapidocr_onnxruntime/utils.py | 29 +++++++++++++++++++++++-----
 python/requirements_ort.txt          |  4 +++-
 2 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/python/rapidocr_onnxruntime/utils.py b/python/rapidocr_onnxruntime/utils.py
index a3b3870b9..79af3227e 100644
--- a/python/rapidocr_onnxruntime/utils.py
+++ b/python/rapidocr_onnxruntime/utils.py
@@ -57,18 +57,37 @@ def __init__(self, config):
         }
 
         EP_list = []
-        if (
-            config["use_cuda"]
-            and get_device() == "GPU"
-            and cuda_ep in get_available_providers()
-        ):
+        is_use_cude =  config["use_cuda"] and get_device() == "GPU" and cuda_ep in get_available_providers()
+        if (is_use_cude):
             EP_list = [(cuda_ep, cuda_provider_options)]
         EP_list.append((cpu_ep, cpu_provider_options))
 
+        # if platform is windows, use directml as primary provider
+        if os.name == "nt":
+            directml_ep = "DmlExecutionProvider"
+            # print (get_available_providers())
+            if directml_ep in get_available_providers():
+                print ("Windows platform detected, try to use DirectML as primary provider")
+                EP_list.insert(0, (directml_ep, 
+                    cuda_provider_options if is_use_cude else cpu_provider_options
+                ))
+        
+
         self._verify_model(config["model_path"])
         self.session = InferenceSession(
             config["model_path"], sess_options=sess_opt, providers=EP_list
         )
+        
+        # TODO: verify this is correct for detecting current_provider
+        current_provider = self.session.get_providers()[0]
+
+        # verify if the DirectML provider is used
+        if os.name == "nt":
+            if current_provider != directml_ep:
+                warnings.warn(
+                    f"DirectML is not available for the current environment, the inference part is automatically shifted to be executed under other EP.\n"
+                )
+
 
         if config["use_cuda"] and cuda_ep not in self.session.get_providers():
             warnings.warn(
diff --git a/python/requirements_ort.txt b/python/requirements_ort.txt
index 4111d1fe8..9b629440c 100644
--- a/python/requirements_ort.txt
+++ b/python/requirements_ort.txt
@@ -1,8 +1,10 @@
 pyclipper>=1.2.0
-onnxruntime>=1.7.0
 opencv_python>=4.5.1.48
 numpy>=1.19.5
 six>=1.15.0
 Shapely>=1.7.1
 PyYAML
 Pillow
+#  install the onnxruntime-directml if on windows platform, notice that the onnxruntime-directml is conflict with onnxruntime, we can only install one of them
+onnxruntime-directml;platform_system=="Windows"
+onnxruntime>=1.7.0;platform_system!="Windows"
\ No newline at end of file