提交 4402e629 编写于 作者: W WenmuZhou

修正export_model里的bug,添加predict_det

上级 89e031f0
......@@ -12,6 +12,13 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import sys
__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(__dir__)
sys.path.append(os.path.abspath(os.path.join(__dir__, '..')))
import argparse
import paddle
......@@ -20,14 +27,11 @@ from paddle.jit import to_static
from ppocr.modeling.architectures import build_model
from ppocr.postprocess import build_post_process
from ppocr.utils.save_load import init_model
from ppocr.utils.logging import get_logger
from tools.program import load_config
from tools.program import merge_config
def parse_args():
def str2bool(v):
return v.lower() in ("true", "t", "1")
parser = argparse.ArgumentParser()
parser.add_argument("-c", "--config", help="configuration file to use")
parser.add_argument(
......@@ -43,7 +47,7 @@ class Model(paddle.nn.Layer):
# Please modify the 'shape' according to actual needs
@to_static(input_spec=[
paddle.static.InputSpec(
shape=[None, 3, 32, None], dtype='float32')
shape=[None, 3, 640, 640], dtype='float32')
])
def forward(self, inputs):
x = self.pre_model(inputs)
......@@ -53,14 +57,13 @@ class Model(paddle.nn.Layer):
def main():
FLAGS = parse_args()
config = load_config(FLAGS.config)
merge_config(FLAGS.opt)
logger = get_logger()
# build post process
post_process_class = build_post_process(config['PostProcess'],
config['Global'])
# build model
#for rec algorithm
# for rec algorithm
if hasattr(post_process_class, 'character'):
char_num = len(getattr(post_process_class, 'character'))
config['Architecture']["Head"]['out_channels'] = char_num
......@@ -69,7 +72,10 @@ def main():
model.eval()
model = Model(model)
paddle.jit.save(model, FLAGS.output_path)
save_path = '{}/{}'.format(FLAGS.output_path,
config['Architecture']['model_type'])
paddle.jit.save(model, save_path)
logger.info('inference model is saved to {}'.format(save_path))
if __name__ == "__main__":
......
......@@ -22,7 +22,6 @@ import cv2
import numpy as np
import time
import sys
import paddle
import tools.infer.utility as utility
......@@ -39,7 +38,7 @@ class TextDetector(object):
postprocess_params = {}
if self.det_algorithm == "DB":
pre_process_list = [{
'ResizeForTest': {
'DetResizeForTest': {
'limit_side_len': args.det_limit_side_len,
'limit_type': args.det_limit_type
}
......@@ -53,7 +52,7 @@ class TextDetector(object):
}, {
'ToCHWImage': None
}, {
'keepKeys': {
'KeepKeys': {
'keep_keys': ['image', 'shape']
}
}]
......@@ -68,8 +67,9 @@ class TextDetector(object):
self.preprocess_op = create_operators(pre_process_list)
self.postprocess_op = build_post_process(postprocess_params)
self.predictor = paddle.jit.load(args.det_model_dir)
self.predictor.eval()
self.predictor, self.input_tensor, self.output_tensors = utility.create_predictor(
args, 'det', logger) # paddle.jit.load(args.det_model_dir)
# self.predictor.eval()
def order_points_clockwise(self, pts):
"""
......@@ -133,11 +133,23 @@ class TextDetector(object):
return None, 0
img = np.expand_dims(img, axis=0)
shape_list = np.expand_dims(shape_list, axis=0)
img = img.copy()
starttime = time.time()
preds = self.predictor(img)
if self.use_zero_copy_run:
self.input_tensor.copy_from_cpu(img)
self.predictor.zero_copy_run()
else:
im = paddle.fluid.core.PaddleTensor(img)
self.predictor.run([im])
outputs = []
for output_tensor in self.output_tensors:
output = output_tensor.copy_to_cpu()
outputs.append(output)
preds = outputs[0]
# preds = self.predictor(img)
post_result = self.postprocess_op(preds, shape_list)
dt_boxes = post_result[0]['points']
dt_boxes = self.filter_tag_det_res(dt_boxes, ori_im.shape)
elapse = time.time() - starttime
......@@ -146,8 +158,6 @@ class TextDetector(object):
if __name__ == "__main__":
args = utility.parse_args()
place = paddle.CPUPlace()
paddle.disable_static(place)
image_file_list = get_image_file_list(args.image_dir)
logger = get_logger()
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册