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

paddlex/repo_apis/Paddle3D_api/bev_fusion/model.py

@@ -0,0 +1,238 @@
+# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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 os.path as osp
+
+from .runner import raise_unsupported_api_error
+from ...base import BaseModel
+from ....utils import logging
+from ...base.utils.arg import CLIArgument
+from ....utils.misc import abspath
+
+
+class BEVFusionModel(BaseModel):
+
+    def train(
+        self,
+        batch_size=None,
+        learning_rate=None,
+        epochs_iters=None,
+        pretrained=None,
+        ips=None,
+        device="gpu",
+        resume_path=None,
+        dy2st=False,
+        amp="OFF",
+        num_workers=None,
+        use_vdl=True,
+        save_dir=None,
+        **kwargs,
+    ):
+        if resume_path is not None:
+            resume_path = abspath(resume_path)
+        if not use_vdl:
+            logging.warning("Currently, VisualDL cannot be disabled during training.")
+        if save_dir is not None:
+            save_dir = abspath(save_dir)
+        else:
+            # `save_dir` is None
+            save_dir = abspath(osp.join("output", "train"))
+
+        if dy2st:
+            raise ValueError(f"`dy2st`={dy2st} is not supported.")
+        if device in ("cpu", "gpu"):
+            logging.warning(
+                f"The device type to use will be automatically determined, which may differ from the sepcified type: {repr(device)}."
+            )
+
+        # Update YAML config file
+        config = self.config.copy()
+        if epochs_iters is not None:
+            config.update_iters(epochs_iters)
+        if amp is not None:
+            if amp != "OFF":
+                config._update_amp(amp)
+
+        # Parse CLI arguments
+        cli_args = []
+        if batch_size is not None:
+            cli_args.append(CLIArgument("--batch_size", batch_size))
+        if learning_rate is not None:
+            cli_args.append(CLIArgument("--learning_rate", learning_rate))
+        if num_workers is not None:
+            cli_args.append(CLIArgument("--num_workers", num_workers))
+        if resume_path is not None:
+            if save_dir is not None:
+                raise ValueError(
+                    "When `resume_path` is not None, `save_dir` must be set to None."
+                )
+            model_dir = osp.dirname(resume_path)
+            cli_args.append(CLIArgument("--resume"))
+            cli_args.append(CLIArgument("--save_dir", model_dir))
+        if save_dir is not None:
+            cli_args.append(CLIArgument("--save_dir", save_dir))
+        if pretrained is not None:
+            cli_args.append(CLIArgument("--model", abspath(pretrained)))
+
+        do_eval = kwargs.pop("do_eval", True)
+
+        profile = kwargs.pop("profile", None)
+        if profile is not None:
+            cli_args.append(CLIArgument("--profiler_options", profile))
+
+        log_interval = kwargs.pop("log_interval", 1)
+        if log_interval is not None:
+            cli_args.append(CLIArgument("--log_interval", log_interval))
+
+        save_interval = kwargs.pop("save_interval", 1)
+        if save_interval is not None:
+            cli_args.append(CLIArgument("--save_interval", save_interval))
+
+        seed = kwargs.pop("seed", None)
+        if seed is not None:
+            cli_args.append(CLIArgument("--seed", seed))
+
+        self._assert_empty_kwargs(kwargs)
+
+        # PDX related settings
+        uniform_output_enabled = kwargs.pop("uniform_output_enabled", True)
+        export_with_pir = kwargs.pop("export_with_pir", False)
+        config.update({"uniform_output_enabled": uniform_output_enabled})
+        config.update({"pdx_model_name": self.name})
+        if export_with_pir:
+            config.update({"export_with_pir": export_with_pir})
+
+        with self._create_new_config_file() as config_path:
+            config.dump(config_path)
+            return self.runner.train(
+                config_path, cli_args, device, ips, save_dir, do_eval=do_eval
+            )
+
+    def evaluate(
+        self,
+        weight_path,
+        batch_size=None,
+        ips=None,
+        device="gpu",
+        amp="OFF",
+        num_workers=None,
+        **kwargs,
+    ):
+        weight_path = abspath(weight_path)
+
+        if device in ("cpu", "gpu"):
+            logging.warning(
+                f"The device type to use will be automatically determined, which may differ from the sepcified type: {repr(device)}."
+            )
+
+        # Update YAML config file
+        config = self.config.copy()
+
+        if amp is not None:
+            if amp != "OFF":
+                raise ValueError("AMP evaluation is not supported.")
+
+        # Parse CLI arguments
+        cli_args = []
+        if weight_path is not None:
+            cli_args.append(CLIArgument("--model", weight_path))
+        if batch_size is not None:
+            cli_args.append(CLIArgument("--batch_size", batch_size))
+            if batch_size != 1:
+                raise ValueError("Batch size other than 1 is not supported.")
+        if num_workers is not None:
+            cli_args.append(CLIArgument("--num_workers", num_workers))
+
+        self._assert_empty_kwargs(kwargs)
+
+        # PDX related settings
+        uniform_output_enabled = kwargs.pop("uniform_output_enabled", True)
+        export_with_pir = kwargs.pop("export_with_pir", False)
+        config.update({"uniform_output_enabled": uniform_output_enabled})
+        config.update({"pdx_model_name": self.name})
+        if export_with_pir:
+            config.update({"export_with_pir": export_with_pir})
+
+        with self._create_new_config_file() as config_path:
+            config.dump(config_path)
+            cp = self.runner.evaluate(config_path, cli_args, device, ips)
+            return cp
+
+    def predict(self, weight_path, input_path, device="gpu", save_dir=None, **kwargs):
+        raise_unsupported_api_error("predict", self.__class__)
+
+    def export(self, weight_path, save_dir, **kwargs):
+        if not weight_path.startswith("http"):
+            weight_path = abspath(weight_path)
+        save_dir = abspath(save_dir)
+
+        # Update YAML config file
+        config = self.config.copy()
+
+        # Parse CLI arguments
+        cli_args = []
+        if weight_path is not None:
+            cli_args.append(CLIArgument("--model", weight_path))
+        if save_dir is not None:
+            cli_args.append(CLIArgument("--save_dir", save_dir))
+
+        cli_args.append(CLIArgument("--save_name", "inference"))
+        cli_args.append(CLIArgument("--save_inference_yml"))
+
+        # PDX related settings
+        uniform_output_enabled = kwargs.pop("uniform_output_enabled", True)
+        export_with_pir = kwargs.pop("export_with_pir", False)
+        config.update({"uniform_output_enabled": uniform_output_enabled})
+        config.update({"pdx_model_name": self.name})
+        if export_with_pir:
+            config.update({"export_with_pir": export_with_pir})
+
+        self._assert_empty_kwargs(kwargs)
+        with self._create_new_config_file() as config_path:
+            config.dump(config_path)
+            return self.runner.export(config_path, cli_args, None)
+
+    def infer(self, model_dir, device="gpu", **kwargs):
+        model_dir = abspath(model_dir)
+
+        # Parse CLI arguments
+        cli_args = []
+        model_file_path = osp.join(model_dir, ".pdmodel")
+        params_file_path = osp.join(model_dir, ".pdiparams")
+        cli_args.append(CLIArgument("--model_file", model_file_path))
+        cli_args.append(CLIArgument("--params_file", params_file_path))
+        if device is not None:
+            device_type, _ = self.runner.parse_device(device)
+            if device_type not in ("cpu", "gpu"):
+                raise ValueError(f"`device`={repr(device)} is not supported.")
+        infer_dir = osp.join(self.runner.runner_root_path, self.model_info["infer_dir"])
+        self._assert_empty_kwargs(kwargs)
+        # The inference script does not require a config file
+        return self.runner.infer(None, cli_args, device, infer_dir, None)
+
+    def compression(
+        self,
+        weight_path,
+        ann_file=None,
+        class_names=None,
+        batch_size=None,
+        learning_rate=None,
+        epochs_iters=None,
+        device="gpu",
+        use_vdl=True,
+        save_dir=None,
+        **kwargs,
+    ):
+        raise_unsupported_api_error("compression", self.__class__)

