Add Support for POT and NNCF Compression in OpenVINO Export

Signed-off-by: Haroon <[email protected]>
openvinotoolkit · Dec 2, 2024 · edf2b51 · edf2b51
1 parent 22d0558
commit edf2b51
Showing 1 changed file with 16 additions and 13 deletions.
diff --git a/tests/integration/cli/test_cli.py b/tests/integration/cli/test_cli.py
@@ -160,33 +160,36 @@ def test_predict_with_image_path(self, project_path: Path) -> None:
         [
             (ExportType.TORCH, None),
             (ExportType.ONNX, None),
-            (ExportType.OPENVIVO, None),
-            (ExportType.OPENVIVO, CompressionType.POT),
-            (ExportType.OPENVIVO, CompressionType.NNCF),
+            (ExportType.OPENVINO, None),
+            (ExportType.OPENVINO, CompressionType.POT),
+            (ExportType.OPENVINO, CompressionType.NNCF),
         ],
     )
     def test_export(
         self,
         project_path: Path,
         export_type: ExportType,
+        compression_type: CompressionType | None,
     ) -> None:
         """Test the export method of the CLI.
 
         Args:
             dataset_path (Path): Root of the synthetic/original dataset.
             project_path (Path): Path to temporary project folder.
             export_type (ExportType): Export type.
+            compression_type (CompressionType): Compression type (if any).
         """
-        AnomalibCLI(
-            args=[
-                "export",
-                "--export_type",
-                export_type,
-                *self._get_common_cli_args(None, project_path),
-                "--ckpt_path",
-                f"{project_path}/Padim/MVTec/dummy/v0/weights/lightning/model.ckpt",
-            ],
-        )
+        args = [
+            "export",
+            "--export_type",
+            export_type,
+            *self._get_common_cli_args(None, project_path),
+            "--ckpt_path",
+            f"{project_path}/Padim/MVTec/dummy/v0/weights/lightning/model.ckpt",
+        ]
+        if compression_type:
+            args += ["--compression_type", str(compression_type)]
+        AnomalibCLI(args=args)
 
     @staticmethod
     def _get_common_cli_args(dataset_path: Path | None, project_path: Path) -> list[str]: