disable gpu & windows

This commit is contained in:
weishengyu 2020-12-19 12:02:20 +08:00
parent cf03889b74
commit 1f9d6d7f31
1 changed files with 4 additions and 1 deletions

View File

@ -15,6 +15,7 @@ import numpy as np
import cv2 import cv2
import math import math
import paddle import paddle
import platform
from arch import style_text_rec from arch import style_text_rec
from utils.sys_funcs import check_gpu from utils.sys_funcs import check_gpu
@ -23,13 +24,15 @@ from utils.logging import get_logger
class StyleTextRecPredictor(object): class StyleTextRecPredictor(object):
def __init__(self, config): def __init__(self, config):
self.logger = get_logger()
algorithm = config['Predictor']['algorithm'] algorithm = config['Predictor']['algorithm']
assert algorithm in ["StyleTextRec" assert algorithm in ["StyleTextRec"
], "Generator {} not supported.".format(algorithm) ], "Generator {} not supported.".format(algorithm)
use_gpu = config["Global"]['use_gpu'] use_gpu = config["Global"]['use_gpu']
if use_gpu and paddle.is_compiled_with_cuda() and platform.system()=="Windows":
self.logger.error("GPU mode on Windows is not supported.")
check_gpu(use_gpu) check_gpu(use_gpu)
paddle.set_device('gpu' if use_gpu else 'cpu') paddle.set_device('gpu' if use_gpu else 'cpu')
self.logger = get_logger()
self.generator = getattr(style_text_rec, algorithm)(config) self.generator = getattr(style_text_rec, algorithm)(config)
self.height = config["Global"]["image_height"] self.height = config["Global"]["image_height"]
self.width = config["Global"]["image_width"] self.width = config["Global"]["image_width"]