paddlex/repo_apis/Paddle3D_api/bev_fusion/register.py

@@ -0,0 +1,55 @@
+# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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 os
+import os.path as osp
+
+from ...base.register import register_model_info, register_suite_info
+from .model import BEVFusionModel
+from .runner import BEVFusionRunner
+from .config import BEVFusionConfig
+
+REPO_ROOT_PATH = os.environ.get("PADDLE_PDX_PADDLE3D_PATH")
+PDX_CONFIG_DIR = osp.abspath(osp.join(osp.dirname(__file__), "..", "configs"))
+
+register_suite_info(
+    {
+        "suite_name": "BEVFusion",
+        "model": BEVFusionModel,
+        "runner": BEVFusionRunner,
+        "config": BEVFusionConfig,
+        "runner_root_path": REPO_ROOT_PATH,
+    }
+)
+
+register_model_info(
+    {
+        "model_name": "BEVFusion",
+        "suite": "BEVFusion",
+        "config_path": osp.join(PDX_CONFIG_DIR, "BEVFusion.yaml"),
+        "auto_compression_config_path": osp.join(PDX_CONFIG_DIR, "None"),
+        "supported_apis": ["train", "evaluate", "export", "infer"],
+        "supported_train_opts": {
+            "device": ["cpu", "gpu_nxcx"],
+            "dy2st": False,
+            "amp": ["O1", "O2"],
+        },
+        "supported_evaluate_opts": {"device": ["cpu", "gpu_nxcx"]},
+        "supported_infer_opts": {"device": ["cpu", "gpu"]},
+        "supported_dataset_types": [],
+        # Additional info
+        "infer_dir": "deploy/bevfusion",
+    }
+)

