paddlex 3.0.0b2__py3-none-any.whl → 3.0.0rc1__py3-none-any.whl

paddlex/modules/formula_recognition/trainer.py

@@ -0,0 +1,120 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import shutil
+from pathlib import Path
+
+from ..base import BaseTrainer
+from .model_list import MODELS
+
+
+class FormulaRecTrainer(BaseTrainer):
+    """Text Recognition Model Trainer"""
+
+    entities = MODELS
+
+    def dump_label_dict(self, src_label_dict_path: str):
+        """dump label dict config
+
+        Args:
+            src_label_dict_path (str): path to label dict file to be saved.
+        """
+        dst_label_dict_path = Path(self.global_config.output).joinpath("label_dict.txt")
+        shutil.copyfile(src_label_dict_path, dst_label_dict_path)
+
+    def update_config(self):
+        """update training config"""
+        if self.train_config.log_interval:
+            self.pdx_config.update_log_interval(self.train_config.log_interval)
+        if self.train_config.eval_interval:
+            self.pdx_config._update_eval_interval_by_epoch(
+                self.train_config.eval_interval
+            )
+        if self.train_config.save_interval:
+            self.pdx_config.update_save_interval(self.train_config.save_interval)
+
+        if self.global_config["model"] == "LaTeX_OCR_rec":
+            self.pdx_config.update_dataset(
+                self.global_config.dataset_dir, "LaTeXOCRDataSet"
+            )
+        elif self.global_config["model"] in (
+            "UniMERNet",
+            "PP-FormulaNet-L",
+            "PP-FormulaNet-S",
+        ):
+            self.pdx_config.update_dataset(
+                self.global_config.dataset_dir, "SimpleDataSet"
+            )
+
+        label_dict_path = Path(self.global_config.dataset_dir).joinpath("dict.txt")
+        if label_dict_path.exists():
+            self.pdx_config.update_label_dict_path(label_dict_path)
+            self.dump_label_dict(label_dict_path)
+
+        if self.train_config.pretrain_weight_path:
+            self.pdx_config.update_pretrained_weights(
+                self.train_config.pretrain_weight_path
+            )
+
+        if self.train_config.batch_size is not None:
+            if self.global_config["model"] == "LaTeX_OCR_rec":
+                self.pdx_config.update_batch_size_pair(
+                    self.train_config.batch_size, mode="train"
+                )
+            else:
+                self.pdx_config.update_batch_size(
+                    self.train_config.batch_size, mode="train"
+                )
+
+        if self.eval_config.batch_size is not None:
+            if self.global_config["model"] == "LaTeX_OCR_rec":
+                self.pdx_config.update_batch_size_pair(
+                    self.eval_config.batch_size, mode="eval"
+                )
+            else:
+                self.pdx_config.update_batch_size(
+                    self.eval_config.batch_size, mode="eval"
+                )
+
+        if self.train_config.learning_rate is not None:
+            self.pdx_config.update_learning_rate(self.train_config.learning_rate)
+
+        if self.train_config.get("delimiter", None) is not None:
+            self.pdx_config.update_delimiter(self.train_config.delimiter, mode="train")
+
+        if self.eval_config.get("delimiter", None) is not None:
+            self.pdx_config.update_delimiter(self.eval_config.delimiter, mode="eval")
+
+        if self.train_config.epochs_iters is not None:
+            self.pdx_config._update_epochs(self.train_config.epochs_iters)
+        if (
+            self.train_config.resume_path is not None
+            and self.train_config.resume_path != ""
+        ):
+            self.pdx_config._update_checkpoints(self.train_config.resume_path)
+        if self.global_config.output is not None:
+            self.pdx_config._update_output_dir(self.global_config.output)
+
+    def get_train_kwargs(self) -> dict:
+        """get key-value arguments of model training function
+
+        Returns:
+            dict: the arguments of training function.
+        """
+        return {
+            "device": self.get_device(),
+            "dy2st": self.train_config.get("dy2st", False),
+            "amp": self.train_config.get("amp", "OFF"),  # amp support 'O1', 'O2', 'OFF'
+        }

paddlex/modules/general_recognition/__init__.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .trainer import ShiTuRecTrainer
 from .dataset_checker import ShiTuRecDatasetChecker
 from .evaluator import ShiTuRecEvaluator
 from .exportor import ShiTuRecExportor
+from .trainer import ShiTuRecTrainer

paddlex/modules/general_recognition/dataset_checker/__init__.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,8 +15,8 @@
 from pathlib import Path
 
 from ...base import BaseDatasetChecker
-from .dataset_src import check, convert, split_dataset, deep_analyse
 from ..model_list import MODELS
+from .dataset_src import check, convert, deep_analyse, split_dataset
 
 
 class ShiTuRecDatasetChecker(BaseDatasetChecker):

paddlex/modules/general_recognition/dataset_checker/dataset_src/__init__.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 
+from .analyse_dataset import deep_analyse
 from .check_dataset import check
 from .convert_dataset import convert
 from .split_dataset import split_dataset
-from .analyse_dataset import deep_analyse

