Merge branch 'hotfix/onnx_export'

pixano · Jul 7, 2023 · 0acbeba · 0acbeba
2 parents 9a3e8eb + 75070fd
commit 0acbeba
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 6 deletions.
diff --git a/RELEASE.md b/RELEASE.md
@@ -1,6 +1,12 @@
 # Release notes
 
 
+# Pixano Inference v0.1.4
+
+### Fixed
+- Fix export to ONNX for SAM
+
+
 # Pixano Inference v0.1.3
 
 ### Fixed

diff --git a/pixano_inference/__version__.py b/pixano_inference/__version__.py
@@ -11,4 +11,4 @@
 #
 # http://www.cecill.info
 
-__version__ = "0.1.3"
+__version__ = "0.1.4"
diff --git a/pixano_inference/segment_anything/segment_anything.py b/pixano_inference/segment_anything/segment_anything.py
@@ -164,16 +164,16 @@ def export_to_onnx(self, library_dir: Path):
         model_dir.mkdir(parents=True, exist_ok=True)
 
         # Put model to CPU for export
-        self.sam.to("cpu")
+        self.model.to("cpu")
 
         # Export settings
-        onnx_model = SamOnnxModel(self.sam, return_single_mask=True)
+        onnx_model = SamOnnxModel(self.model, return_single_mask=True)
         dynamic_axes = {
             "point_coords": {1: "num_points"},
             "point_labels": {1: "num_points"},
         }
-        embed_dim = self.sam.prompt_encoder.embed_dim
-        embed_size = self.sam.prompt_encoder.image_embedding_size
+        embed_dim = self.model.prompt_encoder.embed_dim
+        embed_size = self.model.prompt_encoder.image_embedding_size
         mask_input_size = [4 * x for x in embed_size]
         dummy_inputs = {
             "image_embeddings": torch.randn(
@@ -220,4 +220,4 @@ def export_to_onnx(self, library_dir: Path):
         )
 
         # Put model back to device after export
-        self.sam.to(self.device)
+        self.model.to(self.device)