sam和mobilesam导出预处理的onnx

一、前言

sam或者mobilesam的python推理都存在一些前处理，如下所示：

sam.to(device='cuda')
predictor = SamPredictor(sam)
predictor.set_image(image)
image_embedding = predictor.get_image_embedding().cpu().numpy()

    checkpoint = "./weights/mobile_sam.pt"
    model_type = "vit_t"
    
    sam = sam_model_registry[model_type](checkpoint=checkpoint)
    
    # export onnx
    onnx_model_path = "sam_onnx_example_new.onnx"
    
    onnx_model = SamOnnxModel(sam, return_single_mask=True)
    
    print(checkpoint)
    
    export_onnx_model(onnx_model)
    
    """
    如果需要，还可以对模型进行量化和优化。我们发现，这显著改善了web运行时，
    而性能的变化可以忽略不计。
    """
    
    result_quantized = quantized_model(onnx_model_quantized_path = "sam_onnx_quantized_example.onnx")
    
    
    image = cv2.imread('images/test/picture2.jpg')
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    plt.figure(figsize=(10,10))
    plt.imshow(image)
    plt.axis('on')
    plt.show()
    
    
    ort_session = onnxruntime.InferenceSession(onnx_model_path)
    
    """
    要使用ONNX模型，必须首先使用SAM图像编码器对图像进行预处理。
    这是一个较重的过程ÿ