how_to_export_vim_h_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Divide the embeding model into 2 parts, named as part1 and part2 \n",
    "# The decoder model remains unchanged\n",
    "# export_embedding_model_part_1()\n",
    "# export_embedding_model_part_2()\n",
    "# export_sam_h_model()\n",
    "# onnx_model_example2()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "from segment_anything.modeling import Sam\n",
    "import numpy as np\n",
    "from torchvision.transforms.functional import resize, to_pil_image\n",
    "from typing import Tuple\n",
    "from segment_anything import sam_model_registry, SamPredictor\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import warnings\n",
    "import os\n",
    "os.environ[\"KMP_DUPLICATE_LIB_OK\"]  =  \"TRUE\"\n",
    "import onnxruntime\n",
    "\n",
    "\n",
    "def show_mask(mask, ax, random_color=False):\n",
    "    if random_color:\n",
    "        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)\n",
    "    else:\n",
    "        color = np.array([30/255, 144/255, 255/255, 0.6])\n",
    "    h, w = mask.shape[-2:]\n",
    "    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)\n",
    "    ax.imshow(mask_image)\n",
    "    \n",
    "def show_points(coords, labels, ax, marker_size=375):\n",
    "    pos_points = coords[labels==1]\n",
    "    neg_points = coords[labels==0]\n",
    "    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)\n",
    "    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)   \n",
    "    \n",
    "def show_box(box, ax):\n",
    "    x0, y0 = box[0], box[1]\n",
    "    w, h = box[2] - box[0], box[3] - box[1]\n",
    "    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2))  \n",
    "\n",
    "\n",
    "def get_preprocess_shape(oldh: int, oldw: int, long_side_length: int) -> Tuple[int, int]:\n",
    "    \"\"\"\n",
    "    Compute the output size given input size and target long side length.\n",
    "    \"\"\"\n",
    "    scale = long_side_length * 1.0 / max(oldh, oldw)\n",
    "    newh, neww = oldh * scale, oldw * scale\n",
    "    neww = int(neww + 0.5)\n",
    "    newh = int(newh + 0.5)\n",
    "    return (newh, neww)\n",
    "\n",
    "# @torch.no_grad()\n",
    "def pre_processing(image: np.ndarray, target_length: int, device,pixel_mean,pixel_std,img_size):\n",
    "    target_size = get_preprocess_shape(image.shape[0], image.shape[1], target_length)\n",
    "    input_image = np.array(resize(to_pil_image(image), target_size))\n",
    "    input_image_torch = torch.as_tensor(input_image, device=device)\n",
    "    input_image_torch = input_image_torch.permute(2, 0, 1).contiguous()[None, :, :, :]\n",
    "\n",
    "    # Normalize colors\n",
    "    input_image_torch = (input_image_torch - pixel_mean) / pixel_std\n",
    "\n",
    "    # Pad\n",
    "    h, w = input_image_torch.shape[-2:]\n",
    "    padh = img_size - h\n",
    "    padw = img_size - w\n",
    "    input_image_torch = F.pad(input_image_torch, (0, padw, 0, padh))\n",
    "    return input_image_torch\n",
    "\n",
    "def export_embedding_model():\n",
    "    sam_checkpoint = \"weights/sam_vit_l_0b3195.pth\"\n",
    "    model_type = \"vit_l\"\n",
    "\n",
    "    device = \"cpu\"\n",
    "\n",
    "    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)\n",
    "    sam.to(device=device)\n",
    "\n",
    "    # image_encoder = EmbeddingOnnxModel(sam)\n",
    "    image = cv2.imread('notebooks/images/truck.jpg')\n",
    "    target_length = sam.image_encoder.img_size\n",
    "    pixel_mean = sam.pixel_mean \n",
    "    pixel_std = sam.pixel_std\n",
    "    img_size = sam.image_encoder.img_size\n",
    "    inputs = pre_processing(image, target_length, device,pixel_mean,pixel_std,img_size)\n",
    "    onnx_model_path = model_type+\"_\"+\"embedding.onnx\"\n",
    "    dummy_inputs = {\n",
    "    \"images\": inputs\n",
    "}\n",
    "    output_names = [\"image_embeddings\"]\n",
    "    image_embeddings = sam.image_encoder(inputs).cpu().numpy()\n",
    "    print('image_embeddings', image_embeddings.shape)\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.filterwarnings(\"ignore\", category=torch.jit.TracerWarning)\n",
    "        warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
    "        with open(onnx_model_path, \"wb\") as f:\n",
    "            torch.onnx.export(\n",
    "                sam.image_encoder,\n",
    "                tuple(dummy_inputs.values()),\n",
    "                f,\n",
    "                export_params=True,\n",
    "                verbose=False,\n",
    "                opset_version=13,\n",
    "                do_constant_folding=True,\n",
    "                input_names=list(dummy_inputs.keys()),\n",
    "                output_names=output_names,\n",
    "                # dynamic_axes=dynamic_axes,\n",
    "            )    \n",
    "\n",
    "def export_sam_model():\n",
    "    from segment_anything.utils.onnx import SamOnnxModel\n",
    "    checkpoint = \"weights/sam_vit_l_0b3195.pth\"\n",
    "    model_type = \"vit_l\"\n",
    "    sam = sam_model_registry[model_type](checkpoint=checkpoint)\n",
    "    onnx_model_path = \"sam_onnx_example.onnx\"\n",
    "\n",
    "    onnx_model = SamOnnxModel(sam, return_single_mask=True)\n",
    "\n",
    "    dynamic_axes = {\n",
    "        \"point_coords\": {1: \"num_points\"},\n",
    "        \"point_labels\": {1: \"num_points\"},\n",
    "    }\n",
    "\n",
    "    embed_dim = sam.prompt_encoder.embed_dim\n",
    "    embed_size = sam.prompt_encoder.image_embedding_size\n",
    "    mask_input_size = [4 * x for x in embed_size]\n",
    "    dummy_inputs = {\n",
    "        \"image_embeddings\": torch.randn(1, embed_dim, *embed_size, dtype=torch.float),\n",
    "        \"point_coords\": torch.randint(low=0, high=1024, size=(1, 5, 2), dtype=torch.float),\n",
    "        \"point_labels\": torch.randint(low=0, high=4, size=(1, 5), dtype=torch.float),\n",
    "        \"mask_input\": torch.randn(1, 1, *mask_input_size, dtype=torch.float),\n",
    "        \"has_mask_input\": torch.tensor([1], dtype=torch.float),\n",
    "        # \"orig_im_size\": torch.tensor([1500, 2250], dtype=torch.float),\n",
    "    }\n",
    "    # output_names = [\"masks\", \"iou_predictions\", \"low_res_masks\"]\n",
    "    output_names = [\"masks\", \"scores\"]\n",
    "\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.filterwarnings(\"ignore\", category=torch.jit.TracerWarning)\n",
    "        warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
    "        with open(onnx_model_path, \"wb\") as f:\n",
    "            torch.onnx.export(\n",
    "                onnx_model,\n",
    "                tuple(dummy_inputs.values()),\n",
    "                f,\n",
    "                export_params=True,\n",
    "                verbose=False,\n",
    "                opset_version=13,\n",
    "                do_constant_folding=True,\n",
    "                input_names=list(dummy_inputs.keys()),\n",
    "                output_names=output_names,\n",
    "                dynamic_axes=dynamic_axes,\n",
    "            )  \n",
    "\n",
    "def onnx_model_example():\n",
    "    import os\n",
    "    ort_session_embedding = onnxruntime.InferenceSession('vit_l_embedding.onnx',providers=['CPUExecutionProvider'])\n",
    "    ort_session_sam = onnxruntime.InferenceSession('sam_onnx_example.onnx',providers=['CPUExecutionProvider'])\n",
    "\n",
    "    image = cv2.imread('notebooks/images/truck.jpg')\n",
    "    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "    image2 = image.copy()\n",
    "    prompt_embed_dim = 256\n",
    "    image_size = 1024\n",
    "    vit_patch_size = 16\n",
    "    image_embedding_size = image_size // vit_patch_size\n",
    "    target_length = image_size\n",
    "    pixel_mean=[123.675, 116.28, 103.53],\n",
    "    pixel_std=[58.395, 57.12, 57.375]\n",
    "    pixel_mean = torch.Tensor(pixel_mean).view(-1, 1, 1)\n",
    "    pixel_std = torch.Tensor(pixel_std).view(-1, 1, 1)\n",
    "    device = \"cpu\"\n",
    "    inputs = pre_processing(image, target_length, device,pixel_mean,pixel_std,image_size)\n",
    "    ort_inputs = {\n",
    "    \"images\": inputs.cpu().numpy()\n",
    "    }\n",
    "    image_embeddings = ort_session_embedding.run(None, ort_inputs)[0]\n",
    "\n",
    "    input_point = np.array([[500, 375]])\n",
    "    input_label = np.array([1])\n",
    "\n",
    "    onnx_coord = np.concatenate([input_point, np.array([[0.0, 0.0]])], axis=0)[None, :, :]\n",
    "    onnx_label = np.concatenate([input_label, np.array([-1])], axis=0)[None, :].astype(np.float32)\n",
    "    from segment_anything.utils.transforms import ResizeLongestSide\n",
    "    transf = ResizeLongestSide(image_size)\n",
    "    onnx_coord = transf.apply_coords(onnx_coord, image.shape[:2]).astype(np.float32)\n",
    "    onnx_mask_input = np.zeros((1, 1, 256, 256), dtype=np.float32)\n",
    "    onnx_has_mask_input = np.zeros(1, dtype=np.float32)\n",
    "\n",
    "    ort_inputs = {\n",
    "        \"image_embeddings\": image_embeddings,\n",
    "        \"point_coords\": onnx_coord,\n",
    "        \"point_labels\": onnx_label,\n",
    "        \"mask_input\": onnx_mask_input,\n",
    "        \"has_mask_input\": onnx_has_mask_input,\n",
    "        # \"orig_im_size\": np.array(image.shape[:2], dtype=np.float32)\n",
    "    }\n",
    "\n",
    "    masks, _ = ort_session_sam.run(None, ort_inputs)\n",
    "\n",
    "    from segment_anything.utils.onnx import SamOnnxModel\n",
    "    checkpoint = \"weights/sam_vit_l_0b3195.pth\"\n",
    "    model_type = \"vit_l\"\n",
    "    sam = sam_model_registry[model_type](checkpoint=checkpoint)\n",
    "    onnx_model_path = \"sam_onnx_example.onnx\"\n",
    "\n",
    "    onnx_model = SamOnnxModel(sam, return_single_mask=True)\n",
    "    masks = onnx_model.mask_postprocessing(torch.as_tensor(masks), torch.as_tensor(image.shape[:2]))\n",
    "    masks = masks > 0.0\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    plt.imshow(image)\n",
    "    show_mask(masks, plt.gca())\n",
    "    # show_box(input_box, plt.gca())\n",
    "    show_points(input_point, input_label, plt.gca())\n",
    "    plt.axis('off')\n",
    "    plt.savefig('demo.png')\n",
    "\n",
    "def export_engine_image_encoder(f='vit_l_embedding.onnx'):\n",
    "    import tensorrt as trt\n",
    "    from pathlib import Path\n",
    "    file = Path(f)\n",
    "    f = file.with_suffix('.engine')  # TensorRT engine file\n",
    "    onnx = file.with_suffix('.onnx')\n",
    "    logger = trt.Logger(trt.Logger.INFO)\n",
    "    builder = trt.Builder(logger)\n",
    "    config = builder.create_builder_config()\n",
    "    workspace = 6\n",
    "    print(\"workspace: \", workspace)\n",
    "    config.max_workspace_size = workspace * 1 << 30\n",
    "    flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))\n",
    "    network = builder.create_network(flag)\n",
    "    parser = trt.OnnxParser(network, logger)\n",
    "    if not parser.parse_from_file(str(onnx)):\n",
    "        raise RuntimeError(f'failed to load ONNX file: {onnx}')\n",
    "\n",
    "    inputs = [network.get_input(i) for i in range(network.num_inputs)]\n",
    "    outputs = [network.get_output(i) for i in range(network.num_outputs)]\n",
    "    for inp in inputs:\n",
    "        print(f'input \"{inp.name}\" with shape{inp.shape} {inp.dtype}')\n",
    "    for out in outputs:\n",
    "        print(f'output \"{out.name}\" with shape{out.shape} {out.dtype}')\n",
    "\n",
    "    half = True\n",
    "    print(f'building FP{16 if builder.platform_has_fast_fp16 and half else 32} engine as {f}')\n",
    "    if builder.platform_has_fast_fp16 and half:\n",
    "        config.set_flag(trt.BuilderFlag.FP16)\n",
    "    with builder.build_engine(network, config) as engine, open(f, 'wb') as t:\n",
    "        t.write(engine.serialize())\n",
    "\n",
    "\n",
    "def export_engine_prompt_encoder_and_mask_decoder(f='sam_onnx_example.onnx'):\n",
    "    import tensorrt as trt\n",
    "    from pathlib import Path\n",
    "    file = Path(f)\n",
    "    f = file.with_suffix('.engine')  # TensorRT engine file\n",
    "    onnx = file.with_suffix('.onnx')\n",
    "    logger = trt.Logger(trt.Logger.INFO)\n",
    "    builder = trt.Builder(logger)\n",
    "    config = builder.create_builder_config()\n",
    "    workspace = 10\n",
    "    print(\"workspace: \", workspace)\n",
    "    config.max_workspace_size = workspace * 1 << 30\n",
    "    flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))\n",
    "    network = builder.create_network(flag)\n",
    "    parser = trt.OnnxParser(network, logger)\n",
    "    if not parser.parse_from_file(str(onnx)):\n",
    "        raise RuntimeError(f'failed to load ONNX file: {onnx}')\n",
    "\n",
    "    inputs = [network.get_input(i) for i in range(network.num_inputs)]\n",
    "    outputs = [network.get_output(i) for i in range(network.num_outputs)]\n",
    "    for inp in inputs:\n",
    "        print(f'input \"{inp.name}\" with shape{inp.shape} {inp.dtype}')\n",
    "    for out in outputs:\n",
    "        print(f'output \"{out.name}\" with shape{out.shape} {out.dtype}')\n",
    "\n",
    "    profile = builder.create_optimization_profile()\n",
    "    profile.set_shape('image_embeddings', (1, 256, 64, 64), (1, 256, 64, 64), (1, 256, 64, 64))\n",
    "    profile.set_shape('point_coords', (1, 2,2), (1, 5,2), (1,10,2))\n",
    "    profile.set_shape('point_labels', (1, 2), (1, 5), (1,10))\n",
    "    profile.set_shape('mask_input', (1, 1, 256, 256), (1, 1, 256, 256), (1, 1, 256, 256))\n",
    "    profile.set_shape('has_mask_input', (1,), (1, ), (1, ))\n",
    "    # # profile.set_shape_input('orig_im_size', (416,416), (1024,1024), (1500, 2250))\n",
    "    # profile.set_shape_input('orig_im_size', (2,), (2,), (2, ))\n",
    "    config.add_optimization_profile(profile)\n",
    "\n",
    "    half = True\n",
    "    print(f'building FP{16 if builder.platform_has_fast_fp16 and half else 32} engine as {f}')\n",
    "    if builder.platform_has_fast_fp16 and half:\n",
    "        config.set_flag(trt.BuilderFlag.FP16)\n",
    "    with builder.build_engine(network, config) as engine, open(f, 'wb') as t:\n",
    "        t.write(engine.serialize())\n",
    "\n",
    "\n",
    "def collect_stats(model, device, data_loader, num_batches):\n",
    "    \"\"\"Feed data to the network and collect statistic\"\"\"\n",
    "\n",
    "    # Enable calibrators\n",
    "    for name, module in model.named_modules():\n",
    "        if isinstance(module, quant_nn.TensorQuantizer):\n",
    "            if module._calibrator is not None:\n",
    "                module.disable_quant()\n",
    "                module.enable_calib()\n",
    "            else:\n",
    "                module.disable()\n",
    "\n",
    "    for i, (path, im, im0s, vid_cap, s) in tqdm(enumerate(data_loader), total=num_batches):\n",
    "        im = torch.from_numpy(im).to(device)\n",
    "        im = im.float()\n",
    "        im /= 255  # 0 - 255 to 0.0 - 1.0\n",
    "        if len(im.shape) == 3:\n",
    "            im = im[None]  # expand for batch dim\n",
    "        model(im)\n",
    "        if i >= num_batches:\n",
    "            break\n",
    "\n",
    "    # Disable calibrators\n",
    "    for name, module in model.named_modules():\n",
    "        if isinstance(module, quant_nn.TensorQuantizer):\n",
    "            if module._calibrator is not None:\n",
    "                module.enable_quant()\n",
    "                module.disable_calib()\n",
    "            else:\n",
    "                module.enable()\n",
    "\n",
    "def compute_amax(model, **kwargs):\n",
    "    # Load calib result\n",
    "    for name, module in model.named_modules():\n",
    "        if isinstance(module, quant_nn.TensorQuantizer):\n",
    "            if module._calibrator is not None:\n",
    "                if isinstance(module._calibrator, calib.MaxCalibrator):\n",
    "                    module.load_calib_amax()\n",
    "                else:\n",
    "                    module.load_calib_amax(**kwargs)\n",
    "\n",
    "class embedding_model_part_1(nn.Module):\n",
    "    def __init__(self) :\n",
    "        super().__init__()\n",
    "        sam_checkpoint = \"weights/sam_vit_h_4b8939.pth\"\n",
    "        model_type = \"vit_h\"\n",
    "        device = \"cpu\"\n",
    "        self.sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)\n",
    "        self.sam.to(device=device)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.sam.image_encoder.patch_embed(x)\n",
    "        if self.sam.image_encoder.pos_embed is not None:\n",
    "            x = x + self.sam.image_encoder.pos_embed\n",
    "\n",
    "        val_1 = len(self.sam.image_encoder.blocks) // 2\n",
    "        for blk in self.sam.image_encoder.blocks[0:val_1]:\n",
    "            x = blk(x)\n",
    "\n",
    "        return x\n",
    "    \n",
    "class embedding_model_part_2(nn.Module):\n",
    "    def __init__(self) :\n",
    "        super().__init__()\n",
    "        sam_checkpoint = \"weights/sam_vit_h_4b8939.pth\"\n",
    "        model_type = \"vit_h\"\n",
    "        device = \"cpu\"\n",
    "        self.sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)\n",
    "        self.sam.to(device=device)\n",
    "\n",
    "    def forward(self, x):\n",
    "        val_1 = len(self.sam.image_encoder.blocks) // 2\n",
    "        for blk in self.sam.image_encoder.blocks[val_1:]:\n",
    "            x = blk(x)\n",
    "        x = self.sam.image_encoder.neck(x.permute(0, 3, 1, 2))\n",
    "        return x\n",
    "\n",
    "def export_embedding_model_part_1():\n",
    "    device = \"cpu\"\n",
    "    model_type = \"vit_h\"\n",
    "    model = embedding_model_part_1()\n",
    "    # image_encoder = EmbeddingOnnxModel(sam)\n",
    "    image = cv2.imread('notebooks/images/truck.jpg')\n",
    "    target_length = model.sam.image_encoder.img_size\n",
    "    pixel_mean = model.sam.pixel_mean \n",
    "    pixel_std = model.sam.pixel_std\n",
    "    img_size = model.sam.image_encoder.img_size\n",
    "    inputs = pre_processing(image, target_length, device,pixel_mean,pixel_std,img_size)\n",
    "    onnx_model_path = model_type+\"_\"+\"part_1_embedding.onnx\"\n",
    "    dummy_inputs = {\n",
    "    \"images\": inputs\n",
    "}\n",
    "    output_names = [\"image_embeddings_part_1\"]\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.filterwarnings(\"ignore\", category=torch.jit.TracerWarning)\n",
    "        warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
    "        with open(onnx_model_path, \"wb\") as f:\n",
    "            torch.onnx.export(\n",
    "                model,\n",
    "                tuple(dummy_inputs.values()),\n",
    "                f,\n",
    "                export_params=True,\n",
    "                verbose=False,\n",
    "                opset_version=13,\n",
    "                do_constant_folding=True,\n",
    "                input_names=list(dummy_inputs.keys()),\n",
    "                output_names=output_names,\n",
    "                # dynamic_axes=dynamic_axes,\n",
    "            )   \n",
    "\n",
    "def export_embedding_model_part_2():\n",
    "    device = \"cpu\"\n",
    "    model_type = \"vit_h\"\n",
    "    model = embedding_model_part_2()\n",
    "    \n",
    "    inputs = torch.randn(1, 64, 64, 1280)\n",
    "    onnx_model_path = model_type+\"_\"+\"part_2_embedding.onnx\"\n",
    "    dummy_inputs = {\n",
    "    \"image_embeddings_part_1\": inputs\n",
    "}\n",
    "    output_names = [\"image_embeddings_part_2\"]\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.filterwarnings(\"ignore\", category=torch.jit.TracerWarning)\n",
    "        warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
    "        with open(onnx_model_path, \"wb\") as f:\n",
    "            torch.onnx.export(\n",
    "                model,\n",
    "                tuple(dummy_inputs.values()),\n",
    "                f,\n",
    "                export_params=True,\n",
    "                verbose=False,\n",
    "                opset_version=13,\n",
    "                do_constant_folding=True,\n",
    "                input_names=list(dummy_inputs.keys()),\n",
    "                output_names=output_names,\n",
    "                # dynamic_axes=dynamic_axes,\n",
    "            )   \n",
    "        \n",
    "def export_sam_h_model():\n",
    "    from segment_anything.utils.onnx import SamOnnxModel\n",
    "    checkpoint = \"weights/sam_vit_h_4b8939.pth\"\n",
    "    model_type = \"vit_h\"\n",
    "    sam = sam_model_registry[model_type](checkpoint=checkpoint)\n",
    "    onnx_model_path = \"sam_h_decoder_onnx.onnx\"\n",
    "\n",
    "    onnx_model = SamOnnxModel(sam, return_single_mask=True)\n",
    "\n",
    "    dynamic_axes = {\n",
    "        \"point_coords\": {1: \"num_points\"},\n",
    "        \"point_labels\": {1: \"num_points\"},\n",
    "    }\n",
    "\n",
    "    embed_dim = sam.prompt_encoder.embed_dim\n",
    "    embed_size = sam.prompt_encoder.image_embedding_size\n",
    "    mask_input_size = [4 * x for x in embed_size]\n",
    "    dummy_inputs = {\n",
    "        \"image_embeddings\": torch.randn(1, embed_dim, *embed_size, dtype=torch.float),\n",
    "        \"point_coords\": torch.randint(low=0, high=1024, size=(1, 5, 2), dtype=torch.float),\n",
    "        \"point_labels\": torch.randint(low=0, high=4, size=(1, 5), dtype=torch.float),\n",
    "        \"mask_input\": torch.randn(1, 1, *mask_input_size, dtype=torch.float),\n",
    "        \"has_mask_input\": torch.tensor([1], dtype=torch.float),\n",
    "        # \"orig_im_size\": torch.tensor([1500, 2250], dtype=torch.float),\n",
    "    }\n",
    "    # output_names = [\"masks\", \"iou_predictions\", \"low_res_masks\"]\n",
    "    output_names = [\"masks\", \"scores\"]\n",
    "\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.filterwarnings(\"ignore\", category=torch.jit.TracerWarning)\n",
    "        warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
    "        with open(onnx_model_path, \"wb\") as f:\n",
    "            torch.onnx.export(\n",
    "                onnx_model,\n",
    "                tuple(dummy_inputs.values()),\n",
    "                f,\n",
    "                export_params=True,\n",
    "                verbose=False,\n",
    "                opset_version=13,\n",
    "                do_constant_folding=True,\n",
    "                input_names=list(dummy_inputs.keys()),\n",
    "                output_names=output_names,\n",
    "                dynamic_axes=dynamic_axes,\n",
    "            )  \n",
    "\n",
    "def onnx_model_example2():\n",
    "    import os\n",
    "    ort_session_embedding_part_1 = onnxruntime.InferenceSession('vit_h_part_1_embedding.onnx',providers=['CPUExecutionProvider'])\n",
    "    ort_session_embedding_part_2 = onnxruntime.InferenceSession('vit_h_part_2_embedding.onnx',providers=['CPUExecutionProvider'])\n",
    "    ort_session_sam = onnxruntime.InferenceSession('sam_h_decoder_onnx.onnx',providers=['CPUExecutionProvider'])\n",
    "\n",
    "    image = cv2.imread('notebooks/images/truck.jpg')\n",
    "    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "    image2 = image.copy()\n",
    "    prompt_embed_dim = 256\n",
    "    image_size = 1024\n",
    "    vit_patch_size = 16\n",
    "    image_embedding_size = image_size // vit_patch_size\n",
    "    target_length = image_size\n",
    "    pixel_mean=[123.675, 116.28, 103.53],\n",
    "    pixel_std=[58.395, 57.12, 57.375]\n",
    "    pixel_mean = torch.Tensor(pixel_mean).view(-1, 1, 1)\n",
    "    pixel_std = torch.Tensor(pixel_std).view(-1, 1, 1)\n",
    "    device = \"cpu\"\n",
    "    inputs = pre_processing(image, target_length, device,pixel_mean,pixel_std,image_size)\n",
    "    ort_inputs = {\n",
    "    \"images\": inputs.cpu().numpy()\n",
    "    }\n",
    "    image_embeddings = ort_session_embedding_part_1.run(None, ort_inputs)[0]\n",
    "    ort_inputs = {\n",
    "    \"image_embeddings_part_1\": image_embeddings\n",
    "    }\n",
    "    image_embeddings = ort_session_embedding_part_2.run(None, ort_inputs)[0]\n",
    "\n",
    "    input_point = np.array([[784, 379]])\n",
    "    input_label = np.array([1])\n",
    "\n",
    "    onnx_coord = np.concatenate([input_point, np.array([[0.0, 0.0]])], axis=0)[None, :, :]\n",
    "    onnx_label = np.concatenate([input_label, np.array([-1])], axis=0)[None, :].astype(np.float32)\n",
    "    from segment_anything.utils.transforms import ResizeLongestSide\n",
    "    transf = ResizeLongestSide(image_size)\n",
    "    onnx_coord = transf.apply_coords(onnx_coord, image.shape[:2]).astype(np.float32)\n",
    "    onnx_mask_input = np.zeros((1, 1, 256, 256), dtype=np.float32)\n",
    "    onnx_has_mask_input = np.zeros(1, dtype=np.float32)\n",
    "\n",
    "    ort_inputs = {\n",
    "        \"image_embeddings\": image_embeddings,\n",
    "        \"point_coords\": onnx_coord,\n",
    "        \"point_labels\": onnx_label,\n",
    "        \"mask_input\": onnx_mask_input,\n",
    "        \"has_mask_input\": onnx_has_mask_input,\n",
    "        # \"orig_im_size\": np.array(image.shape[:2], dtype=np.float32)\n",
    "    }\n",
    "\n",
    "    masks, _ = ort_session_sam.run(None, ort_inputs)\n",
    "\n",
    "    from segment_anything.utils.onnx import SamOnnxModel\n",
    "    checkpoint = \"weights/sam_vit_h_4b8939.pth\"\n",
    "    model_type = \"vit_h\"\n",
    "    sam = sam_model_registry[model_type](checkpoint=checkpoint)\n",
    "\n",
    "    onnx_model = SamOnnxModel(sam, return_single_mask=True)\n",
    "    masks = onnx_model.mask_postprocessing(torch.as_tensor(masks), torch.as_tensor(image.shape[:2]))\n",
    "    masks = masks > 0.0\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    plt.imshow(image)\n",
    "    show_mask(masks, plt.gca())\n",
    "    # show_box(input_box, plt.gca())\n",
    "    show_points(input_point, input_label, plt.gca())\n",
    "    plt.axis('off')\n",
    "    plt.savefig('demo.png')\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    with torch.no_grad():\n",
    "        # export_embedding_model()\n",
    "        # export_sam_model()\n",
    "        # onnx_model_example()\n",
    "        # export_engine_image_encoder()\n",
    "        # export_engine_prompt_encoder_and_mask_decoder()\n",
    "        export_embedding_model_part_1()\n",
    "        export_embedding_model_part_2()\n",
    "        export_sam_h_model()\n",
    "        onnx_model_example2()\n",
    "        "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}