def __init__(self, args):
self.text_detector = TextDetectorHelper(args)
self.text_recognizer = TextRecognizerHelper(args)
self.use_angle_cls = args.use_angle_cls
if self.use_angle_cls:
self.clas_client = Debugger()
self.clas_client.load_model_config(
global_args.cls_model_dir, gpu=True, profile=False)
self.text_classifier = TextClassifierHelper(args)
self.det_client = Debugger()
self.det_client.load_model_config(
global_args.det_model_dir, gpu=True, profile=False)
self.fetch = ["ctc_greedy_decoder_0.tmp_0", "softmax_0.tmp_0"]
class TextSystemHelper(TextSystem):
def __init__(self, args):
self.text_detector = TextDetectorHelper(args)
self.text_recognizer = TextRecognizerHelper(args)
self.use_angle_cls = args.use_angle_cls
if self.use_angle_cls:
self.clas_client = Debugger()
self.clas_client.load_model_config(global_args.cls_server_dir,
gpu=True,
profile=False)
self.text_classifier = TextClassifierHelper(args)
self.det_client = Debugger()
self.det_client.load_model_config(global_args.det_server_dir,
gpu=True,
profile=False)
self.fetch = [
"save_infer_model/scale_0.tmp_0", "save_infer_model/scale_1.tmp_0"
]
def preprocess(self, img):
feed, fetch, self.tmp_args = self.text_detector.preprocess(img)
fetch_map = self.det_client.predict(feed, fetch)
outputs = [fetch_map[x] for x in fetch]
dt_boxes = self.text_detector.postprocess(outputs, self.tmp_args)
if dt_boxes is None:
return None, None
img_crop_list = []
dt_boxes = sorted_boxes(dt_boxes)
self.dt_boxes = dt_boxes
for bno in range(len(dt_boxes)):
tmp_box = copy.deepcopy(dt_boxes[bno])
img_crop = self.get_rotate_crop_image(img, tmp_box)
img_crop_list.append(img_crop)
if self.use_angle_cls:
feed, fetch, self.tmp_args = self.text_classifier.preprocess(
img_crop_list)
fetch_map = self.clas_client.predict(feed, fetch)
outputs = [fetch_map[x] for x in self.text_classifier.fetch]
for x in fetch_map.keys():
if ".lod" in x:
self.tmp_args[x] = fetch_map[x]
img_crop_list, _ = self.text_classifier.postprocess(
outputs, self.tmp_args)
feed, fetch, self.tmp_args = self.text_recognizer.preprocess(
img_crop_list)
return feed, self.fetch, self.tmp_args
def postprocess(self, outputs, args):
return self.text_recognizer.postprocess(outputs, args)
def __init__(self, args):
self.text_detector = TextDetectorHelper(args)
self.text_recognizer = TextRecognizerHelper(args)
self.use_angle_cls = args.use_angle_cls
if self.use_angle_cls:
self.clas_client = Debugger()
self.clas_client.load_model_config(global_args.cls_server_dir,
gpu=True,
profile=False)
self.text_classifier = TextClassifierHelper(args)
self.det_client = Debugger()
self.det_client.load_model_config(global_args.det_server_dir,
gpu=True,
profile=False)
self.fetch = [
"save_infer_model/scale_0.tmp_0", "save_infer_model/scale_1.tmp_0"
]