diff --git a/PPOCRLabel.py b/PPOCRLabel.py index cab409f..d85c4b4 100644 --- a/PPOCRLabel.py +++ b/PPOCRLabel.py @@ -134,6 +134,10 @@ def __init__( default_filename=None, default_predefined_class_file=None, default_save_dir=None, + det_model_dir=None, + rec_model_dir=None, + rec_char_dict_path=None, + cls_model_dir=None, ): super(MainWindow, self).__init__() self.setWindowTitle(__appname__) @@ -165,15 +169,27 @@ def getStr(strId): self.key_dialog_tip = getStr("keyDialogTip") self.defaultSaveDir = default_save_dir - self.ocr = PaddleOCR( - use_pdserving=False, - use_angle_cls=True, - det=True, - cls=True, - use_gpu=gpu, - lang=lang, - show_log=False, - ) + + params = { + "use_pdserving": False, + "use_angle_cls": True, + "det": True, + "cls": True, + "use_gpu": gpu, + "lang": lang, + "show_log": False, + } + + if det_model_dir is not None: + params["det_model_dir"] = det_model_dir + if rec_model_dir is not None: + params["rec_model_dir"] = rec_model_dir + if rec_char_dict_path is not None: + params["rec_char_dict_path"] = rec_char_dict_path + if cls_model_dir is not None: + params["cls_model_dir"] = cls_model_dir + + self.ocr = PaddleOCR(**params) self.table_ocr = PPStructure( use_pdserving=False, use_gpu=gpu, lang=lang, layout=False, show_log=False ) @@ -3528,6 +3544,11 @@ def get_main_app(argv=[]): ), nargs="?", ) + arg_parser.add_argument("--det_model_dir", type=str, default=None, nargs="?") + arg_parser.add_argument("--rec_model_dir", type=str, default=None, nargs="?") + arg_parser.add_argument("--rec_char_dict_path", type=str, default=None, nargs="?") + arg_parser.add_argument("--cls_model_dir", type=str, default=None, nargs="?") + args = arg_parser.parse_args(argv[1:]) win = MainWindow( @@ -3536,6 +3557,10 @@ def get_main_app(argv=[]): img_list_natural_sort=args.img_list_natural_sort, kie_mode=args.kie, default_predefined_class_file=args.predefined_classes_file, + det_model_dir=args.det_model_dir, + rec_model_dir=args.rec_model_dir, + rec_char_dict_path=args.rec_char_dict_path, + cls_model_dir=args.cls_model_dir, ) win.show() return app, win diff --git a/README.md b/README.md index 1b431ca..1aaa5cd 100644 --- a/README.md +++ b/README.md @@ -19,6 +19,11 @@ PPOCRLabelv2 is a semi-automatic graphic annotation tool suitable for OCR field, - 2024.09: - Added `Re-recognition` and `Auto Save Unsaved changes` features. For usage details, please refer to the "11. Additional Feature Description" in the "2.1 Operational Steps" section below. - Added the parameter `--img_list_natural_sort`, which defaults to natural sorting for the left image list. After configuring this parameter, character sorting will be used to easily locate images based on character order. + - Add 4 custom model parameters: + - `det_model_dir`: Path to the detection model directory + - `rec_model_dir`: Path to the recognition model directory + - `rec_char_dict_path`: Path to the recognition model dictionary file + - `cls_model_dir`: Path to the classification model directory - 2022.05: Add table annotations, follow `2.2 Table Annotations` for more information (by [whjdark](https://github.com/peterh0323); [Evezerest](https://github.com/Evezerest)) - 2022.02: (by [PeterH0323](https://github.com/peterh0323)) - Add KIE Mode by using `--kie`, for [detection + identification + keyword extraction] labeling. diff --git a/README_ch.md b/README_ch.md index f08230c..6638543 100644 --- a/README_ch.md +++ b/README_ch.md @@ -19,6 +19,11 @@ PPOCRLabel是一款适用于OCR领域的半自动化图形标注工具,内置P - 2024.09: - 新增`自动重新识别`和`自动保存未提交变更`功能,使用方法详见下方`2.1 操作步骤`的`11. 补充功能说明`。 - 新增`--img_list_natural_sort`参数,默认左侧图片列表使用自然排序,配置该参数后,将使用字符排序,方便根据字符顺序定位图片。 + - 新增4个自定义模型的参数: + - `det_model_dir` :检测模型目录路径 + - `rec_model_dir` :识别模型目录路径 + - `rec_char_dict_path` :识别模型字典文件路径 + - `cls_model_dir` :分类模型目录路径 - 2022.05:**新增表格标注**,使用方法见下方`2.2 表格标注`(by [whjdark](https://github.com/peterh0323); [Evezerest](https://github.com/Evezerest)) - 2022.02:**新增关键信息标注**、优化标注体验(by [PeterH0323](https://github.com/peterh0323) ) - 新增:使用 `--kie` 进入 KIE 功能,用于打【检测+识别+关键字提取】的标签