"Environment Variable CUDA_VISIBLE_DEVICES is not set correctly. If you wanna use gpu, please set CUDA_VISIBLE_DEVICES via export CUDA_VISIBLE_DEVICES=cuda_device_id."
)
self.config.ir_optim=True
self.config.gpu_mem=8000
#params for text detector
self.config.det_algorithm=det_algorithm
self.config.det_model_dir=det_model_dir
# self.config.det_model_dir = "./inference/det/"
#DB parmas
self.config.det_db_thresh=0.3
self.config.det_db_box_thresh=0.5
self.config.det_db_unclip_ratio=2.0
#EAST parmas
self.config.det_east_score_thresh=0.8
self.config.det_east_cover_thresh=0.1
self.config.det_east_nms_thresh=0.2
defread_images(self,paths=[]):
images=[]
forimg_pathinpaths:
assertos.path.isfile(
img_path),"The {} isn't a valid file.".format(img_path)
img=cv2.imread(img_path)
ifimgisNone:
logger.info("error in loading image:{}".format(img_path))
continue
images.append(img)
returnimages
defdet_text(self,
images=[],
paths=[],
det_max_side_len=960,
draw_img_save='ocr_det_result',
visualization=False):
"""
Get the text box in the predicted images.
Args:
images (list(numpy.ndarray)): images data, shape of each is [H, W, C]. If images not paths
paths (list[str]): The paths of images. If paths not images
use_gpu (bool): Whether to use gpu. Default false.
output_dir (str): The directory to store output images.
visualization (bool): Whether to save image or not.
box_thresh(float): the threshold of the detected text box's confidence
Returns:
res (list): The result of text detection box and save path of images.
"Environment Variable CUDA_VISIBLE_DEVICES is not set correctly. If you wanna use gpu, please set CUDA_VISIBLE_DEVICES via export CUDA_VISIBLE_DEVICES=cuda_device_id."
)
self.config.ir_optim=True
self.config.gpu_mem=8000
#params for text recognizer
self.config.rec_algorithm=rec_algorithm
self.config.rec_model_dir=rec_model_dir
# self.config.rec_model_dir = "./inference/rec/"
self.config.rec_image_shape="3, 32, 320"
self.config.rec_char_type='ch'
self.config.rec_batch_num=rec_batch_num
self.config.rec_char_dict_path=rec_char_dict_path
self.config.use_space_char=True
defread_images(self,paths=[]):
images=[]
forimg_pathinpaths:
assertos.path.isfile(
img_path),"The {} isn't a valid file.".format(img_path)
img=cv2.imread(img_path)
ifimgisNone:
logger.info("error in loading image:{}".format(img_path))
continue
images.append(img)
returnimages
defrec_text(self,
images=[],
paths=[]):
"""
Get the text box in the predicted images.
Args:
images (list(numpy.ndarray)): images data, shape of each is [H, W, C]. If images not paths
paths (list[str]): The paths of images. If paths not images
Returns:
res (list): The result of text detection box and save path of images.
"Environment Variable CUDA_VISIBLE_DEVICES is not set correctly. If you wanna use gpu, please set CUDA_VISIBLE_DEVICES via export CUDA_VISIBLE_DEVICES=cuda_device_id."
)
self.config.ir_optim=True
self.config.gpu_mem=8000
#params for text detector
self.config.det_algorithm=det_algorithm
self.config.det_model_dir=det_model_dir
# self.config.det_model_dir = "./inference/det/"
#DB parmas
self.config.det_db_thresh=0.3
self.config.det_db_box_thresh=0.5
self.config.det_db_unclip_ratio=2.0
#EAST parmas
self.config.det_east_score_thresh=0.8
self.config.det_east_cover_thresh=0.1
self.config.det_east_nms_thresh=0.2
#params for text recognizer
self.config.rec_algorithm=rec_algorithm
self.config.rec_model_dir=rec_model_dir
# self.config.rec_model_dir = "./inference/rec/"
self.config.rec_image_shape="3, 32, 320"
self.config.rec_char_type='ch'
self.config.rec_batch_num=rec_batch_num
self.config.rec_char_dict_path=rec_char_dict_path
self.config.use_space_char=True
defread_images(self,paths=[]):
images=[]
forimg_pathinpaths:
assertos.path.isfile(
img_path),"The {} isn't a valid file.".format(img_path)
img=cv2.imread(img_path)
ifimgisNone:
logger.info("error in loading image:{}".format(img_path))
continue
images.append(img)
returnimages
defrecognize_text(self,
images=[],
paths=[],
det_max_side_len=960,
draw_img_save='ocr_result',
visualization=False,
text_thresh=0.5):
"""
Get the chinese texts in the predicted images.
Args:
images (list(numpy.ndarray)): images data, shape of each is [H, W, C]. If images not paths
paths (list[str]): The paths of images. If paths not images
use_gpu (bool): Whether to use gpu.
batch_size(int): the program deals once with one
output_dir (str): The directory to store output images.
visualization (bool): Whether to save image or not.
box_thresh(float): the threshold of the detected text box's confidence
text_thresh(float): the threshold of the recognize chinese texts' confidence
Returns:
res (list): The result of chinese texts and save path of images.