OpenVINO(TM) Toolkit support

README.md

@@ -114,12 +114,39 @@ reader = easyocr.Reader(['ch_sim','en'], gpu=False)
 
 For more information, read the [tutorial](https://www.jaided.ai/easyocr/tutorial) and [API Documentation](https://www.jaided.ai/easyocr/documentation).
 
+#### Run with [OpenVINO(TM) Toolkit](https://github.com/openvinotoolkit/openvino) backend 
+To run EasyOCR with OpenVINO(TM) Toolkit on on Intel(R) CPU, Intel(R) Processor Graphics or Intel(R) Discrete Graphics, pass the inference device name as a Reader input argument. Choose between ‘ov_cpu’, ‘ov_gpu.<Intel GPU#>’ or 'ov_auto' (ov_+ device/plugin name as it's called in [OpenVINO](https://docs.openvino.ai/2024/about-openvino/compatibility-and-support/supported-devices.html))
+
+The following example sets up model inference to OpenVINO+Intel(R) CPU:
+
+```python
+reader = easyocr.Reader(['ch_sim','en'], 'ov_cpu')
+```
+Example of running inference on OpenVINO+Intel(R) Processor Graphics:
+
+```python
+reader = easyocr.Reader(['ch_sim','en'], 'ov_gpu')
+```
+Running inference on OpenVINO+Intel(R) Discrete Graphics:
+
+```python
+reader = easyocr.Reader(['ch_sim','en'], 'ov_gpu.1')
+```
+
 #### Run on command line
 
 ```shell
 $ easyocr -l ch_sim en -f chinese.jpg --detail=1 --gpu=True
 ```
 
+#### Run on command line with [OpenVINO(TM) Toolkit](https://github.com/openvinotoolkit/openvino) backend 
+
+Choose target device (--gpu) from ov_cpu, ov_gpu, ov_gpu.0, ov_gpu.1, ov_gpu.2,ov_gpu.3, ov_auto.
+
+```shell
+$ easyocr -l ch_sim en -f chinese.jpg --detail=1 --gpu=ov_cpu
+```
+
 ## Train/use your own model
 
 For recognition model, [Read here](https://github.com/JaidedAI/EasyOCR/blob/master/custom_model.md).

easyocr/cli.py

@@ -14,10 +14,9 @@ def parse_args():
     )
     parser.add_argument(
         "--gpu",
-        type=bool,
-        choices=[True, False],
+        choices=[True, False, 'ov_cpu', 'ov_gpu', 'ov_gpu.0', 'ov_gpu.1', 'ov_gpu.2', 'ov_gpu.3', 'ov_auto'],
         default=True,
-        help="Using GPU (default: True)",
+        help="Using GPU (default: True) or OpenVINO(TM) Toolkit on Intel(R) CPU, Intel(R) Processor Graphics or Intel(R) Discrete Graphics. Choose from True, False, ov_cpu, ov_gpu, ov_gpu.0, ov_gpu.1, ov_gpu.2, ov_gpu.3, ov_auto" ,
     )
     parser.add_argument(
         "--model_storage_directory",

easyocr/detection.py

@@ -10,6 +10,10 @@
 from .imgproc import resize_aspect_ratio, normalizeMeanVariance
 from .craft import CRAFT
 
+import openvino as ov
+import re
+import os
+
 def copyStateDict(state_dict):
     if list(state_dict.keys())[0].startswith("module"):
         start_idx = 1
@@ -39,11 +43,18 @@ def test_net(canvas_size, mag_ratio, net, image, text_threshold, link_threshold,
     x = [np.transpose(normalizeMeanVariance(n_img), (2, 0, 1))
          for n_img in img_resized_list]
     x = torch.from_numpy(np.array(x))
-    x = x.to(device)
+    if 'ov_' in device:
+        x=x.to('cpu')
+    else:
+        x = x.to(device)
 
     # forward pass
-    with torch.no_grad():
-        y, feature = net(x)
+    if 'ov_' in device:
+        res=net.infer_new_request({0: x})
+        y=torch.tensor(res[0])
+    else:
+        with torch.no_grad():
+            y, feature = net(x)
 
     boxes_list, polys_list = [], []
     for out in y:
@@ -81,12 +92,24 @@ def get_detector(trained_model, device='cpu', quantize=True, cudnn_benchmark=Fal
                 torch.quantization.quantize_dynamic(net, dtype=torch.qint8, inplace=True)
             except:
                 pass
+        net.eval()
+    elif 'ov_' in device:
+        ov_device=re.sub('ov_','',device).upper()
+        net.load_state_dict(copyStateDict(torch.load(trained_model, map_location='cpu')))
+        dummy_inp = torch.rand(1, 3, 608, 800)
+        net_ov = ov.convert_model(net, example_input=dummy_inp)
+        core = ov.Core()
+        if 'GPU' in ov_device:
+            cache_dir=os.path.expanduser('~/.EasyOCR/cache')
+            core.set_property({'CACHE_DIR': cache_dir})
+        net=core.compile_model(net_ov, ov_device)
+        print("Text detection model is running with OpenVINO on Intel", ov_device)
     else:
         net.load_state_dict(copyStateDict(torch.load(trained_model, map_location=device)))
         net = torch.nn.DataParallel(net).to(device)
         cudnn.benchmark = cudnn_benchmark
-
-    net.eval()
+        net.eval()
+
     return net
 
 def get_textbox(detector, image, canvas_size, mag_ratio, text_threshold, link_threshold, low_text, poly, device, optimal_num_chars=None, **kwargs):

easyocr/recognition.py

@@ -9,6 +9,9 @@
 import importlib
 from .utils import CTCLabelConverter
 import math
+import openvino as ov
+import os
+import re
 
 def custom_mean(x):
     return x.prod()**(2.0/np.sqrt(len(x)))
@@ -96,9 +99,21 @@ def __call__(self, batch):
         image_tensors = torch.cat([t.unsqueeze(0) for t in resized_images], 0)
         return image_tensors
 
+def copyStateDict(state_dict):
+    new_state_dict = OrderedDict()
+    for key, value in state_dict.items():
+        new_key = key[7:]
+        new_state_dict[new_key] = value
+    return new_state_dict
+
 def recognizer_predict(model, converter, test_loader, batch_max_length,\
                        ignore_idx, char_group_idx, decoder = 'greedy', beamWidth= 5, device = 'cpu'):
-    model.eval()
+    ov_device=''
+    if 'ov_' not in device:
+        model.eval()
+    else:
+        ov_device=device
+        device='cpu'
     result = []
     with torch.no_grad():
         for image_tensors in test_loader:
@@ -108,7 +123,12 @@ def recognizer_predict(model, converter, test_loader, batch_max_length,\
             length_for_pred = torch.IntTensor([batch_max_length] * batch_size).to(device)
             text_for_pred = torch.LongTensor(batch_size, batch_max_length + 1).fill_(0).to(device)
 
-            preds = model(image, text_for_pred)
+            if ov_device!='':
+                res = model.infer_new_request({0: image})
+                preds = next(iter(res.values()))
+                preds=torch.tensor(preds)
+            else:
+                preds = model(image, text_for_pred)
 
             # Select max probabilty (greedy decoding) then decode index to character
             preds_size = torch.IntTensor([preds.size(1)] * batch_size)
@@ -167,16 +187,26 @@ def get_recognizer(recog_network, network_params, character,\
 
     if device == 'cpu':
         state_dict = torch.load(model_path, map_location=device)
-        new_state_dict = OrderedDict()
-        for key, value in state_dict.items():
-            new_key = key[7:]
-            new_state_dict[new_key] = value
+        new_state_dict = copyStateDict(state_dict)
         model.load_state_dict(new_state_dict)
         if quantize:
             try:
                 torch.quantization.quantize_dynamic(model, dtype=torch.qint8, inplace=True)
             except:
                 pass
+    elif 'ov_' in device:
+        state_dict = torch.load(model_path, map_location="cpu")
+        new_state_dict = copyStateDict(state_dict)
+        model.load_state_dict(new_state_dict)
+        ov_device = re.sub('ov_','',device).upper()
+        core = ov.Core()
+        if 'GPU' in ov_device:
+            cache_dir = os.path.expanduser('~/.EasyOCR/cache')
+            core.set_property({'CACHE_DIR': cache_dir})
+        dummy_inp = torch.zeros(1, 1, 64, 320),torch.zeros(1,33)
+        model_ov = ov.convert_model(model,example_input=dummy_inp)
+        model = core.compile_model(model_ov, ov_device)
+        print('Text recognition model is running with OpenVINO on Intel ', ov_device)
     else:
         model = torch.nn.DataParallel(model).to(device)
         model.load_state_dict(torch.load(model_path, map_location=device))

requirements.txt

@@ -10,3 +10,4 @@ PyYAML
 Shapely
 pyclipper
 ninja
+openvino