paddlex/repo_apis/Paddle3D_api/bev_fusion/runner.py

@@ -0,0 +1,104 @@
+# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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 os
+
+from ...base import BaseRunner
+
+
+def raise_unsupported_api_error(api_name, cls=None):
+    # TODO: Automatically extract `api_name` and `cls` from stack frame
+    if cls is not None:
+        name = f"{cls.__name__}.{api_name}"
+    else:
+        name = api_name
+    raise UnsupportedAPIError(f"The API `{name}` is not supported.")
+
+
+class UnsupportedAPIError(Exception):
+    pass
+
+
+class BEVFusionRunner(BaseRunner):
+    def train(self, config_path, cli_args, device, ips, save_dir, do_eval=True):
+        args, env = self.distributed(device, ips, log_dir=save_dir)
+        cmd = [*args, "tools/train.py"]
+        if do_eval:
+            cmd.append("--do_eval")
+        cmd.extend(["--config", config_path, *cli_args])
+        return self.run_cmd(
+            cmd,
+            env=env,
+            switch_wdir=True,
+            echo=True,
+            silent=False,
+            capture_output=True,
+            log_path=self._get_train_log_path(save_dir),
+        )
+
+    def evaluate(self, config_path, cli_args, device, ips):
+        args, env = self.distributed(device, ips)
+        cmd = [*args, "tools/evaluate.py", "--config", config_path, *cli_args]
+        cp = self.run_cmd(
+            cmd, env=env, switch_wdir=True, echo=True, silent=False, capture_output=True
+        )
+        if cp.returncode == 0:
+            metric_dict = _extract_eval_metrics(cp.stdout)
+            cp.metrics = metric_dict
+        return cp
+
+    def predict(self, config_path, cli_args, device):
+        raise_unsupported_api_error("predict", self.__class__)
+
+    def export(self, config_path, cli_args, device):
+        # `device` unused
+        cmd = [self.python, "tools/export.py", "--config", config_path, *cli_args]
+        return self.run_cmd(cmd, switch_wdir=True, echo=True, silent=False)
+
+    def infer(self, config_path, cli_args, device, infer_dir, save_dir=None):
+        # `config_path` and `device` unused
+        cmd = [self.python, "infer.py", *cli_args]
+        python_infer_dir = os.path.join(infer_dir, "python")
+        cp = self.run_cmd(cmd, switch_wdir=python_infer_dir, echo=True, silent=False)
+        return cp
+
+    def compression(
+        self, config_path, train_cli_args, export_cli_args, device, train_save_dir
+    ):
+        raise_unsupported_api_error("compression", self.__class__)
+
+
+def _extract_eval_metrics(stdout):
+    import re
+
+    _DP = r"[-+]?[0-9]*\.?[0-9]+(?:[eE][-+]?[0-9]+)?"
+    metrics = ["mAP", "NDS"]
+    patterns = {}
+    for metric in metrics:
+        pattern = f"{metric}: (_dp)".replace("_dp", _DP)
+        patterns[metric] = pattern
+
+    metric_dict = dict()
+
+    # TODO: Use lazy version to make it more efficient
+    lines = stdout.splitlines()
+    for line in lines:
+        for m in patterns:
+            p = re.compile(patterns[m])
+            match = p.search(line)
+            if match:
+                metric_dict[m] = float(match.groups()[0])
+
+    return metric_dict