paddlex/modules/general_recognition/dataset_checker/dataset_src/analyse_dataset.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,22 +13,20 @@
 # limitations under the License.
 
 import os
-import json
-import math
 import platform
-from pathlib import Path
 
-from collections import defaultdict
-from PIL import Image
 import numpy as np
-import matplotlib.pyplot as plt
-from matplotlib import font_manager
-from matplotlib.backends.backend_agg import FigureCanvasAgg
 
+from .....utils.deps import function_requires_deps, is_dep_available
 from .....utils.file_interface import custom_open
 from .....utils.fonts import PINGFANG_FONT_FILE_PATH
 
+if is_dep_available("matplotlib"):
+    import matplotlib.pyplot as plt
+    from matplotlib import font_manager
 
+
+@function_requires_deps("matplotlib")
 def deep_analyse(dataset_path, output, dataset_type="ShiTuRec"):
     """class analysis for dataset"""
     tags = ["train", "gallery", "query"]

paddlex/modules/general_recognition/dataset_checker/dataset_src/check_dataset.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,10 +15,11 @@
 import os
 import os.path as osp
 import random
-from PIL import Image, ImageOps
 from collections import defaultdict
 
-from .....utils.errors import DatasetFileNotFoundError, CheckFailedError
+from PIL import Image, ImageOps
+
+from .....utils.errors import CheckFailedError, DatasetFileNotFoundError
 from .utils.visualizer import draw_label
 
 
@@ -35,9 +36,7 @@ def check(dataset_dir, output, sample_num=10, dataset_type="ShiTuRec"):
     valid_num_parts = 2
 
     sample_cnts = dict()
-    label_map_dict = dict()
     sample_paths = defaultdict(list)
-    labels = []
 
     for tag in tags:
         file_list = osp.join(dataset_dir, f"{tag}.txt")

paddlex/modules/general_recognition/dataset_checker/dataset_src/convert_dataset.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,16 +13,17 @@
 # limitations under the License.
 
 
-import os
 import json
-from .....utils.file_interface import custom_open
+import os
+
 from .....utils.errors import ConvertFailedError
+from .....utils.file_interface import custom_open
 
 
 def check_src_dataset(root_dir, dataset_type):
     """check src dataset format validity"""
     if dataset_type in ("LabelMe"):
