From f44f43b452777e806cf1129ca78ccc5675d4ba86 Mon Sep 17 00:00:00 2001 From: avbelova Date: Fri, 29 Dec 2023 19:07:34 +0000 Subject: [PATCH 1/4] Add Intel OpenVINO support for Intel CPU, iGPU and dGPU --- easyocr/detection.py | 33 ++++++++++++++++++++++++++++----- easyocr/recognition.py | 27 +++++++++++++++++++++++++-- 2 files changed, 53 insertions(+), 7 deletions(-) diff --git a/easyocr/detection.py b/easyocr/detection.py index 072178a233..aa7f0748a5 100644 --- a/easyocr/detection.py +++ b/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): diff --git a/easyocr/recognition.py b/easyocr/recognition.py index 147370f7d0..c238ce2327 100644 --- a/easyocr/recognition.py +++ b/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))) @@ -98,7 +101,12 @@ def __call__(self, batch): 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 +116,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) @@ -177,6 +190,16 @@ def get_recognizer(recog_network, network_params, character,\ torch.quantization.quantize_dynamic(model, dtype=torch.qint8, inplace=True) except: pass + elif 'ov_' in device: + 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}) + ov_model_path=os.path.expanduser('~/.EasyOCR/openvino_model/1_recognition_model.onnx') + model_ov = core.read_model(ov_model_path) + 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)) From 6693c234cfe40ca7e43e16de37b4987a495ee77f Mon Sep 17 00:00:00 2001 From: avbelova Date: Tue, 26 Mar 2024 11:59:48 +0000 Subject: [PATCH 2/4] Add direct import of PyTorch recognition model --- easyocr/recognition.py | 23 +++++++++++++++-------- requirements.txt | 1 + 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/easyocr/recognition.py b/easyocr/recognition.py index c238ce2327..8437fdd8f1 100644 --- a/easyocr/recognition.py +++ b/easyocr/recognition.py @@ -99,6 +99,13 @@ 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'): ov_device='' @@ -180,10 +187,7 @@ 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: @@ -191,13 +195,16 @@ def get_recognizer(recog_network, network_params, character,\ except: pass elif 'ov_' in device: - ov_device=re.sub('ov_','',device).upper() + 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') + cache_dir = os.path.expanduser('~/.EasyOCR/cache') core.set_property({'CACHE_DIR': cache_dir}) - ov_model_path=os.path.expanduser('~/.EasyOCR/openvino_model/1_recognition_model.onnx') - model_ov = core.read_model(ov_model_path) + 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: diff --git a/requirements.txt b/requirements.txt index 9d0fb7b710..1973d1539d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,3 +10,4 @@ PyYAML Shapely pyclipper ninja +openvino From d94118679949f58e71a9281a866f21393458b8b4 Mon Sep 17 00:00:00 2001 From: avbelova Date: Tue, 26 Mar 2024 15:00:56 +0000 Subject: [PATCH 3/4] Add OpenVINO(TM) Toolkit support through CLI --- easyocr/cli.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/easyocr/cli.py b/easyocr/cli.py index e8520b8b1d..c825c757b4 100644 --- a/easyocr/cli.py +++ b/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", From 31273155f67c5fca246ee3c0632a8a50eff41ebf Mon Sep 17 00:00:00 2001 From: Anna Gubenkova Date: Tue, 26 Mar 2024 20:33:07 +0100 Subject: [PATCH 4/4] Update README.md Update readme --- README.md | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/README.md b/README.md index eab77bbdad..8eb5a968be 100644 --- a/README.md +++ b/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.’ 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).