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

paddlex/modules/object_detection/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.
@@ -13,11 +13,8 @@
 # limitations under the License.
 
 
-from pathlib import Path
-
-from ..base import BaseTrainer
-from ...utils.config import AttrDict
 from ...utils import logging
+from ..base import BaseTrainer
 from .model_list import MODELS
 
 
@@ -28,7 +25,19 @@ class DetTrainer(BaseTrainer):
 
     def _update_dataset(self):
         """update dataset settings"""
-        self.pdx_config.update_dataset(self.global_config.dataset_dir, "COCODetDataset")
+        metric = self.pdx_config.metric if "metric" in self.pdx_config else "COCO"
+        data_fields = (
+            self.pdx_config.TrainDataset["data_fields"]
+            if "data_fields" in self.pdx_config.TrainDataset
+            else None
+        )
+
+        self.pdx_config.update_dataset(
+            self.global_config.dataset_dir,
+            "COCODetDataset",
+            data_fields=data_fields,
+            metric=metric,
+        )
 
     def update_config(self):
         """update training config"""
@@ -57,6 +66,8 @@ class DetTrainer(BaseTrainer):
             epochs_iters = self.train_config.epochs_iters
         else:
             epochs_iters = self.pdx_config.get_epochs_iters()
+        if self.train_config.warmup_steps is not None:
+            self.pdx_config.update_warmup_steps(self.train_config.warmup_steps)
         if self.global_config.output is not None:
             self.pdx_config.update_save_dir(self.global_config.output)
 