paddlex/repo_apis/Paddle3D_api/pp3d_config.py

@@ -0,0 +1,144 @@
+# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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 codecs
+import yaml
+from ...utils.misc import abspath
+
+from ..base import BaseConfig
+
+
+class PP3DConfig(BaseConfig):
+    # Refer to https://github.com/PaddlePaddle/Paddle3D/blob/release/1.0/paddle3d/apis/config.py
+    def update(self, dict_like_obj):
+        def _merge_config_dicts(dict_from, dict_to):
+            # According to
+            # https://github.com/PaddlePaddle/Paddle3D/blob/3cf884ecbc94330be0e2db780434bb60b9b4fe8c/paddle3d/apis/config.py#L90
+            for key, val in dict_from.items():
+                if isinstance(val, dict) and key in dict_to:
+                    dict_to[key] = _merge_config_dicts(val, dict_to[key])
+                else:
+                    dict_to[key] = val
+            return dict_to
+
+        dict_ = _merge_config_dicts(dict_like_obj, self.dict)
+        self.reset_from_dict(dict_)
+
+    def load(self, config_path):
+        with codecs.open(config_path, "r", "utf-8") as file:
+            dic = yaml.load(file, Loader=yaml.FullLoader)
+        dict_ = dic
+        self.reset_from_dict(dict_)
+
+    def dump(self, config_path):
+        with open(config_path, "w", encoding="utf-8") as f:
+            yaml.dump(self.dict, f)
+
+    def update_learning_rate(self, learning_rate):
+        if "lr_scheduler" not in self:
+            raise RuntimeError(
+                "Not able to update learning rate, because no LR scheduler config was found."
+            )
+
+        # Some lr_scheduler in Paddle3D has not learning_rate parameter.
+        if self.lr_scheduler["type"] == "OneCycle":
+            self.lr_scheduler["lr_max"] = learning_rate
+        elif self.lr_scheduler["type"] == "OneCycleWarmupDecayLr":
+            self.lr_scheduler["base_learning_rate"] = learning_rate
+        else:
+            self.lr_scheduler["learning_rate"] = learning_rate
+
+    def update_batch_size(self, batch_size, mode="train"):
+        if mode == "train":
+            self.set_val("batch_size", batch_size)
+        else:
+            raise ValueError(
+                f"Setting `batch_size` in {repr(mode)} mode is not supported."
+            )
+
+    def update_epochs(self, epochs, mode="train"):
+        if mode == "train":
+            self.set_val("epochs", epochs)
+        else:
+            raise ValueError(f"Setting `epochs` in {repr(mode)} mode is not supported.")
+
+    def update_pretrained_weights(self, weight_path, is_backbone=False):
+        raise NotImplementedError
+
+    def get_epochs_iters(self):
+        if "iters" in self:
+            return self.iters
+        else:
+            assert "epochs" in self
+            return self.epochs
+
+    def get_learning_rate(self):
+        if "lr_scheduler" not in self or "learning_rate" not in self.lr_scheduler:
+            # Default lr
+            return 0.0001
+        else:
+            lr = self.lr_scheduler["learning_rate"]
+            while isinstance(lr, dict):
+                lr = lr["learning_rate"]
+            return lr
+
+    def get_batch_size(self, mode="train"):
+        if "batch_size" in self:
+            return self.batch_size
+        else:
+            # Default batch size
+            return 1
+
+    def get_qat_epochs_iters(self):
+        assert (
+            "finetune_config" in self
+        ), "QAT training yaml should contain finetune_config key"
+        if "iters" in self.finetune_config:
+            return self.finetune_config["iters"]
+        else:
+            assert "epochs" in self.finetune_config
+            return self.finetune_config["epochs"]
+
+    def get_qat_learning_rate(self):
+        assert (
+            "finetune_config" in self
+        ), "QAT training yaml should contain finetune_config key"
+        cfg = self.finetune_config
+        if "lr_scheduler" in cfg or "learning_rate" not in cfg.lr_scheduler:
+            # Default lr
+            return 1.25e-4
+        else:
+            lr = cfg.lr_scheduler["learning_rate"]
+            while isinstance(lr, dict):
+                lr = lr["learning_rate"]
+            return lr
+
+    def update_warmup_steps(self, steps):
+        self.lr_scheduler["warmup_steps"] = steps
+
+    def update_end_lr(self, learning_rate):
+        self.lr_scheduler["end_lr"] = learning_rate
+
+    def update_iters(self, iters):
+        self.set_val("iters", iters)
+        if "epochs" in self:
+            self.set_val("epochs", None)
+
+    def update_finetune_iters(self, iters):
+        self.finetune_config["iters"] = iters
+        if "epochs" in self.finetune_config:
+            self.finetune_config["epochs"] = None
+
+    def update_save_dir(self, save_dir: str):
+        self["save_dir"] = abspath(save_dir)

