提交 d27360a9 编写于 作者: W WenmuZhou

move draw_result to utilitu.py

上级 dec76eb7
...@@ -25,27 +25,23 @@ from pathlib import Path ...@@ -25,27 +25,23 @@ from pathlib import Path
from ppocr.utils.logging import get_logger from ppocr.utils.logging import get_logger
from ppstructure.predict_system import OCRSystem, save_res from ppstructure.predict_system import OCRSystem, save_res
from utility import init_args from ppstructure.table.predict_table import to_excel
from ppstructure.utility import init_args, draw_result
logger = get_logger() logger = get_logger()
from ppocr.utils.utility import check_and_read_gif, get_image_file_list from ppocr.utils.utility import check_and_read_gif, get_image_file_list
from ppocr.utils.network import maybe_download, download_with_progressbar from ppocr.utils.network import maybe_download, download_with_progressbar
__all__ = ['PaddleStructure'] __all__ = ['PaddleStructure', 'draw_result', 'to_excel']
VERSION = '2.1' VERSION = '2.1'
BASE_DIR = os.path.expanduser("~/.paddlestructure/") BASE_DIR = os.path.expanduser("~/.paddlestructure/")
model_urls = { model_urls = {
'det': { 'det': 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_det_infer.tar',
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_det_infer.tar', 'rec': 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_rec_infer.tar',
}, 'structure': 'https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_structure_infer.tar'
'rec': {
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_det_infer.tar',
},
'structure': {
'https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_det_infer.tar',
},
} }
...@@ -143,21 +139,8 @@ def main(): ...@@ -143,21 +139,8 @@ def main():
img_name = os.path.basename(img_path).split('.')[0] img_name = os.path.basename(img_path).split('.')[0]
logger.info('{}{}{}'.format('*' * 10, img_path, '*' * 10)) logger.info('{}{}{}'.format('*' * 10, img_path, '*' * 10))
result = structure_engine(img_path) result = structure_engine(img_path)
save_res(result, args.output, os.path.basename(img_path).split('.')[0])
for item in result: for item in result:
logger.info(item['res']) logger.info(item['res'])
save_res(result, save_folder, img_name) save_res(result, save_folder, img_name)
logger.info('result save to {}'.format(os.path.join(save_folder, img_name))) logger.info('result save to {}'.format(os.path.join(save_folder, img_name)))
if __name__ == '__main__':
table_engine = PaddleStructure(
det_model_dir='/Users/zhoujun20/Desktop/工作相关/table/table_pr/PaddleOCR/inference/table/ch_ppocr_mobile_v2.0_table_det_infer',
rec_model_dir='/Users/zhoujun20/Desktop/工作相关/table/table_pr/PaddleOCR/inference/table/ch_ppocr_mobile_v2.0_table_rec_infer',
structure_model_dir='/Users/zhoujun20/Desktop/工作相关/table/table_pr/PaddleOCR/inference/table/ch_ppocr_mobile_v2.0_table_structure_infer',
output='/Users/zhoujun20/Desktop/工作相关/table/table_pr/PaddleOCR/output/table',
show_log=True)
img = cv2.imread('/Users/zhoujun20/Desktop/工作相关/table/table_pr/PaddleOCR/ppstructure/test_imgs/table_1.png')
result = table_engine(img)
for line in result:
print(line)
...@@ -31,7 +31,7 @@ from ppocr.utils.utility import get_image_file_list, check_and_read_gif ...@@ -31,7 +31,7 @@ from ppocr.utils.utility import get_image_file_list, check_and_read_gif
from ppocr.utils.logging import get_logger from ppocr.utils.logging import get_logger
from tools.infer.predict_system import TextSystem from tools.infer.predict_system import TextSystem
from ppstructure.table.predict_table import TableSystem, to_excel from ppstructure.table.predict_table import TableSystem, to_excel
from ppstructure.utility import parse_args from ppstructure.utility import parse_args,draw_result
logger = get_logger() logger = get_logger()
...@@ -39,7 +39,8 @@ logger = get_logger() ...@@ -39,7 +39,8 @@ logger = get_logger()
class OCRSystem(object): class OCRSystem(object):
def __init__(self, args): def __init__(self, args):
args.det_pad = True args.det_pad = True
args.det_pad_size = 640 args.det_pad_size = 960
args.drop_score = 0
self.text_system = TextSystem(args) self.text_system = TextSystem(args)
self.table_system = TableSystem(args, self.text_system.text_detector, self.text_system.text_recognizer) self.table_system = TableSystem(args, self.text_system.text_detector, self.text_system.text_recognizer)
self.table_layout = lp.PaddleDetectionLayoutModel("lp://PubLayNet/ppyolov2_r50vd_dcn_365e_publaynet/config", self.table_layout = lp.PaddleDetectionLayoutModel("lp://PubLayNet/ppyolov2_r50vd_dcn_365e_publaynet/config",
...@@ -107,6 +108,8 @@ def main(args): ...@@ -107,6 +108,8 @@ def main(args):
starttime = time.time() starttime = time.time()
res = structure_sys(img) res = structure_sys(img)
save_res(res, save_folder, img_name) save_res(res, save_folder, img_name)
draw_img = draw_result(img,res, args.vis_font_path)
cv2.imwrite(os.path.join(save_folder, img_name, 'show.jpg'), draw_img)
logger.info('result save to {}'.format(os.path.join(save_folder, img_name))) logger.info('result save to {}'.format(os.path.join(save_folder, img_name)))
elapse = time.time() - starttime elapse = time.time() - starttime
logger.info("Predict time : {:.3f}s".format(elapse)) logger.info("Predict time : {:.3f}s".format(elapse))
......
...@@ -11,9 +11,10 @@ ...@@ -11,9 +11,10 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import logging
from tools.infer.utility import str2bool, init_args as infer_args from PIL import Image
import numpy as np
from tools.infer.utility import draw_ocr_box_txt, init_args as infer_args
def init_args(): def init_args():
...@@ -38,3 +39,21 @@ def init_args(): ...@@ -38,3 +39,21 @@ def init_args():
def parse_args(): def parse_args():
parser = init_args() parser = init_args()
return parser.parse_args() return parser.parse_args()
def draw_result(image, result, font_path):
if isinstance(image, np.ndarray):
image = Image.fromarray(image)
boxes, txts, scores = [], [], []
for region in result:
if region['type'] == 'Table':
pass
elif region['type'] == 'Figure':
pass
else:
for box, rec_res in zip(region['res'][0], region['res'][1]):
boxes.append(np.array(box).reshape(-1, 2))
txts.append(rec_res[0])
scores.append(rec_res[1])
im_show = draw_ocr_box_txt(image, boxes, txts, scores, font_path=font_path,drop_score=0)
return im_show
\ No newline at end of file
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册