@@ -82,4 +93,6 @@ class DetTrainer(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/open_vocabulary_detection/__init__.py

@@ -0,0 +1,18 @@
+# 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.
+
+from .dataset_checker import OVDetDatasetChecker
+from .evaluator import OVDetEvaluator
+from .exportor import OVDetExportor
+from .trainer import OVDetTrainer

paddlex/modules/open_vocabulary_detection/dataset_checker.py

@@ -0,0 +1,29 @@
+# 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.
+
+from ...utils.errors import UnsupportedAPIError
+from ..base import BaseDatasetChecker
+from .model_list import MODELS
+
+
+class OVDetDatasetChecker(BaseDatasetChecker):
+    """Dataset Checker for Open Vocabulary Detection Model"""
+
+    entities = MODELS
+
+    def __init__(self, config):
+        # not support for now
+        raise UnsupportedAPIError(
+            "open vocabulary detection models do not support data check for now."
+        )

paddlex/modules/open_vocabulary_detection/evaluator.py

@@ -0,0 +1,29 @@
+# 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.
+
+from ...utils.errors import UnsupportedAPIError
+from ..base import BaseEvaluator
+from .model_list import MODELS
+
+
+class OVDetEvaluator(BaseEvaluator):
+    """Open Vocabulary Detection Model Evaluator"""
+
+    entities = MODELS
+
+    def __init__(self, config):
+        # not support for now
+        raise UnsupportedAPIError(
+            "open vocabulary detection models do not support evaluate for now."
+        )

paddlex/modules/open_vocabulary_detection/exportor.py

@@ -0,0 +1,29 @@
+# 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.
+
+from ...utils.errors import UnsupportedAPIError
+from ..base import BaseExportor
+from .model_list import MODELS
+
+
+class OVDetExportor(BaseExportor):
+    """Open Vocabulary Detection Model Exportor"""
+
+    entities = MODELS
+
+    def __init__(self, config):
+        # not support for now
+        raise UnsupportedAPIError(
+            "open vocabulary detection models do not support export for now."
+        )

paddlex/modules/open_vocabulary_detection/model_list.py

@@ -0,0 +1,16 @@
+# 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.
+
+
+MODELS = ["GroundingDINO-T", "YOLO-Worldv2-L"]

paddlex/modules/open_vocabulary_detection/trainer.py

@@ -0,0 +1,44 @@
+# 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.
+
+from ...utils.errors import UnsupportedAPIError
+from ..base import BaseTrainer
+from .model_list import MODELS
+
+
+class OVDetTrainer(BaseTrainer):
+    """Open Vocabulary Detection Model Trainer"""
+
+    entities = MODELS
+
+    def __init__(self, config):
+        # not support for now
+        raise UnsupportedAPIError(
+            "open vocabulary detection models do not support train for now."
+        )
+
+    def update_config(self):
+        """update training config"""
+
+    def get_train_kwargs(self) -> dict:
+        """get key-value arguments of model training function
+
+        Returns:
+            dict: the arguments of training function.
+        """
+        train_args = {
+            "device": self.get_device(),
+            "amp": self.train_config.get("amp", "OFF"),  # amp support 'O1', 'O2', 'OFF'
+        }
+        return train_args

paddlex/modules/open_vocabulary_segmentation/__init__.py

@@ -0,0 +1,18 @@
+# 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.
+
+from .dataset_checker import OVSegDatasetChecker
+from .evaluator import OVSegEvaluator
+from .exportor import OVSegExportor
+from .trainer import OVSegTrainer

paddlex/modules/open_vocabulary_segmentation/dataset_checker.py

@@ -0,0 +1,29 @@
+# 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.
+
+from ...utils.errors import UnsupportedAPIError
+from ..base import BaseDatasetChecker
+from .model_list import MODELS
+
+
+class OVSegDatasetChecker(BaseDatasetChecker):
+    """Dataset Checker for Open Vocabulary Segmentation Model"""
+
+    entities = MODELS
+
+    def __init__(self, config):
+        # not support for now
+        raise UnsupportedAPIError(
+            "open vocabulary segmentation models do not support data check for now."
+        )

paddlex/modules/open_vocabulary_segmentation/evaluator.py

@@ -0,0 +1,29 @@
+# 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.
+
+from ...utils.errors import UnsupportedAPIError
+from ..base import BaseEvaluator
+from .model_list import MODELS
+
+
+class OVSegEvaluator(BaseEvaluator):
+    """Open Vocabulary Segmentation Model Evaluator"""
+
+    entities = MODELS
+
+    def __init__(self, config):
+        # not support for now
+        raise UnsupportedAPIError(
+            "open vocabulary segmentation models do not support evaluate for now."
+        )

paddlex/modules/open_vocabulary_segmentation/exportor.py

@@ -0,0 +1,29 @@
+# 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.
+
+from ...utils.errors import UnsupportedAPIError
+from ..base import BaseExportor
+from .model_list import MODELS
+
+
+class OVSegExportor(BaseExportor):
+    """Open Vocabulary Segmentation Model Exportor"""
+
+    entities = MODELS
+
+    def __init__(self, config):
+        # not support for now
+        raise UnsupportedAPIError(
+            "open vocabulary segmentation models do not support export for now."
+        )

paddlex/modules/open_vocabulary_segmentation/model_list.py

@@ -0,0 +1,19 @@
+# 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.
+
+
+MODELS = [
+    "SAM-H_point",
+    "SAM-H_box",
+]

paddlex/modules/open_vocabulary_segmentation/trainer.py

@@ -0,0 +1,44 @@
+# 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.
+
+from ...utils.errors import UnsupportedAPIError
+from ..base import BaseTrainer
+from .model_list import MODELS
+
+
+class OVSegTrainer(BaseTrainer):
+    """Open Vocabulary Segmentation Model Trainer"""
+
+    entities = MODELS
+
+    def __init__(self, config):
+        # not support for now
+        raise UnsupportedAPIError(
+            "open vocabulary segmentation models do not support train for now."
+        )
+
+    def update_config(self):
+        """update training config"""
+
+    def get_train_kwargs(self) -> dict:
+        """get key-value arguments of model training function
+
+        Returns:
+            dict: the arguments of training function.
+        """
+        train_args = {
+            "device": self.get_device(),
+            "amp": self.train_config.get("amp", "OFF"),  # amp support 'O1', 'O2', 'OFF'
+        }
+        return train_args

paddlex/modules/semantic_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 SegDatasetChecker
-from .trainer import SegTrainer
 from .evaluator import SegEvaluator
 from .exportor import SegExportor
+from .trainer import SegTrainer

paddlex/modules/semantic_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.
@@ -15,11 +15,11 @@
 
 import os
 import os.path as osp
+from pathlib import Path
 
 from ...base import BaseDatasetChecker
-from .dataset_src import check_dataset, convert_dataset, split_dataset, anaylse_dataset
-
 from ..model_list import MODELS
+from .dataset_src import anaylse_dataset, check_dataset, convert_dataset, split_dataset
 
 
 class SegDatasetChecker(BaseDatasetChecker):
@@ -28,6 +28,20 @@ class SegDatasetChecker(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/semantic_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 anaylse_dataset
 from .check_dataset import check_dataset
 from .convert_dataset import convert_dataset
 from .split_dataset import split_dataset
-from .analyse_dataset import anaylse_dataset

paddlex/modules/semantic_segmentation/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.
@@ -16,14 +16,17 @@
 import os
 import os.path as osp
 
-import matplotlib.pyplot as plt
 import numpy as np
 from PIL import Image, ImageOps
 
+from .....utils.deps import function_requires_deps, is_dep_available
 from .....utils.file_interface import custom_open
-from .....utils.logging import info
 
+if is_dep_available("matplotlib"):
+    import matplotlib.pyplot as plt
 
+
+@function_requires_deps("matplotlib")
 def anaylse_dataset(dataset_dir, output):
     """class analysis for dataset"""
 

paddlex/modules/semantic_segmentation/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.
@@ -19,10 +19,10 @@ import os.path as osp
 import numpy as np
 from PIL import Image, ImageOps
 
-from .utils.visualizer import visualize
 from .....utils.errors import DatasetFileNotFoundError
 from .....utils.file_interface import custom_open
 from .....utils.logging import info
+from .utils.visualizer import visualize
 
 
 def check_dataset(dataset_dir, output, sample_num=10):

paddlex/modules/semantic_segmentation/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,18 +13,20 @@
 # limitations under the License.
 
 
-import glob
 import json
 import os
 import os.path as osp
 import shutil
 
-import cv2
 import numpy as np
 from PIL import Image, ImageDraw
 
-from .....utils.file_interface import custom_open
 from .....utils import logging
+from .....utils.deps import function_requires_deps, is_dep_available
+from .....utils.file_interface import custom_open
+
+if is_dep_available("opencv-contrib-python"):
+    import cv2
 
 
 def convert_dataset(dataset_type, input_dir):
@@ -35,6 +37,7 @@ def convert_dataset(dataset_type, input_dir):
         raise NotImplementedError(dataset_type)
 
 
+@function_requires_deps("opencv-contrib-python")
 def convert_labelme_dataset(input_dir):
     """convert labelme format to paddlex official format"""
     bg_name = "_background_"

paddlex/modules/semantic_segmentation/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.
@@ -18,8 +18,8 @@ import os.path as osp
 import random
 import shutil
 
-from .....utils.file_interface import custom_open
 from .....utils import logging
+from .....utils.file_interface import custom_open
 
 
 def split_dataset(root_dir, train_percent, val_percent):

paddlex/modules/semantic_segmentation/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/semantic_segmentation/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.
@@ -13,9 +13,10 @@
 # limitations under the License.
 
 
-import cv2
 import numpy as np
 
+from ......utils.deps import function_requires_deps
+
 
 def get_color_map_list(length):
     """Returns the color map for visualizing the segmentation mask"""
@@ -34,8 +35,11 @@ def get_color_map_list(length):
     return color_map
 
 
+@function_requires_deps("opencv-contrib-python")
 def visualize(image, result, weight=0.6, use_multilabel=False):
     """Convert predict result to color image, and save added image."""
+    import cv2
+
     color_map = get_color_map_list(256)
     color_map = [color_map[i : i + 3] for i in range(0, len(color_map), 3)]
     color_map = np.array(color_map).astype("uint8")

paddlex/modules/semantic_segmentation/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/semantic_segmentation/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.
@@ -20,3 +20,12 @@ class SegExportor(BaseExportor):
     """Semantic Segmentation Model Exportor"""
 
     entities = MODELS
+
+    def get_export_kwargs(self):
+        """get key-value arguments of model export function"""
+        kwargs = super().get_export_kwargs()
+
+        input_shape = self.export_config.get("input_shape")
+        if input_shape is not None:
+            kwargs["input_shape"] = input_shape
+        return kwargs

paddlex/modules/semantic_segmentation/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.
@@ -32,4 +32,6 @@ MODELS = [
     "SeaFormer_tiny",
     "SeaFormer_small",
     "SeaFormer_large",
+    "MaskFormer_tiny",
+    "MaskFormer_small",
 ]