Skip to content

Commit

Permalink
add other devices in ppocr
Browse files Browse the repository at this point in the history
  • Loading branch information
YanhuiDua committed Jun 26, 2024
1 parent 4336771 commit 17c7f22
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 1 deletion.
5 changes: 4 additions & 1 deletion paddleocr.py
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ def _import_file(module_name, file_path, make_importable=False):
is_link,
confirm_model_dir_url,
)
from tools.infer.utility import draw_ocr, str2bool, check_gpu
from tools.infer.utility import draw_ocr, str2bool, check_gpu, check_xpu, check_npu, check_mlu
from ppstructure.utility import init_args, draw_structure_result
from ppstructure.predict_system import StructureSystem, save_structure_res, to_excel

Expand Down Expand Up @@ -619,6 +619,9 @@ def __init__(self, **kwargs):
SUPPORT_OCR_MODEL_VERSION, params.ocr_version
)
params.use_gpu = check_gpu(params.use_gpu)
params.use_xpu = check_xpu(params.use_xpu)
params.use_npu = check_npu(params.use_npu)
params.use_mlu = check_mlu(params.use_mlu)

if not params.show_log:
logger.setLevel(logging.INFO)
Expand Down
20 changes: 20 additions & 0 deletions tools/infer/utility.py
Original file line number Diff line number Diff line change
Expand Up @@ -805,6 +805,26 @@ def check_gpu(use_gpu):
use_gpu = False
return use_gpu

def check_xpu(use_xpu):
if use_xpu and (
not paddle.is_compiled_with_xpu() or paddle.device.get_device() == "cpu"
):
use_xpu = False
return use_xpu

def check_npu(use_npu):
if use_npu and (
not paddle.is_compiled_with_custom_device("npu") or paddle.device.get_device() == "cpu"
):
use_npu = False
return use_npu

def check_mlu(use_mlu):
if use_mlu and (
not paddle.is_compiled_with_custom_device("mlu") or paddle.device.get_device() == "cpu"
):
use_mlu = False
return use_mlu

if __name__ == "__main__":
pass

0 comments on commit 17c7f22

Please sign in to comment.