"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
cfg.ir_optim=True
#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
self.text_detector=TextDetector(cfg)
defread_images(self,paths=[]):
images=[]
...
...
@@ -83,10 +66,9 @@ class OCRDet(hub.Module):
images.append(img)
returnimages
defdet_text(self,
defpredict(self,
images=[],
paths=[],
det_max_side_len=960,
draw_img_save='ocr_det_result',
visualization=False):
"""
...
...
@@ -94,10 +76,8 @@ class OCRDet(hub.Module):
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.
draw_img_save (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.
"""
...
...
@@ -111,8 +91,6 @@ class OCRDet(hub.Module):
assertpredicted_data!=[],"There is not any image to be predicted. Please check the input data."
"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/"
cfg.ir_optim=True
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
self.text_recognizer=TextRecognizer(cfg)
defread_images(self,paths=[]):
images=[]
...
...
@@ -81,7 +66,7 @@ class OCRRec(hub.Module):
images.append(img)
returnimages
defrec_text(self,
defpredict(self,
images=[],
paths=[]):
"""
...
...
@@ -102,14 +87,13 @@ class OCRRec(hub.Module):
assertpredicted_data!=[],"There is not any image to be predicted. Please check the input data."
"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
cfg.ir_optim=True
self.text_sys=TextSystem(cfg)
defread_images(self,paths=[]):
images=[]
...
...
@@ -99,10 +66,9 @@ class OCRSystem(hub.Module):
images.append(img)
returnimages
defrecognize_text(self,
defpredict(self,
images=[],
paths=[],
det_max_side_len=960,
draw_img_save='ocr_result',
visualization=False,
text_thresh=0.5):
...
...
@@ -111,11 +77,8 @@ class OCRSystem(hub.Module):
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.
draw_img_save (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.
...
...
@@ -130,8 +93,6 @@ class OCRSystem(hub.Module):
assertpredicted_data!=[],"There is not any image to be predicted. Please check the input data."
self.config.det_max_side_len=det_max_side_len
text_sys=TextSystem(self.config)
cnt=0
all_results=[]
forimginpredicted_data:
...
...
@@ -142,7 +103,7 @@ class OCRSystem(hub.Module):
all_results.append(result)
continue
starttime=time.time()
dt_boxes,rec_res=text_sys(img)
dt_boxes,rec_res=self.text_sys(img)
elapse=time.time()-starttime
cnt+=1
print("Predict time of image %d: %.3fs"%(cnt,elapse))