paddlex/repo_apis/PaddleClas_api/cls/model.py

@@ -115,8 +115,11 @@ class ClsModel(BaseModel):
             # PDX related settings
             device_type = device.split(":")[0]
             uniform_output_enabled = kwargs.pop("uniform_output_enabled", True)
+            export_with_pir = kwargs.pop("export_with_pir", False)
             config.update([f"Global.uniform_output_enabled={uniform_output_enabled}"])
             config.update([f"Global.pdx_model_name={self.name}"])
+            if export_with_pir:
+                config.update([f"Global.export_with_pir={export_with_pir}"])
 
             config.dump(config_path)
             self._assert_empty_kwargs(kwargs)
@@ -231,8 +234,11 @@ class ClsModel(BaseModel):
                 config.update_device(device)
             # PDX related settings
             uniform_output_enabled = kwargs.pop("uniform_output_enabled", True)
+            export_with_pir = kwargs.pop("export_with_pir", False)
             config.update([f"Global.uniform_output_enabled={uniform_output_enabled}"])
             config.update([f"Global.pdx_model_name={self.name}"])
+            if export_with_pir:
+                config.update([f"Global.export_with_pir={export_with_pir}"])
 
             config.dump(config_path)
 

paddlex/repo_apis/PaddleClas_api/cls/register.py

@@ -874,7 +874,6 @@ register_model_info(
         "config_path": osp.join(PDX_CONFIG_DIR, "MobileFaceNet.yaml"),
         "supported_apis": ["train", "evaluate", "predict", "export", "infer"],
         "infer_config": "deploy/configs/inference_cls.yaml",
-        "hpi_config_path": None,
     }
 )
 
