Simplify SAM 2.1 decoder output and fix ONNX export warnings

models/mask-object-sam21-base-plus/demo.py

@@ -93,13 +93,13 @@ def run_inference(encoder_path, decoder_path, image_path, output_path, points):
         "mask_input": mask_input,
         "has_mask_input": has_mask_input,
     })
-    masks, iou_predictions, low_res_masks = dec_outputs
+    masks, iou_predictions = dec_outputs
     t_dec = time.perf_counter()
     print(f"  Decoder:       {t_dec - t_enc:.3f}s")
     print(f"  Masks shape:   {masks.shape}")
     print(f"  IoU scores:    {iou_predictions[0]}")
 
-    # masks shape: [B, 3, 1024, 1024] - select best by IoU
+    # Select best mask by IoU
     best_idx = int(np.argmax(iou_predictions[0]))
     mask_logits = masks[0, best_idx]
     print(f"  Best mask:     #{best_idx} (IoU={iou_predictions[0, best_idx]:.4f})")

models/mask-object-sam21-small/convert.py

@@ -14,8 +14,6 @@
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
-
 from sam2.build_sam import build_sam2
 from sam2.modeling.sam2_base import SAM2Base
 
@@ -103,9 +101,8 @@ class SAM2DecoderOnnxModel(nn.Module):
         has_mask_input:   (1,)
 
     Outputs:
-        masks:           (B, num_masks, 1024, 1024)
+        masks:           (B, num_masks, 256, 256)
         iou_predictions: (B, num_masks)
-        low_res_masks:   (B, num_masks, 256, 256)
     """
 
     def __init__(self, sam_model: SAM2Base, multimask_output: bool):
@@ -132,7 +129,7 @@ def forward(
 
         high_res_feats = [high_res_feats_0, high_res_feats_1]
 
-        low_res_masks, iou_predictions, _, _ = self.mask_decoder.predict_masks(
+        masks, iou_predictions, _, _ = self.mask_decoder.predict_masks(
             image_embeddings=image_embed,
             image_pe=self.prompt_encoder.get_dense_pe(),
             sparse_prompt_embeddings=sparse_embedding,
@@ -141,13 +138,9 @@ def forward(
             high_res_features=high_res_feats,
         )
 
-        # Upscale from 256x256 to full 1024x1024 resolution
-        masks = F.interpolate(low_res_masks, size=(1024, 1024), mode="bilinear", align_corners=False)
-
         if self.multimask_output:
             masks = masks[:, 1:, :, :]
             iou_predictions = iou_predictions[:, 1:]
-            low_res_masks = low_res_masks[:, 1:, :, :]
         else:
             masks, iou_predictions = (
                 self.mask_decoder._dynamic_multimask_via_stability(
@@ -157,7 +150,7 @@ def forward(
 
         masks = torch.clamp(masks, -32.0, 32.0)
 
-        return masks, iou_predictions, low_res_masks
+        return masks, iou_predictions
 
     def _embed_points(
         self, point_coords: torch.Tensor, point_labels: torch.Tensor
@@ -295,7 +288,7 @@ def run_decoder_export(
                 opset_version=opset,
                 do_constant_folding=True,
                 input_names=list(dummy_inputs.keys()),
-                output_names=["masks", "iou_predictions", "low_res_masks"],
+                output_names=["masks", "iou_predictions"],
                 dynamic_axes=dynamic_axes,
                 dynamo=False,
             )
@@ -316,6 +309,16 @@ def simplify_model(model_path):
     onnx_model = onnx.load(model_path)
     model_simp, check = onnxsim.simplify(onnx_model)
     if check:
+        # Remove value_info entries with empty shapes that cause
+        # onnxruntime MergeShapeInfo warnings (e.g. conv_s0, conv_s1)
+        to_remove = [
+            vi for vi in model_simp.graph.value_info
+            if len(vi.type.tensor_type.shape.dim) == 0
+        ]
+        for vi in to_remove:
+            model_simp.graph.value_info.remove(vi)
+        if to_remove:
+            print(f"Removed {len(to_remove)} empty shape annotations.")
         onnx.save(model_simp, model_path)
         print("Graph simplification passed.")
     else:

models/mask-object-sam21-small/demo.py

@@ -93,13 +93,13 @@ def run_inference(encoder_path, decoder_path, image_path, output_path, points):
         "mask_input": mask_input,
         "has_mask_input": has_mask_input,
     })
-    masks, iou_predictions, low_res_masks = dec_outputs
+    masks, iou_predictions = dec_outputs
     t_dec = time.perf_counter()
     print(f"  Decoder:       {t_dec - t_enc:.3f}s")
     print(f"  Masks shape:   {masks.shape}")
     print(f"  IoU scores:    {iou_predictions[0]}")
 
-    # masks shape: [B, 3, 1024, 1024] - select best by IoU
+    # Select best mask by IoU
     best_idx = int(np.argmax(iou_predictions[0]))
     mask_logits = masks[0, best_idx]
     print(f"  Best mask:     #{best_idx} (IoU={iou_predictions[0, best_idx]:.4f})")

models/mask-object-sam21-tiny/demo.py

@@ -93,13 +93,13 @@ def run_inference(encoder_path, decoder_path, image_path, output_path, points):
         "mask_input": mask_input,
         "has_mask_input": has_mask_input,
     })
-    masks, iou_predictions, low_res_masks = dec_outputs
+    masks, iou_predictions = dec_outputs
     t_dec = time.perf_counter()
     print(f"  Decoder:       {t_dec - t_enc:.3f}s")
     print(f"  Masks shape:   {masks.shape}")
     print(f"  IoU scores:    {iou_predictions[0]}")
 
-    # masks shape: [B, 3, 1024, 1024] - select best by IoU
+    # Select best mask by IoU
     best_idx = int(np.argmax(iou_predictions[0]))
     mask_logits = masks[0, best_idx]
     print(f"  Best mask:     #{best_idx} (IoU={iou_predictions[0, best_idx]:.4f})")

pyproject.toml

@@ -20,6 +20,7 @@ core = [
     "numpy<2.0",
     "onnx>=1.10.0",
     "onnxruntime",
+    "onnxsim",
     "opencv-python",
 ]
 nafnet = [
@@ -29,7 +30,6 @@ nafnet = [
 nind = [
     {include-group = "core"},
     "onnxconverter-common",
-    "onnxsim",
 ]
 segnext = [
     {include-group = "core"},