diff --git a/models/mask-object-sam21-base-plus/demo.py b/models/mask-object-sam21-base-plus/demo.py index 1501498..dd98149 100644 --- a/models/mask-object-sam21-base-plus/demo.py +++ b/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})") diff --git a/models/mask-object-sam21-small/convert.py b/models/mask-object-sam21-small/convert.py index af0af7b..39712cc 100644 --- a/models/mask-object-sam21-small/convert.py +++ b/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: diff --git a/models/mask-object-sam21-small/demo.py b/models/mask-object-sam21-small/demo.py index 1501498..dd98149 100644 --- a/models/mask-object-sam21-small/demo.py +++ b/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})") diff --git a/models/mask-object-sam21-tiny/demo.py b/models/mask-object-sam21-tiny/demo.py index 1501498..dd98149 100644 --- a/models/mask-object-sam21-tiny/demo.py +++ b/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})") diff --git a/pyproject.toml b/pyproject.toml index 10f04ac..560b2bd 100644 --- a/pyproject.toml +++ b/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"},