@@ -885,6 +884,25 @@ register_model_info(
         "config_path": osp.join(PDX_CONFIG_DIR, "ResNet50_face.yaml"),
         "supported_apis": ["train", "evaluate", "predict", "export", "infer"],
         "infer_config": "deploy/configs/inference_cls.yaml",
-        "hpi_config_path": None,
+    }
+)
+
+register_model_info(
+    {
+        "model_name": "PP-LCNet_x0_25_textline_ori",
+        "suite": "Cls",
+        "config_path": osp.join(PDX_CONFIG_DIR, "PP-LCNet_x0_25_textline_ori.yaml"),
+        "supported_apis": ["train", "evaluate", "predict", "export", "infer"],
+        "infer_config": "deploy/configs/inference_cls.yaml",
+    }
+)
+
+register_model_info(
+    {
+        "model_name": "PP-LCNet_x1_0_table_cls",
+        "suite": "Cls",
+        "config_path": osp.join(PDX_CONFIG_DIR, "PP-LCNet_x1_0_table_cls.yaml"),
+        "supported_apis": ["train", "evaluate", "predict", "export", "infer"],
+        "infer_config": "deploy/configs/inference_cls.yaml",
     }
 )

paddlex/repo_apis/PaddleDetection_api/instance_seg/config.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import List
 from ...base import BaseConfig
 from ....utils.misc import abspath
 from ..config_helper import PPDetConfigMixin
@@ -51,7 +52,7 @@ class InstanceSegConfig(DetConfig):
         dataset_path: str,
         dataset_type: str = None,
         *,
-        data_fields: list[str] = None,
+        data_fields: List[str] = None,
         image_dir: str = "images",
         train_anno_path: str = "annotations/instance_train.json",
         val_anno_path: str = "annotations/instance_val.json",
@@ -95,7 +96,7 @@ class InstanceSegConfig(DetConfig):
     def _make_dataset_config(
         self,
         dataset_root_path: str,
-        data_fields: list[str,] = None,
+        data_fields: List[str,] = None,
         image_dir: str = "images",
         train_anno_path: str = "annotations/instance_train.json",
         val_anno_path: str = "annotations/instance_val.json",
@@ -210,7 +211,7 @@ class InstanceSegConfig(DetConfig):
             if sch[key] == "CosineDecay":
                 sch["max_epochs"] = max_epochs
 
-    def update_milestone(self, milestones: list[int]):
+    def update_milestone(self, milestones: List[int]):
         """update milstone of `PiecewiseDecay` learning scheduler
 
         Args:
@@ -270,6 +271,9 @@ class InstanceSegConfig(DetConfig):
         elif device_type.lower() == "mlu":
             self["use_mlu"] = True
             self["use_gpu"] = False
+        elif device_type.lower() == "gcu":
+            self["use_gcu"] = True
+            self["use_gpu"] = False
         else:
             assert device_type.lower() == "cpu"
             self["use_gpu"] = False
@@ -343,7 +347,7 @@ class InstanceSegConfig(DetConfig):
         """
         self["num_classes"] = num_classes
 
-    def update_random_size(self, randomsize: list[list[int, int]]):
+    def update_random_size(self, randomsize):
         """update `target_size` of `BatchRandomResize` op in TestReader
 
         Args:

paddlex/repo_apis/PaddleDetection_api/instance_seg/model.py

@@ -127,8 +127,11 @@ class InstanceSegModel(BaseModel):
 
         # PDX related settings
         uniform_output_enabled = kwargs.pop("uniform_output_enabled", True)
+        export_with_pir = kwargs.pop("export_with_pir", False)
         config.update({"uniform_output_enabled": uniform_output_enabled})
         config.update({"pdx_model_name": self.name})
+        if export_with_pir:
+            config.update({"export_with_pir": export_with_pir})
 
         self._assert_empty_kwargs(kwargs)
 
@@ -278,8 +281,11 @@ class InstanceSegModel(BaseModel):
 
         # PDX related settings
         uniform_output_enabled = kwargs.pop("uniform_output_enabled", True)
+        export_with_pir = kwargs.pop("export_with_pir", False)
         config.update({"uniform_output_enabled": uniform_output_enabled})
         config.update({"pdx_model_name": self.name})
+        if export_with_pir:
+            config.update({"export_with_pir": export_with_pir})
 
         self._assert_empty_kwargs(kwargs)