diff --git a/src/anomalib/data/utils/image.py b/src/anomalib/data/utils/image.py index 64a27724cc..85d2cfe2df 100644 --- a/src/anomalib/data/utils/image.py +++ b/src/anomalib/data/utils/image.py @@ -440,7 +440,7 @@ def save_image(filename: Path | str, image: np.ndarray | Figure, root: Path | No # if file_path is absolute, then root is ignored # so we remove the top level directory from the path if file_path.is_absolute() and root: - file_path = Path(*file_path.parts[2:]) # OS-AGNOSTIC + file_path = Path(*file_path.parts[-2:]) # OS-AGNOSTIC if root: file_path = root / file_path diff --git a/src/anomalib/engine/engine.py b/src/anomalib/engine/engine.py index b537819729..eef5734e8d 100644 --- a/src/anomalib/engine/engine.py +++ b/src/anomalib/engine/engine.py @@ -876,6 +876,7 @@ def export( model: AnomalyModule, export_type: ExportType | str, export_root: str | Path | None = None, + model_file_name: str = "model", input_size: tuple[int, int] | None = None, transform: Transform | None = None, compression_type: CompressionType | None = None, @@ -891,6 +892,8 @@ def export( export_type (ExportType): Export type. export_root (str | Path | None, optional): Path to the output directory. If it is not set, the model is exported to trainer.default_root_dir. Defaults to None. + model_file_name (str = "model"): Name of the exported model file. If it is not set, the model is + is called "model". Defaults to "model". input_size (tuple[int, int] | None, optional): A statis input shape for the model, which is exported to ONNX and OpenVINO format. Defaults to None. transform (Transform | None, optional): Input transform to include in the exported model. If not provided, @@ -951,12 +954,14 @@ def export( if export_type == ExportType.TORCH: exported_model_path = model.to_torch( export_root=export_root, + model_file_name=model_file_name, transform=transform, task=self.task, ) elif export_type == ExportType.ONNX: exported_model_path = model.to_onnx( export_root=export_root, + model_file_name=model_file_name, input_size=input_size, transform=transform, task=self.task, @@ -964,6 +969,7 @@ def export( elif export_type == ExportType.OPENVINO: exported_model_path = model.to_openvino( export_root=export_root, + model_file_name=model_file_name, input_size=input_size, transform=transform, task=self.task, diff --git a/src/anomalib/models/components/base/export_mixin.py b/src/anomalib/models/components/base/export_mixin.py index 327cb87e02..af0dc95bd3 100644 --- a/src/anomalib/models/components/base/export_mixin.py +++ b/src/anomalib/models/components/base/export_mixin.py @@ -44,6 +44,7 @@ class ExportMixin: def to_torch( self, export_root: Path | str, + model_file_name: str, transform: Transform | None = None, task: TaskType | None = None, ) -> Path: @@ -51,6 +52,7 @@ def to_torch( Args: export_root (Path): Path to the output folder. + model_file_name (str): Name of the exported model transform (Transform, optional): Input transforms used for the model. If not provided, the transform is taken from the model. Defaults to ``None``. @@ -85,7 +87,7 @@ def to_torch( inference_model = InferenceModel(model=self.model, transform=transform) export_root = _create_export_root(export_root, ExportType.TORCH) metadata = self._get_metadata(task=task) - pt_model_path = export_root / "model.pt" + pt_model_path = export_root / (model_file_name + ".pt") torch.save( obj={"model": inference_model, "metadata": metadata}, f=pt_model_path, @@ -95,6 +97,7 @@ def to_torch( def to_onnx( self, export_root: Path | str, + model_file_name: str, input_size: tuple[int, int] | None = None, transform: Transform | None = None, task: TaskType | None = None, @@ -103,6 +106,7 @@ def to_onnx( Args: export_root (Path): Path to the root folder of the exported model. + model_file_name (str): Name of the exported model. input_size (tuple[int, int] | None, optional): Image size used as the input for onnx converter. Defaults to None. transform (Transform, optional): Input transforms used for the model. If not provided, the transform is @@ -147,7 +151,7 @@ def to_onnx( else {"input": {0: "batch_size", 2: "height", 3: "weight"}, "output": {0: "batch_size"}} ) _write_metadata_to_json(self._get_metadata(task), export_root) - onnx_path = export_root / "model.onnx" + onnx_path = export_root / (model_file_name + ".onnx") torch.onnx.export( inference_model, input_shape.to(self.device), @@ -163,6 +167,7 @@ def to_onnx( def to_openvino( self, export_root: Path | str, + model_file_name: str, input_size: tuple[int, int] | None = None, transform: Transform | None = None, compression_type: CompressionType | None = None, @@ -175,6 +180,7 @@ def to_openvino( Args: export_root (Path): Path to the export folder. + model_file_name (str): Name of the exported model input_size (tuple[int, int] | None, optional): Input size of the model. Used for adding metadata to the IR. Defaults to None. transform (Transform, optional): Input transforms used for the model. If not provided, the transform is @@ -252,9 +258,9 @@ def to_openvino( import openvino as ov with TemporaryDirectory() as onnx_directory: - model_path = self.to_onnx(onnx_directory, input_size, transform, task) + model_path = self.to_onnx(onnx_directory, model_file_name, input_size, transform, task) export_root = _create_export_root(export_root, ExportType.OPENVINO) - ov_model_path = export_root / "model.xml" + ov_model_path = export_root / (model_file_name + ".xml") ov_args = {} if ov_args is None else ov_args model = ov.convert_model(model_path, **ov_args)