-        anno_suffix = ".json"
+        pass
     else:
         raise ConvertFailedError(
             message=f"数据格式转换失败！不支持{dataset_type}格式数据集。当前仅支持 LabelMe 格式。"
@@ -53,7 +54,7 @@ def convert(dataset_type, input_dir):
 
 
 def convert_labelme_dataset(root_dir):
-    image_dir = os.path.join(root_dir, "images")
+    os.path.join(root_dir, "images")
     anno_path = os.path.join(root_dir, "annotations")
     label_path = os.path.join(root_dir, "label.txt")
     train_rate = 50

paddlex/modules/general_recognition/dataset_checker/dataset_src/split_dataset.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

paddlex/modules/general_recognition/dataset_checker/dataset_src/utils/__init__.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

paddlex/modules/general_recognition/dataset_checker/dataset_src/utils/visualizer.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,12 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
 import numpy as np
-import json
-from pathlib import Path
 import PIL
-from PIL import Image, ImageDraw, ImageFont
+from PIL import ImageDraw, ImageFont
 
 from ......utils.fonts import PINGFANG_FONT_FILE_PATH
 

paddlex/modules/general_recognition/evaluator.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

paddlex/modules/general_recognition/exportor.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

paddlex/modules/general_recognition/model_list.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

paddlex/modules/general_recognition/trainer.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

paddlex/modules/image_classification/__init__.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .trainer import ClsTrainer
 from .dataset_checker import ClsDatasetChecker
 from .evaluator import ClsEvaluator
 from .exportor import ClsExportor
+from .trainer import ClsTrainer

paddlex/modules/image_classification/dataset_checker/__init__.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,8 +15,8 @@
 from pathlib import Path
 
 from ...base import BaseDatasetChecker
-from .dataset_src import check, convert, split_dataset, deep_analyse
 from ..model_list import MODELS
+from .dataset_src import check, convert, deep_analyse, split_dataset
 
 
 class ClsDatasetChecker(BaseDatasetChecker):

paddlex/modules/image_classification/dataset_checker/dataset_src/__init__.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 
+from .analyse_dataset import deep_analyse
 from .check_dataset import check
 from .convert_dataset import convert
 from .split_dataset import split_dataset
-from .analyse_dataset import deep_analyse

paddlex/modules/image_classification/dataset_checker/dataset_src/analyse_dataset.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,22 +13,21 @@
 # limitations under the License.
 
 import os
-import json
-import math
 import platform
-from pathlib import Path
-
 from collections import defaultdict
-from PIL import Image
+
 import numpy as np
-import matplotlib.pyplot as plt
-from matplotlib import font_manager
-from matplotlib.backends.backend_agg import FigureCanvasAgg
 
+from .....utils.deps import function_requires_deps, is_dep_available
 from .....utils.file_interface import custom_open
 from .....utils.fonts import PINGFANG_FONT_FILE_PATH
 
+if is_dep_available("matplotlib"):
+    import matplotlib.pyplot as plt
+    from matplotlib import font_manager
+
 
+@function_requires_deps("matplotlib")
 def deep_analyse(dataset_path, output):
     """class analysis for dataset"""
     tags = ["train", "val"]

paddlex/modules/image_classification/dataset_checker/dataset_src/check_dataset.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,10 +15,11 @@
 import os
 import os.path as osp
 import random
-from PIL import Image, ImageOps
 from collections import defaultdict
 
-from .....utils.errors import DatasetFileNotFoundError, CheckFailedError
+from PIL import Image, ImageOps
+
+from .....utils.errors import CheckFailedError, DatasetFileNotFoundError
 from .utils.visualizer import draw_label
 
 

paddlex/modules/image_classification/dataset_checker/dataset_src/convert_dataset.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import argparse
-import os
 import json
+import os
+
 from .....utils.file_interface import custom_open
 
 
@@ -34,7 +34,7 @@ def convert(input_dir):
     with custom_open(os.path.join(input_dir, "label.txt"), "w") as f:
         f.write("".join(label_content))
     anno_path = os.path.join(input_dir, "annotations")
-    jsons_path = os.listdir()
+    os.listdir()
     train_list = os.path.join(input_dir, "train.txt")
     val_list = os.path.join(input_dir, "val.txt")
     label_info = []

paddlex/modules/image_classification/dataset_checker/dataset_src/split_dataset.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

paddlex/modules/image_classification/dataset_checker/dataset_src/utils/__init__.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

paddlex/modules/image_classification/dataset_checker/dataset_src/utils/visualizer.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,12 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
 import numpy as np
-import json
-from pathlib import Path
 import PIL
-from PIL import Image, ImageDraw, ImageFont
+from PIL import ImageDraw, ImageFont
 
 from ......utils.fonts import PINGFANG_FONT_FILE_PATH
 

paddlex/modules/image_classification/evaluator.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

paddlex/modules/image_classification/exportor.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

paddlex/modules/image_classification/model_list.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -56,6 +56,7 @@ MODELS = [
     "PP-HGNetV2-B5",
     "PP-HGNetV2-B6",
     "PP-LCNet_x0_25",
+    "PP-LCNet_x0_25_textline_ori",
     "PP-LCNet_x0_35",
     "PP-LCNet_x0_5",
     "PP-LCNet_x0_75",
@@ -94,4 +95,5 @@ MODELS = [
     "FasterNet-T0",
     "FasterNet-T1",
     "FasterNet-T2",
+    "PP-LCNet_x1_0_table_cls",
 ]

paddlex/modules/image_classification/trainer.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,13 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import json
 import shutil
 from pathlib import Path
 
 from ..base import BaseTrainer
 from .model_list import MODELS
-from ...utils.config import AttrDict
 
 
 class ClsTrainer(BaseTrainer):
@@ -79,4 +77,6 @@ class ClsTrainer(BaseTrainer):
         ):
             train_args["resume_path"] = self.train_config.resume_path
         train_args["dy2st"] = self.train_config.get("dy2st", False)
+        # amp support 'O1', 'O2', 'OFF'
+        train_args["amp"] = self.train_config.get("amp", "OFF")
         return train_args

paddlex/modules/image_unwarping/__init__.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

paddlex/modules/image_unwarping/model_list.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

paddlex/modules/instance_segmentation/__init__.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,6 +13,6 @@
 # limitations under the License.
 
 from .dataset_checker import COCOInstSegDatasetChecker
-from .trainer import InstanceSegTrainer
 from .evaluator import InstanceSegEvaluator
 from .exportor import InstanceSegExportor
+from .trainer import InstanceSegTrainer

paddlex/modules/instance_segmentation/dataset_checker/__init__.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,11 +13,11 @@
 # limitations under the License.
 
 import os
+from pathlib import Path
 
-from .dataset_src import check, convert, split_dataset, deep_analyse
 from ...base import BaseDatasetChecker
-
 from ..model_list import MODELS
+from .dataset_src import check, convert, deep_analyse, split_dataset
 
 
 class COCOInstSegDatasetChecker(BaseDatasetChecker):
@@ -26,6 +26,20 @@ class COCOInstSegDatasetChecker(BaseDatasetChecker):
     entities = MODELS
     sample_num = 10
 
+    def get_dataset_root(self, dataset_dir: str) -> str:
+        """find the dataset root dir
+
+        Args:
+            dataset_dir (str): the directory that contain dataset.
+
+        Returns:
+            str: the root directory of dataset.
+        """
+        anno_dirs = list(Path(dataset_dir).glob("**/images"))
+        assert len(anno_dirs) == 1
+        dataset_dir = anno_dirs[0].parent.as_posix()
+        return dataset_dir
+
     def convert_dataset(self, src_dataset_dir: str) -> str:
         """convert the dataset from other type to specified type
 

paddlex/modules/instance_segmentation/dataset_checker/dataset_src/__init__.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 
+from .analyse_dataset import deep_analyse
 from .check_dataset import check
 from .convert_dataset import convert
 from .split_dataset import split_dataset
-from .analyse_dataset import deep_analyse