SegmentAnything导出Onnx模型分割图片

1.SegmentAnything导出Onnx模型分割图片
2.官方SegmentAnything导出Onnx模型,不全面。官方只提供了mask的decoder模型,还缺失image的encoder模型
3.b站视频地址 添加链接描述
4.导出模型的核心代码

 #导出onnx
    def ExportOnnx(self):
        # Export masks decoder from SAM model to ONNX
        onnx_model = SamOnnxModel(self.Model, return_single_mask=True)
        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(1, embed_dim, *embed_size, dtype=torch.float),
            "point_coords": torch.randint(low=0, high=1024, size=(1, 5, 2), dtype=torch.float),
            "point_labels": torch.randint(low=0, high=4, size=(1, 5), dtype=torch.float),
            "mask_input": torch.randn(1, 1, *mask_input_size, dtype=torch.float),
            "has_mask_input": torch.tensor([1], dtype=torch.float),
            "orig_im_size": torch.tensor([1500, 2250], dtype=torch.float),
        }
        output_names = ["masks", "iou_predictions", "low_res_masks"]
        torch.onnx.export(
            f="vit_b_decoder.onnx",
            model=onnx_model,
            args=tuple(dummy_inputs.values()),
            input_names=list(dummy_inputs.keys()),
            output_names=output_names,
            dynamic_axes={
                "point_coords": {1: "num_points"},
                "point_labels": {1: "num_points"}
            },
            export_params=True,
            opset_version=17,
            do_constant_folding=True
        )

        # Export images encoder from SAM model to ONNX
        torch.onnx.export(
            f="vit_b_encoder.onnx",
            model=self.Model.image_encoder,
            args=torch.randn(1, 3, 1024, 1024),
            input_names=["images"],
            output_names=["embeddings"],
            export_params=True)

        pass

相关推荐

  1. SegmentAnything导出Onnx模型分割图片

    2024-04-03 09:48:02       5 阅读
  2. Mxnet导出onnx模型

    2024-04-03 09:48:02       22 阅读
  3. yolov5导出onnx模型问题

    2024-04-03 09:48:02       17 阅读

最近更新

  1. flink mysql数据表同步SQL CDC

    2024-04-03 09:48:02       0 阅读
  2. 【QT进阶】Qt http编程之json解析的简单介绍

    2024-04-03 09:48:02       0 阅读
  3. 从0开始深入理解Spring(1)--SpringApplication构造

    2024-04-03 09:48:02       0 阅读
  4. kubeadm 升级 k8s集群 1.17到1.20

    2024-04-03 09:48:02       0 阅读
  5. SpringCloudAlibaba+RocketMQ实现对消息中间件的整合

    2024-04-03 09:48:02       0 阅读
  6. Docker中Kafka容器创建/更新Topic支持多分区

    2024-04-03 09:48:02       0 阅读

热门阅读

  1. 120.单例模式(C++设计模式)

    2024-04-03 09:48:02       3 阅读
  2. 【Node】使用Node.js构建简单的静态页面生成器

    2024-04-03 09:48:02       5 阅读
  3. 工厂方法模式:灵活的创建对象实例

    2024-04-03 09:48:02       5 阅读
  4. WPF —— 关键帧动画

    2024-04-03 09:48:02       3 阅读
  5. Yolov5封装detect.py面向对象

    2024-04-03 09:48:02       2 阅读
  6. Rancher(v2.6.3)——Rancher部署Minio(单机版)

    2024-04-03 09:48:02       3 阅读
  7. STM32为什么不能跑Linux?

    2024-04-03 09:48:02       4 阅读
  8. 菜鸟笔记-Python函数-ones

    2024-04-03 09:48:02       5 阅读
  9. 14、Lua 模块与包

    2024-04-03 09:48:02       2 阅读
  10. 基于单片机的LED 灯调光系统的研究

    2024-04-03 09:48:02       3 阅读
  11. Git 多人协作开发

    2024-04-03 09:48:02       4 阅读