Initial commit
This commit is contained in:
Eric Mintun committed
commit
571794162e
37 files changed
+5830
No files matched your search
@@ -0,0 +1,454 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"id": "5fa21d44",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Copyright (c) Meta Platforms, Inc. and affiliates."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "b7c0041e",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Automatically generating object masks with SAM"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "289bb0b4",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Since SAM can efficiently process prompts, masks for the entire image can be generated by sampling a large number of prompts over an image. This method was used to generate the dataset SA-1B. \n",
|
||||
"\n",
|
||||
"The class `SamAutomaticMaskGenerator` implements this capability. It works by sampling single-point input prompts in a grid over the image, from each of which SAM can predict multiple masks. Then, masks are filtered for quality and deduplicated using non-maximal suppression. Additional options allow for further improvement of mask quality and quantity, such as running prediction on multiple crops of the image or postprocessing masks to remove small disconnected regions and holes."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"id": "072e25b8",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
"<a target=\"_blank\" href=\"https://colab.research.google.com/github/facebookresearch/segment-anything/blob/main/notebooks/automatic_mask_generator_example.ipynb\">\n",
|
||||
" <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
|
||||
"</a>\n"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from IPython.display import display, HTML\n",
|
||||
"display(HTML(\n",
|
||||
"\"\"\"\n",
|
||||
"<a target=\"_blank\" href=\"https://colab.research.google.com/github/facebookresearch/segment-anything/blob/main/notebooks/automatic_mask_generator_example.ipynb\">\n",
|
||||
" <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
|
||||
"</a>\n",
|
||||
"\"\"\"\n",
|
||||
"))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "c0b71431",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Environment Set-up"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "47e5a78f",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"If running locally using jupyter, first install `segment_anything` in your environment using the [installation instructions](https://github.com/facebookresearch/segment-anything#installation) in the repository. If running from Google collab, set `using_collab=True` below and run the cell."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"id": "4fe300fb",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"using_collab = False"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"id": "0685a2f5",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"if using_collab:\n",
|
||||
" import torch\n",
|
||||
" import torchvision\n",
|
||||
" print(\"PyTorch version:\", torch.__version__)\n",
|
||||
" print(\"Torchvision version:\", torchvision.__version__)\n",
|
||||
" print(\"CUDA is available:\", torch.cuda.is_available())\n",
|
||||
" import sys\n",
|
||||
" !{sys.executable} -m pip install opencv-python matplotlib\n",
|
||||
" !{sys.executable} -m pip install 'git+https://github.com/facebookresearch/segment-anything.git'\n",
|
||||
" \n",
|
||||
" !mkdir images\n",
|
||||
" !wget -P images https://github.com/facebookresearch/segment-anything/blob/main/notebooks/images/dog.jpg\n",
|
||||
" \n",
|
||||
" !wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "fd2bc687",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Set-up"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"id": "560725a2",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import numpy as np\n",
|
||||
"import torch\n",
|
||||
"import matplotlib.pyplot as plt\n",
|
||||
"import cv2"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"id": "74b6e5f0",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def show_anns(anns):\n",
|
||||
" if len(anns) == 0:\n",
|
||||
" return\n",
|
||||
" sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)\n",
|
||||
" ax = plt.gca()\n",
|
||||
" ax.set_autoscale_on(False)\n",
|
||||
" polygons = []\n",
|
||||
" color = []\n",
|
||||
" for ann in sorted_anns:\n",
|
||||
" m = ann['segmentation']\n",
|
||||
" img = np.ones((m.shape[0], m.shape[1], 3))\n",
|
||||
" color_mask = np.random.random((1, 3)).tolist()[0]\n",
|
||||
" for i in range(3):\n",
|
||||
" img[:,:,i] = color_mask[i]\n",
|
||||
" ax.imshow(np.dstack((img, m*0.35)))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "27c41445",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Example image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"id": "ad354922",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"image = cv2.imread('images/dog.jpg')\n",
|
||||
"image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"id": "e0ac8c67",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"image/png": "iVBORw0KGgoAAAANSUhEUgAABiIAAAQeCAYAAABVBSJEAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOz93ZbkOrImiH0GgKS7R2Tmzl27qk6fVk+vkWaNrqU1egjpvaUnkNa8gy66T1ftnRnuJGBzYWaAgaRHeERG7r9yy+XJcDoJgoDBYP9GzMy4wx3ucIc73OEOd7jDHe5whzvc4Q53uMMd7nCHO9zhDnf4DhB+6w7c4Q53uMMd7nCHO9zhDne4wx3ucIc73OEOd7jDHe5whz8v3A0Rd7jDHe5whzvc4Q53uMMd7nCHO9zhDne4wx3ucIc73OG7wd0QcYc73OEOd7jDHe5whzvc4Q53uMMd7nCHO9zhDne4wx2+G9wNEXe4wx3ucIc73OEOd7jDHe5whzvc4Q53uMMd7nCHO9zhu8HdEHGHO9zhDne4wx3ucIc73OEOd7jDHe5whzvc4Q53uMMdvhvcDRF3uMMd7nCHO9zhDne4wx3ucIc73OEOd7jDHe5whzvc4bvB3RBxhzvc4Q53uMMd7nCHO9zhDne4wx3ucIc73OEOd7jDHb4b3A0Rd7jDHe5whzvc4Q53uMMd7nCHO9zhDne4wx3ucIc73OG7Qbr1wv/X//ZfN+eYGUDxZ+pfpB8AKAwUZsg/u1cuZ7Z29D4ixJhAFIBAABGIqP5en8ksj9N7mblrp/WodP3yUNvldY/bs7iUeq3vx3Ng/Sic9Y0DwBEMQmHSsSggYhCxPpbB2cZS7EOZL2AGygJwAXIGgAiiCXEYMRw+YJhOmB5/wHR6wOnjDzicTjicHvAwPeBhfMDhcMA0TTgcjxinEcMwIKWElBJijEgpIYSAGCOICCklgAAOaLO1Gp/tOOhk1uGket7Ggu0aLm3yFSOYM5gZpSyGECilYJ4vKCXjks8oJWPJZ3DJKHlGKRklLyjLgrLMyEtGyRklX/T3glIKGAWsz+Qi7dq8Mpf63Y6ZF3BhLEvRPswozMiloBRGXgpyKVgKI+dcn1OY5dXKFtcIQAAp2rL0sxScz2eUUrBcpL9fL4ylZDydn5BLxnm+IOeCy2XGkgsuc67DFwNhSIRxHDBOCQ/HCafjiGkMmMaIYwLGoONdGHkBSgFmTmCKGE4fEIcRp48/YJyO+Pzjf8Y4Tnh4eMCQEo6HE4aUcDocEWPCOEwIIYCi4EoIASAS1CgBXAKWsmDJCxbOyCUj63HR98056/gQCIQQFOdIcHAcR4QQME3yrGEYEGPsjyEihlBxkFIExbBZn0YPPE1Y0weus7PCZ6Ut7GjLHlTcwvZa7uiSXMfI/fXZfhf8Bxs+zchlxrI8IecFS74gLzMW96GSQaXhH5esOK3vjaJrrR+LnBc5IgptKf16BApKyXJdKeCi/Udrg0oBMSMFIEXCIRLGGHCIwBAJUwJiJIxDQIyEYYoIMSCEgECEAwYEIsEnIqE9ils2jxQCKEQQRQj9hK7PomPFrW/s1rCOQV172l4ibTsO8l3Qt5t9ec/ixqogF8blcsHlcsHX86yfJzxdzvh6OeO8zLjodSDZt4ZhQgwJx+mElBI+fHhAGgY8HB4RY8I0TUJ7xwExJQyHI2IakKYTYhyR0hExTkjpAAoJoIiSRnAcQAiyNxKBgvaeCAyyrVDmtZt32aV7nNzul1QJ+ArvqcD2+fW+zbquSylYFqHTl1mPlzNKXjBf/inHWej2fH4SHFuEvi+Lo+dG03Ve8zIjL0ulH8uSlSZesCwLLk9n5GXB09MTlmXB+fxVaM6yoJQMxgUERgqEFAjHMWIIAYckuJgICIEQA9k2jHkRmnWeZ1wWob3LknFeMual4Lww5lxwXgqWDMwZyIWwFNnfbTRHwzGiSittzeXCWHKpfFGbAUHKEAJCIKQUkVLEOIwYUkKKJPTP5ouMDTKcleMFC5Zc8PMvZ+Sl4HwuyLng6TJjHBM+fjzi8+fP+J/+y3/Bxw8/4C9/+Rsejh/x4fQDHh4ecDo9YJomjOOIw+mIcZowTGPlHRr9p4pbOQsdW5ZL2ztL0fktMic5Y1kuOpeLo2kLcha6k8sCzoYPWeYxy3wuF6WJikt5FlxaZvl9Xs4oXGpf7OjB6E7/sfP9bzIP7UNyIShQt+fUvce+O5pSSq6L0mhq5Rv0s+S50mbD9WxHx2ewp9WNzQLc+l//TkQdnSN3vl6/Obe+Fqvz7b0j7f++P87r9re/7fLZAUrq1m2ESrfW97ej3bvt40s8/V5/v+V6//0WceL6uGzPv9S39b743PGWNtb3PTeeezzUtWvXfen6HQEwUJYEAiPGRR8gMmZGRmZgYeD//H/5v+H//r/9P1F4QCkD/vf/7/8b//v/7/+Dpy9nnJ8uIBIZ7MPjgGlKmAbha7nScF1HUWVIL+reCHvv8NK1b/39tdd9S1u3vM9L110bm628/zxc7+N2wq7h+zVe/3rb6xOA5yYzMopKwuEKDbhKG0y38Q7w0vh96zyu36VQQeGCyBHR/FyJwfpp9LrI6gpHMI1Y0n9CoRPm4a84hyP+/4e/4hIHnNMJREDUxVcogAMDxIjESAACJxASxlAwBkZwdIgbM1y/Z84oXEQu9fKYk909jywyknxY27KZVo4Ny7w4HqaXyxoPlFHKovoIaJ/k+aPDHdE3ZNUBzMoPLBUvyN5FeYmlzCicpSlTdzEws/AQT8uCJWfMXBxvVICs7zU/ASWD5gVUCqDyJfJF5coZXDLATT+DInoFKjOozE2e4CZ31o7Y9da5qgTcx7X1kpN1vHMxJYASPj4+4vHhAZ8+POJ4OGCehZ/8xz//gfPljJ9/+SfmRd7BdEG1bT2mmPDh8SPGccTj4yNyzlXu+/LlC4riTNvnAkCh6tCGYURKQ9WvNX5R1oDhw9evX7EsS+WNbe2M41ivy3nB5XKp9xgvVWXd0tM1r7+y9rZjStovkSWY4fjzRWWQ4uigjHlb93IMFEEIEElZdZhg6JKsV+YQnR4QiIBcJ0pMFADDww/4y//1/4H04S+Y/u1/FR68XLBcnnD+8k/MJeDMEZkJGUDIZ8RyBrGwgxEAkekQC5Yi+F4Uz6E6jHyRNZcNF6Fy6zwLXi5zG6eKf7reTeeBDLDoWFAyiBnEBeAs+G4Kth00lXXhn+Fokh573myNnbfQabfWduBb9/G3Qim3MU03GyJ6+HYmZO/a5xgoItpRnGx7sm+M2Gcu+mvp6iS2dq70bef5cvQ95E5B366XxdR/ty/UGBwCKIpiLsYJcZgwHo4YpiMOxyOmwxGH6YBpnDCNkyguBiWMw4A0pM4AYcpkL2CTauaICEwMgo45aV+wRegqGJNukvrKxW9EYN1DbZPSc7oxlWKGCBHYzUiQVfmQlwWFTTkq37k4RYV+cs7gbMJ9VuJsxKQpAMwY0BkickFBQdb7cjVELMq0qCEiqyEimwJBlOuFS2eIsNkHgEhBlZ4BpAJTwxNjehRP9bjBWf+Hb7zb1RlXV5BS2EABTEEU+tHjQ6hK/lCPspEGaoqX4BhPQ02P4fWlKm7YoxW/AkAcnKJCnmvGsBDkE2PQv7eKIFJFHMjYQDcE5AcJXT9a1+y7vsN6ZVcm8bqAUmfpynXczbEavcwgBtmwSjbczKrUzRu8z2UR5Zwp6fQYSgFxw3NpqzEppRo32kbnFXNZ+cGcM+oGpn1jVuaExThqG729byi6CUcVuiggg1ECoRRCYYAKgZlsqTskocaxdIjc6ITRetJ1a/NZleA7UybosNU2dYYNIgTD42AKQ79mmjBg/QD6fakq5iqaCC3JtuFSABBQAiPFCxiMyzyDAVzijMSMmKK8eYlAYURj4phUCWKfAIYyvXAP9sfVWLy0Qz6/JzchobtOp8nWDuvCN6PPWqnafThXem74WexYr2nrowk0bf3sKbvavK6/Bxm1yryv7l3PNjX6hPo81GP37KrNNLxVvNRB4oqnBKrjBHcdObrgmMZGjmR5XNljjacQxphUsN9SfL/10qqNGAgpRjFgK38wDAOGNDgjQ3T0WfBvPXbs/3IPlHlu9G/zgf/u5lv7WEdzJQC1Ud724Cp4FHb85TVes98bep7T99PmttIpu46oKQmudmqPa1WFhmsPOwro1h9PQ1uTHY9ccZhq+1dpAe1h0a3Q+Bj/vKtXPyMsf1/w6/E6/Pr9+n3BLQrH5/aQ7zu/be9pfJ6c36WDLI4X5iBg/KLQsyB+bsHRBf9/3Te+scfvKOjfriTv52DPAPVS/15jQHlveEkXsHe9h3bvPq31z/h+76e842tu+RXH+rvOax12/YOhvIo+u/6/5cVgPNbOwFUxQtur/ChY+Wfdp0l3O8dDXMVz2c2fwZQbYWe/tj7vvYN/vY1IvwOmhhEZrR2fnUdeHZ+7xn2tY1FlgZ3O3AJ89cvNcPX9bEAIVbbrDeGO17yhD2x38DVcoc31ldfjZrCyI5HIRwgOL+xzZS9tfHF77zXfym7+93nY/vyWll7hg7jx2O19r1zojuwQZnNHd4I7vKGYENKAEBNCjE7y1htNlmdScuBkPHetfMtgonZf1xcbf5XvnJy3C1XEo9XbN/mHybdFSkOonq9kDK/fX0zn8K8GbzREAO8xWBWlX5isrfcCKbK8zEysFbrXL9xuBSZcWl+vbRTr91gzfoXN/q/n7C4WJWBVXrNvh8Al6bsmUIxIwwFxmHB8/AHT4QEffvgJh8MJjx9+wDQdcDw9YhxGjNMBh+mIw+GIaZwwjEMXCeGVutXLTzf/ECxbV3FjR/r+2/GRqAOGN8UzN+cCdsrN5rGtildeqrAgv4vSKqv3gHkRzBYRUcRCvywzULJ6TWaURSyfJWdwWcA8d54JfpNo3rtqLV1M+SuKsMsyi+Fh0evmjFwYc85qiFiQM2NWz9xcDRpNISeLine truncated
|
||||
"text/plain": [
|
||||
"<Figure size 2000x2000 with 1 Axes>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"plt.figure(figsize=(20,20))\n",
|
||||
"plt.imshow(image)\n",
|
||||
"plt.axis('off')\n",
|
||||
"plt.show()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "b8c2824a",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Automatic mask generation"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "d9ef74c5",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"To run automatic mask generation, provide a SAM model to the `SamAutomaticMaskGenerator` class. Set the path below to the SAM checkpoint. Running on CUDA and with the default model is recommended."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"id": "17ade22d",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"sam_checkpoint = \"sam_vit_h_4b8939.pth\"\n",
|
||||
"\n",
|
||||
"device = \"cuda\"\n",
|
||||
"model_type = \"default\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 10,
|
||||
"id": "1848a108",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import sys\n",
|
||||
"sys.path.append(\"..\")\n",
|
||||
"from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor\n",
|
||||
"\n",
|
||||
"sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)\n",
|
||||
"sam.to(device=device)\n",
|
||||
"\n",
|
||||
"mask_generator = SamAutomaticMaskGenerator(sam)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "d6b1ea21",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"To generate masks, just run `generate` on an image."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 11,
|
||||
"id": "391771c1",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"masks = mask_generator.generate(image)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "e36a1a39",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Mask generation returns a list over masks, where each mask is a dictionary containing various data about the mask. These keys are:\n",
|
||||
"* `segmentation` : the mask\n",
|
||||
"* `area` : the area of the mask in pixels\n",
|
||||
"* `bbox` : the boundary box of the mask in XYWH format\n",
|
||||
"* `predicted_iou` : the model's own prediction for the quality of the mask\n",
|
||||
"* `point_coords` : the sampled input point that generated this mask\n",
|
||||
"* `stability_score` : an additional measure of mask quality\n",
|
||||
"* `crop_box` : the crop of the image used to generate this mask in XYWH format"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 13,
|
||||
"id": "4fae8d66",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"dict_keys(['segmentation', 'area', 'bbox', 'predicted_iou', 'point_coords', 'stability_score', 'crop_box'])\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"print(len(masks))\n",
|
||||
"print(masks[0].keys())"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "53009a1f",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Show all the masks overlayed on the image."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 28,
|
||||
"id": "77ac29c5",
|
||||
"metadata": {
|
||||
"scrolled": false
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"image/png": "iVBORw0KGgoAAAANSUhEUgAABiIAAAQeCAYAAABVBSJEAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOz9eZNkSZIfiP3U7B1+xJVHVdbV3dPdg8FggIFgQYJLQoQU4VfhP/wA+/lIkaVQlksuuQBkd4C50D0z3V1nV+UZEX68Z8o/zNRMzd57Hu6REZmRVa5Vke7+nt2mpqaHqRoxM+MIRzjCEY5whCMc4QhHOMIRjnCEIxzhCEc4whGOcIQjHOEewLzvBhzhCEc4whGOcIQjHOEIRzjCEY5whCMc4QhHOMIRjnCEHy8cDRFHOMIRjnCEIxzhCEc4whGOcIQjHOEIRzjCEY5whCMc4d7gaIg4whGOcIQjHOEIRzjCEY5whCMc4QhHOMIRjnCEIxzhCPcGR0PEEY5whCMc4QhHOMIRjnCEIxzhCEc4whGOcIQjHOEIR7g3OBoijnCEIxzhCEc4whGOcIQjHOEIRzjCEY5whCMc4QhHOMK9wdEQcYQjHOEIRzjCEY5whCMc4QhHOMIRjnCEIxzhCEc4whHuDY6GiCMc4QhHOMIRjnCEIxzhCEc4whGOcIQjHOEIRzjCEY5wb3A0RBzhCEc4whGOcIQjHOEIRzjCEY5whCMc4QhHOMIRjnCEe4Nq34T/3f/1/zJ8yPtXxCExpwf5ewY2WKBnCyLyD+nGQiNY2qLmFWiQZ0cjqUxSZj6gg6Mg+QlggEP5eiQo/sNgVR0z4apzcOzHxm0drl93uH7zHF/901/BVBb17BR1u0B7coF2scTi7AKn5x/h9OxjtHWLtmpxcX2G09UctqpgrEF1/Qq2W4MIIBDIEACCiWNOO8Z914TosRoM7O7v/n8wu6xExw5wDIce4PAbfkCY2eOUY4D9OAEO7Dil0SMtg8sBFzm8Df9ISubQnvDesZPiIHMk73xan9OpsvJhYVy/uMRqvcF/nf89XlQv4Poezjms12s459BttnDO4XrD6FyP1XqF3vVYbzfoe4fNZouud9hse7DzVVpDqCtC09Ro2grLeYvFvEHbGLSNxbwCGuPHBo7Rd4BzwJYrMFnUi1PYusHi7AJNO8ejx5+jaVosl0vUVYX5bIG6qrCYzWFthaZuYYwBWQNrLYwxAJHvrTNgZ9C5Dl3foeMevevRh88u9Lfv+zA/BALBGL/WK6pgjEHTNDDGoG1btJc1nvzTEtXVNaoXX4EAXz9R+BNUM/47EbLFz2qOJ2DyHZcJxlMmDEOGX/pLwj2O/8UXEac4lsFwYOfAcHCuBzsHxw7M4Tc7OBeQIOKmWgsQbE64yQovnfNrzAlN4rKzYW0JogFYVQ2u5wv88eQ7fHv2Fcg5EDMqA1SW8Ok3C5y/XsASYIhgDWAIMJZgQCDr6YyfJoKFgRHCR4H2xDn18+sTG4RcsR8MF5vLLP0Mo6vWYwZEweJOIGPg6V0sNlGrkfKYAed69L1DF/869H2PzvXonYNjod0GIAS8NqhsBUMGdVPDGIPK1jCGYG0FIoKxBt2yw/Uvr7C8eISPfvXPUFUtqmoOa1tU1QxkKoAsXNWAbQ379wTzewp9kdb7dcjOYfOH38BdvUmkiDmhGjQ+6LmXccjpdnrL2a+IXpLOyT7VgR1HfO37HswM128i3vrfnR9fF94HfBa88/jv53LdzrCqZ3DMcNzjq7Pf44fZ99hsNui6DpvVGn3XYbVaoes6rNfX6KkDPnHgBgB6AB4fiYDKGhgQKo+A8Tmp3vfOt6N3/s+3kcNc+/f60/llCMeUoZ+NOJbvmxGv4twkPiGStbAeiAjGGFhjYIwJ7RVshrAWQ5yFb+t224Fd6NPrHt1/vca6Y7whg0ePHuHnP/sZzk4v8OTJx3j6F7/CyZ/9Ceq6QlXVsNb6v8A7GGthyMAY8m2ghBV+6QXaEsesD2RKaBin54InzuOFvHcy9y6lZ+cge7vQRGa/7uA6P45lnewST1GShLBPEFE2PSR7icwFSZrEJ8leQyj4pILxTCuJ414kAxV5CGbQ869gvv8nfPf7l/jj71/GMeh7T1v6vkcf9k8C40//mz+DeTLDu4YhX63exX/Kp5SnGStnPPPOyuiQ9GN1Dlp1SGFv9foWaadSDfn0A8Sl8ek6uJR90h8OO0skRFlKS1MA0L5cAd/8Dovzczz57FN88skv8Se/+jdgtmC2+MPv/hZ/+P3fYrvpsN12ftfkHtXffY36q0t8/b/5HG+ezYou8n108VbwfpoxVesuzvpDgjFZ9W3L0nv6rlTvYkbvZp6mSpnSmFD2doxmMEwHNC8cuJmjqz6FowW29UdYmzm+nX2E9upL/Iu/+n+g7rZ5ycRZScK9EXiE1hMaMrCUztw6kceUXMSihBjrj8gFAHrnsIVDBYMKiQ+LvApc3N8R8rL+07y0qk+fCHYEdL5QmD7xQbFHIteFf7RMpMt0gR/rmcGOwU0FJp0mZHCprfmnU78HQrFSeuR6mwH9LL/fgJLf2xr/LzqFC3N2uX2NHzZ/xPq7F7j8p29TQqoAqnB2coKT5RL//tm/wSftF+gDf7nZbNBzj8UZwza+4lLmiUURRXnNVBZg4U/ZyzFFoykpICI/LPwjGRO4H6WTYIYDw/UuyD4uDrPk8WPqccQ5P67MUDwPK3zMYaJbGZgo85hQFYd+jmUWWXGs36ot4EQDtIg4wg/3HePVtxs4foN+9ld4yYS/aX8NLJ+g/eTPQAQYt0FvVuDZNXoYbNlgebLB2fkGxL3/k/5klXKQETl+j/qRIF+6bKA4PofL8ZeCPjbJuCLwal6ewzqUNaTWYDmSnK8RiumSXKCHeYxfncJblWKy/qkyy3oBgC9fAc+/v7Gqhf3fo3af4JL+P+j4y8mk//P/7YayAuxtiLgdHLAJhgU3DRkFvbs6b5t3L5hqM2ffqEjDTHBZKoJjh+32Gp3boGoamKpGM5ujbueYzeeYzZdYzE/QtjNUVY0GDRrXoCYbFBgMy32ueIlGiLzJfks/ROAo045sPqlz8XlSYpT9D4vWiTIh/pM25ahoFobChX2zfJ4WaFRWqt+R4LC0iWN9DhwU58JoIFNucGy75BmbbdltGFVXoeE2KF96bHijlBKJGcoU1gEyOTFbK5pB4OlZCgoUQwZMBtZYWFuhqqqgcArPjIGJn35zNWS8IpmSwYrECKGVUXGs801MFNBE5HW1bKKCjUJbjAkGDkdoXm/RrBjUNSDukswtSiFFVbP+3khDkOPErsScfxmb14wpGEvPOq1mggTnNI4GZhFBGRuVay4pZ4NCVJRSJGVEPC2NGgrXi/XDMKGNahNWeZxTRg7TgzcbuO0G266DcQxiFzStiAwbG4ULsh6EiZK5YSpwl4rv6hlLa3I6mo9tAUTjuzpp5jBVkanLSE+ZpoLj+CYzzZqZC4uCCHDkAAP0zivDHfUADMg4EAySUjLx9MzSIgJgwAhK3w35RF2opt+Ct71qEfk5K5iqEnbzMzzybSSNLHwg9N0X7MliUoZ7epaMx4lmwzPcQtcjbU7GL3lOzsG4PmAFoXI12m7mjZzUR7oU/wJNYXihdIKKDn6Rwpv4lUp0Gt/PbxxSGj4afSDLZIJr1HSWBlzDeIbY4rCWe8dwCAbsukYtf1XtDQ1h/KDorKwbrQDmop74Ncwv63mVlaIE2gGelOWM1zRe6Y7+66808Us/9uOv0jADRN5oQH59EnHER9XxjNufbl2BQ7YG13OQvfJ7KnMwxKrxLz4fJOyxpQJpLT2MruzBMPzkYHJjVZ97rL37ggOq9jqGHiALkEn8o/znmcp8Hcd1H34+IPQou/5umvYe5/qdwI+9fx8GBPHUf3MOs8vXaK+vDsfxAUn3e7VWtkYZRyvQJ5iyXIegdBaqHF+WlDn2PjM/jO4641hIYBCYgiAlbVLKbc6T5wUSIeozzMRIvjP0378iC+CcHFbMuILXP9RUoyOriuMoJpE6oKjro4phHQ4ilPlM7UqnJFcW/UTiGZkCh54YysEOWoqq+XxOyBn7MP3veNNi+K7fxBkQgKoGNj3h9ZZxTQRjKsDaMJUMWIa1DNQMgtdBWMvZ2OWfIqWH2vOPmJZlsAeDPtLIlCt7EfUB2QSm95m2QKs1DoD3zhcbA9R1arfrvU7BGMBYkGtA3IBQ+/eVBbgG+u4m5cJOuGdDxJ5wY/sfCBN8SxhjIFm90Yp4ZsKqDwf92YS5Nej7K3z//T8AAE4//hTtbInTi6eYzRY4Ob3AfH6O+ckTVLZGVTc4XS1xspqjIgu7sLCrNzDra7/gbZUaUih2k4IuKDr4BsKWWctTn1z2s2ACokU/bNdOfgXCzWI1TopZcGAc4ukAB7hUXjp5CPhTCcWJh0zhpdJJ3eGdnCIQA4QLaYWRSac0VVlqDplZESIv4phZhXlj8emLT/F4+xSGCB06/HX119jgEp3zlvKuCyesnQsWcU+YjDFwDBgTlJ5gUNj0tGGG4PkNOzJjFE5p13YGsjXmJ2eo2xlOzs7RNDMsFyeo6waL+RJVVWHWtqhshbaZwVqLuq5zZZ+x3lATNk+HHo5dOInrwNzHsSdjYDw3BACoQlusrWGNQVO3sMbXUb9Y4cn/8gJ2C1i88RucsfAeFLmCZkqZ5D8VE5ajYJGYAu/AWTJSczrJpAq+Fs9SnaKEE6VrwC12QUhOJ30jHjODISeGewDhJHn481Z/B2IHghgBPC77qpQiMLZLGdeCwUI8IuLpoOhNkU4R+OXj4NZr1HSJfvscL/gVDBwMM5rKoLKE9brGtm/A7A1dYMRTIt7Wyd4ARZ7hQFBeG0N+rZPxRpU4bC5Mi4ntl0nxSnvWh6KyCRLvCsENEW7kNIg/tTJUqQpNlrqZ/XT5U1aU5YnoF4xCXdeFk/PC63hPH2u9p0/d1zBkUNXLine truncated
|
||||
"text/plain": [
|
||||
"<Figure size 2000x2000 with 1 Axes>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"plt.figure(figsize=(20,20))\n",
|
||||
"plt.imshow(image)\n",
|
||||
"show_anns(masks)\n",
|
||||
"plt.axis('off')\n",
|
||||
"plt.show() "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "00b3d6b2",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Automatic mask generation options"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "183de84e",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"There are several tunable parameters in automatic mask generation that control how densely points are sampled and what the thresholds are for removing low quality or duplicate masks. Additionally, generation can be automatically run on crops of the image to get improved performance on smaller objects, and post-processing can remove stray pixels and holes. Here is an example configuration that samples more masks:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 24,
|
||||
"id": "68364513",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"mask_generator_2 = SamAutomaticMaskGenerator(\n",
|
||||
" model=sam,\n",
|
||||
" points_per_side=32,\n",
|
||||
" pred_iou_thresh=0.86,\n",
|
||||
" stability_score_thresh=0.92,\n",
|
||||
" crop_n_layers=1,\n",
|
||||
" crop_n_points_downscale_factor=2,\n",
|
||||
" min_mask_region_area=100, # Requires open-cv to run post-processing\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 25,
|
||||
"id": "bebcdaf1",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"masks2 = mask_generator_2.generate(image)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 26,
|
||||
"id": "b8473f3c",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"90"
|
||||
]
|
||||
},
|
||||
"execution_count": 26,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"len(masks2)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 27,
|
||||
"id": "fb702ae3",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"image/png": "iVBORw0KGgoAAAANSUhEUgAABiIAAAQeCAYAAABVBSJEAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOz9WbMst7ImiH0OREQOa9oDN3lInlN3ruFWVXe1SWZtpne96cdIv0q/QE967ze1tVRmfau67lRnuDzk4bCHNWVGAK4HwAEHIiIz17QHMp3cKzMjMMPh8AHuIGZmHOEIRzjCEY5whCMc4QhHOMIRjnCEIxzhCEc4whGOcIQjPAGYD92AIxzhCEc4whGOcIQjHOEIRzjCEY5whCMc4QhHOMIRjvDzhaMh4ghHOMIRjnCEIxzhCEc4whGOcIQjHOEIRzjCEY5whCM8GRwNEUc4whGOcIQjHOEIRzjCEY5whCMc4QhHOMIRjnCEIxzhyeBoiDjCEY5whCMc4QhHOMIRjnCEIxzhCEc4whGOcIQjHOEITwZHQ8QRjnCEIxzhCEc4whGOcIQjHOEIRzjCEY5whCMc4QhHeDI4GiKOcIQjHOEIRzjCEY5whCMc4QhHOMIRjnCEIxzhCEc4wpPB0RBxhCMc4QhHOMIRjnCEIxzhCEc4whGOcIQjHOEIRzjCEZ4MjoaIIxzhCEc4whGOcIQjHOEIRzjCEY5whCMc4QhHOMIRjvBk0Bya8P/xf/2/j54xGABPpqci3USq4iED5GF+/RNo3cfMuQRShbHOo4t7u4L/5qKqeTrtqOCZ1+G5vJwqdx+wKpqK7jIAIq6TJxjA+AkbOHiwBww7tP4S17ev8bvv/wtM06BdnqFdrLE4fYbF+gTr82dY2q/R8a+xaBdYNAu0bYe2bdB2HWxjcftsi345wBgDIkqfRASAYEzspymnR48BTQ4Fq2808Sx+Z67GFQB8fOTTO88M7z3ADMcDmBmeHcAMZg+WT+8B78Fennv1nnPNsV4GBxxiH5sj6eJ3MMAM78Nz70PbJB370Fqffue8Un45LIztT38H/PQt4EM67xy899hsNvDeY9j28N7jZssYvMPt5hbOO2z6LZzz2G57DM5j27swRAxYQ2gbQte16BYNTlYLrFcdFp3BorNYNUBnOPTTM9wAeA/03IDJol2fwbYd1ufP0C1WeP7ia3TdAicnJ2ibBqvlGm3TYL1cwdoGXbsIuGINrLUwxgAUcdobsDcY/IDBDRjYwXkHFz+H2F/nXBwvAoFgjAURoaEGxhh0XQdjDBaLUFfbttist/jx1VuQpVAvGRgyGQetARkTSQbB/v4C9od1mncupmNqoc+sa5V5jjxoIsbM6L/6CcOL6zTvNU4xfMI1gBM+MCLes4dnD2YH7x287+G9g/MDvAtj6d0A5xyIPYh9wlNmD7APaM6c28a6DYD3LuBgtEF7rwlxXC/sQzof2t1836P7XZ/WAHkPYkZjgMYSXuEZzrBGYwBDBEOAMYAxFH5bijQmzFEDAwIBhPAcBJjwXp4DBDImzg+lMQ19QRpbqDXORV/KOTYIdROZ1JbxzGeaIJ+BvDg47zEMHoP3GOIcDM7Bs4fjSFNgAAKMsTAwaJoGRAZd14CMRWvDb9s0gfba0BZrG/RLws1zi4uzc3z9+VdomiWaZgkyDUAWvunAtgXBgCisPTJxbIjAIPSDxz/8yzvcXPm0x0DoXUZJeboLscsMlMcmkXF5neaE4X2gz845MBjODYGWui2YGS7Sc++GiGMeQKSxkWbfnnyHd89+m34vvunRfbNN9GMYXKSJWwzDgO3tBm4YcHt7iy0c3v6nLzBcdLHMsLYIHMabgCbipDUBJ0jhoOCD9x4+ttd5B2aEefYBH5zn+J7hOdBVBsEXQ0qwceykZI23Hhy3u3oS8h4b2hzm2ZpAc42BaikyErPMSfjiEPrQ96H9zvn06X7/Gtv/5R/w4tlz/Kvf/Aa/tqf4S7cCLb4Gul+j/WyF5mUX8NgQjLUgMmEdI7Sn7pPGqYwPXOAHfOh3uYcLvWHV/rimPQPwGc/AYF/t76osog3axT/D0DbRiUmaQDLG0guhReG77CVCi0yYjJyeFI2qmSGqpoUBZg/3ww1u/+sPgI/zzmF+Xq8J358SBtfDuSHxHM4FmuOcgx8cfrMxOHeM01cruOV9eNFxGw95Nc3rHViFykz1wEzmuUPDDsu1Oy/termn2gNhlP++AzpRGM29eADsK2maWsVvj9eMsnCeekbhxSgNoz85wfazz2HbDl23Shm32w2GfgOvZBACsPjhT+iurgIPchf58aOFx56Ih8CusZuaYP18X/6p9Hete1+e+45l3tnHJcyUOTccTwqPVVluPIPjOpqax5zOwuDMvIJfneDbz7/GtjtD377Cxqzw3fIVTu0Gf3vyLRryZX4CEHUnJpUZ+Dk5WZtUOxP6oiQLceYNRioJlXlOd0Hqb9ZHJGkkZufEA4RqvapHc4t11Zz4zlCOn21foGdePc8fzAwX80dVRe639CXy4WAP8uEZxTJDGT7xZHXdgAd51qOoysbE78Px7erNBt/+97eQht/c3uDt1WVOQBaAwWq5xHKxwMUaWLWBf/I+yAjee2z7bZQ1Mi9ZAxmDru1gjUXTtYD3cN7DO4d+GPK8JTbGRJ7RRL7cVN8p8ZPSb8+AG4bUPlZ4ZIxN4+qjbJR1U8JnZp1WAcxxiHZQHELUmwR+Pky7VzourQPIn/X6ET1heCLtj58a76hc/5LDxLn0ANgu4V/+DezyHM2zLwNb7QeYYYOF34TxYgOO6ck7EByIs1aASK/njK9hOETe9Lm1apy88zGNWjtFq73qf1wjIk8wpz5B1l+iDyXkHkv5SVhLCUpNcZmz/LYLpusH5tnOiaXwqPD//H//vw5Kd7AhooTHb/19SqSpfHfhW4pZoEfu1v0KYzAcAJcWTlhynnts3C0G7tG0C5i2RbdcoV2ssFytsFicYNmcY2FO0PIKLXVoqUNDLSzaoLyyFoYbWEcgjko5H4ipLOqg5CpobgWjLXMyRbmEOPYj/uaydKEN+dOH9GnntHGRh39pM2EAngAmkGeATSYsShlLjKSYCApcwLNJbSGELKQIfdhcwyZNsU4SRSwDJiophVhRKj/33NMATx6MFq5dALe3oKjkysrOzHQkRU78b3LUZXDTA71x8MycIVFtQwZMBtZYWNugaRpYa2FtfGYMTPqMmyuZoHyKijEgbEgcy80tEMIqfWNVddxMDUCsjWChXkMG1i9AZODIgqMiyLOF2bYgEwwgkk86RYaiAImAB4Mor/X43XEtqk1i31bAOt1gQTc2PU9sAYsBwkQ8iXl8NByAFF5T3PMIcD7/9gR4AjkCOQMTcdxwZiZkARFkMw74RIrxpKgsJmHhvSiYGUwOrtmE8lzIz3EdcOOA3gNDYCaIPWADLngTqFXah6kirayUFGGRyMhUg07IQgyqQop9e3I6iADmeu5jaVHJQFGJSEUKKTfkJaI0r3U7qc7DYrBk9SKsDfIexgDOB7OPJwMy4TkZE2iMJXjL8JawaFo0pkFms0wYDTLwg4HvTVKa99bDkU9pGYDT7XgolFaLYklxlUbwRCudtWE207NojNMUQ62HwOg1sJtFzOFhYIG2AbYOlj3AwQg3wIPIJbok45LqhjDNYT5I/dWQUmi0Kb4LMk+p3nVCztQ3IihTQHVWXGBJVea4xqmHsp/Rjr25KpVVCznQG1F2d12Lrm3Rti3adgljTmGaDtxYUBPHSVVSClXzFDUZAHQjil4LHmQyWySVVCMOeT9eMwzYd/Bx/TJ7ANsorFRb5py+UdIyo8iY5rRUVQYDCJVlTLW2MTDrFuw48hGBftsVoVkSeOOTkZiU4QNQePwQJfbHAlOk/0NBsRcJfAwN+3Awt8r0vvdIO8wjwJ7WJHpSK0UoUqD5/B8Tmt4d6tZP8zLjdxo+ZM/37XBT6TUckveuddwVIi95lyzvdWF9+FXMAAYa4P0WYJ95DALaxqExfj7jFBqTcGGRUysZixnI8/TwEZGyhL484ijPsEM7yz+kch7/HI3FqJz3gzvGGiwWNhpOgK1TAVyYS200TXH3U3zkLsiHBscw1mqRwrn8IPO3pOVe1bbEzmF2WiebUPRF+NMD+jQaldkhmZNtdqeVru/OUab3MGCyILL5IGng4GHySBWtIRK9ZVkyRVkLI9k/vA3yfKz5gIUusmCZjGLW8QqRVR+qZ1EBBDpwR00/HdLAnyHc0xABPOZg3aUk0QtQmv59+Q8sfTbZfZiVEkl9eqYXlsJ0ypuWB+MdBgwAfDRCEBls/ICfrn4PsoyzL77EYnmCs2efYblc4/TsGRq8hHVfYeEusBieRQVzA0vhdO6WHBwGtNeE9lZaohY1pSexLZmgzo8FqwXL9eNiLDSBz6dx4+kJUfTHNHJ6IBkT2AEM+HCEPCq5Si+KrPgSyyQXln8vLUonL1U+ZCWqZzmlG9vn48kE6asYDeT0OdTpY1Emxzl923yHjb0Cui9Az17i9rv/HXz9UzjlCkonagcfFL/DEE5BBos5p2kxxsAzYExUeoKjN40oXDgRQEOAnZgpMuFUa2uXINtidXqOdrHE6fkFum6Jk/Up2rbDenWCpmmwXCzQ2AaLbglrLdq2zQo/Y0DGBpViVDx5hNPhHvlUv8w9GQMTFfEA0MS2WNvCGoOuXcC6Dt3bz+ANcNXegA2w9QbUA+ur82iAkPWgdvIaL70woaTwrETP8me9OfEoUWYOeJRHG1/Mn9ag79fxhShkgUQBfD4VwxCLfFTOipENwcsnGLrCSQwrHj/egaPXBGTdyCljMZwJjsY2SetlLYgHj5ze1utg077D27O/B2NQJzI8YBn8JYP+uIH50y0MQvu6xqCxhIvVGiu7ANjAEAA2hVHTUDgFTPGkiBjdjLSOAPIAmzSUoX/FqYpAq7xeYxMTKzgiIKeYDRlJAEI4HT+Line truncated
|
||||
"text/plain": [
|
||||
"<Figure size 2000x2000 with 1 Axes>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"plt.figure(figsize=(20,20))\n",
|
||||
"plt.imshow(image)\n",
|
||||
"show_anns(masks2)\n",
|
||||
"plt.axis('off')\n",
|
||||
"plt.show() "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "8c937160",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": []
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"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.10.10"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 98 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 164 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 265 KiB |
@@ -0,0 +1,774 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "901c8ef3",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Copyright (c) Meta Platforms, Inc. and affiliates."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "1662bb7c",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Produces masks from prompts using an ONNX model"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "7fcc21a0",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"SAM's prompt encoder and mask decoder are very lightweight, which allows for efficient computation of a mask given user input. This notebook shows an example of how to export and use this lightweight component of the model in ONNX format, allowing it to run on a variety of platforms that support an ONNX runtime."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"id": "86daff77",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
"<a target=\"_blank\" href=\"https://colab.research.google.com/github/facebookresearch/segment-anything/blob/main/notebooks/onnx_model_example.ipynb\">\n",
|
||||
" <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
|
||||
"</a>\n"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from IPython.display import display, HTML\n",
|
||||
"display(HTML(\n",
|
||||
"\"\"\"\n",
|
||||
"<a target=\"_blank\" href=\"https://colab.research.google.com/github/facebookresearch/segment-anything/blob/main/notebooks/onnx_model_example.ipynb\">\n",
|
||||
" <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
|
||||
"</a>\n",
|
||||
"\"\"\"\n",
|
||||
"))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "55ae4e00",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Environment Set-up"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "109a5cc2",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"If running locally using jupyter, first install `segment_anything` in your environment using the [installation instructions](https://github.com/facebookresearch/segment-anything#installation) in the repository. The latest stable versions of PyTorch and ONNX are recommended for this notebook. If running from Google collab, set `using_collab=True` below and run the cell."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"id": "39b99fc4",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"using_collab = False"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"id": "296a69be",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"if using_collab:\n",
|
||||
" import torch\n",
|
||||
" import torchvision\n",
|
||||
" print(\"PyTorch version:\", torch.__version__)\n",
|
||||
" print(\"Torchvision version:\", torchvision.__version__)\n",
|
||||
" print(\"CUDA is available:\", torch.cuda.is_available())\n",
|
||||
" import sys\n",
|
||||
" !{sys.executable} -m pip install opencv-python matplotlib onnx onnxruntime\n",
|
||||
" !{sys.executable} -m pip install 'git+https://github.com/facebookresearch/segment-anything.git'\n",
|
||||
" \n",
|
||||
" !mkdir images\n",
|
||||
" !wget -P images https://github.com/facebookresearch/segment-anything/blob/main/notebooks/images/truck.jpg\n",
|
||||
" \n",
|
||||
" !wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "dc4a58be",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Set-up"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "42396e8d",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Note that this notebook requires both the `onnx` and `onnxruntime` optional dependencies, in addition to `opencv-python` and `matplotlib` for visualization."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "2c712610",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import torch\n",
|
||||
"import numpy as np\n",
|
||||
"import cv2\n",
|
||||
"import matplotlib.pyplot as plt\n",
|
||||
"from segment_anything import sam_model_registry, SamPredictor\n",
|
||||
"from segment_anything.utils.onnx import SamOnnxModel\n",
|
||||
"\n",
|
||||
"import onnxruntime\n",
|
||||
"from onnxruntime.quantization import QuantType\n",
|
||||
"from onnxruntime.quantization.quantize import quantize_dynamic"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "f29441b9",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def show_mask(mask, ax):\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)) "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "bd0f6b2b",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Export an ONNX model"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "1540f719",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Set the path below to a SAM model checkpoint, then load the model. This will be needed to both export the model and to calculate embeddings for the model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "76fc53f4",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"checkpoint = \"sam_vit_h_4b8939.pth\"\n",
|
||||
"model_type = \"default\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "11bfc8aa",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"sam = sam_model_registry[model_type](checkpoint=checkpoint)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "450c089c",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The script `segment-anything/scripts/export_onnx_model.py` can be used to export the necessary portion of SAM. Alternatively, run the following code to export an ONNX model. If you have already exported a model, set the path below and skip to the next section. Assure that the exported ONNX model aligns with the checkpoint and model type set above. This notebook expects the model was exported with the parameter `return_single_mask=True`."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "38a8add8",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"onnx_model_path = None # Set to use an already exported model, then skip to the next section."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "7da638ba",
|
||||
"metadata": {
|
||||
"scrolled": false
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import warnings\n",
|
||||
"\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",
|
||||
"\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=17,\n",
|
||||
" do_constant_folding=True,\n",
|
||||
" input_names=list(dummy_inputs.keys()),\n",
|
||||
" output_names=output_names,\n",
|
||||
" dynamic_axes=dynamic_axes,\n",
|
||||
" ) "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "c450cf1a",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"If desired, the model can additionally be quantized and optimized. We find this improves web runtime significantly for negligible change in qualitative performance. Run the next cell to quantize the model, or skip to the next section otherwise."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "235d39fe",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"onnx_model_quantized_path = \"sam_onnx_quantized_example.onnx\"\n",
|
||||
"quantize_dynamic(\n",
|
||||
" model_input=onnx_model_path,\n",
|
||||
" model_output=onnx_model_quantized_path,\n",
|
||||
" optimize_model=True,\n",
|
||||
" per_channel=False,\n",
|
||||
" reduce_range=False,\n",
|
||||
" weight_type=QuantType.QUInt8,\n",
|
||||
")\n",
|
||||
"onnx_model_path = onnx_model_quantized_path"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "927a928b",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Example Image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "6be6eb55",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"image = cv2.imread('images/truck.jpg')\n",
|
||||
"image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "b7e9a27a",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"plt.figure(figsize=(10,10))\n",
|
||||
"plt.imshow(image)\n",
|
||||
"plt.axis('on')\n",
|
||||
"plt.show()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "027b177b",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Using an ONNX model"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "778d4593",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Here as an example, we use `onnxruntime` in python on CPU to execute the ONNX model. However, any platform that supports an ONNX runtime could be used in principle. Launch the runtime session below:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "9689b1bf",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"ort_session = onnxruntime.InferenceSession(onnx_model_path)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "7708ead6",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"To use the ONNX model, the image must first be pre-processed using the SAM image encoder. This is a heavier weight process best performed on GPU. SamPredictor can be used as normal, then `.get_image_embedding()` will retreive the intermediate features."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "26e067b4",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"sam.to(device='cuda')\n",
|
||||
"predictor = SamPredictor(sam)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "7ad3f0d6",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"predictor.set_image(image)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "8a6f0f07",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"image_embedding = predictor.get_image_embedding().cpu().numpy()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "5e112f33",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"image_embedding.shape"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "6337b654",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The ONNX model has a different input signature than `SamPredictor.predict`. The following inputs must all be supplied. Note the special cases for both point and mask inputs. All inputs are `np.float32`.\n",
|
||||
"* `image_embeddings`: The image embedding from `predictor.get_image_embedding()`. Has a batch index of length 1.\n",
|
||||
"* `point_coords`: Coordinates of sparse input prompts, corresponding to both point inputs and box inputs. Boxes are encoded using two points, one for the top-left corner and one for the bottom-right corner. *Coordinates must already be transformed to long-side 1024.* Has a batch index of length 1.\n",
|
||||
"* `point_labels`: Labels for the sparse input prompts. 0 is a negative input point, 1 is a positive input point, 2 is a top-left box corner, 3 is a bottom-right box corner, and -1 is a padding point. *If there is no box input, a single padding point with label -1 and coordinates (0.0, 0.0) should be concatenated.*\n",
|
||||
"* `mask_input`: A mask input to the model with shape 1x1x256x256. This must be supplied even if there is no mask input. In this case, it can just be zeros.\n",
|
||||
"* `has_mask_input`: An indicator for the mask input. 1 indicates a mask input, 0 indicates no mask input.\n",
|
||||
"* `orig_im_size`: The size of the input image in (H,W) format, before any transformation. \n",
|
||||
"\n",
|
||||
"Additionally, the ONNX model does not threshold the output mask logits. To obtain a binary mask, threshold at `sam.mask_threshold` (equal to 0.0)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "bf5a9f55",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Example point input"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "1c0deef0",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"input_point = np.array([[500, 375]])\n",
|
||||
"input_label = np.array([1])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "7256394c",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Add a batch index, concatenate a padding point, and transform."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "4f69903e",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"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",
|
||||
"\n",
|
||||
"onnx_coord = predictor.transform.apply_coords(onnx_coord, image.shape[:2]).astype(np.float32)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "b188dc53",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Create an empty mask input and an indicator for no mask."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "5cb52bcf",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"onnx_mask_input = np.zeros((1, 1, 256, 256), dtype=np.float32)\n",
|
||||
"onnx_has_mask_input = np.zeros(1, dtype=np.float32)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "a99c2cc5",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Package the inputs to run in the onnx model"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "b1d7ea11",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"ort_inputs = {\n",
|
||||
" \"image_embeddings\": image_embedding,\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",
|
||||
"}"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "4b6409c9",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Predict a mask and threshold it."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "dc4cc082",
|
||||
"metadata": {
|
||||
"scrolled": false
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"masks, _, low_res_logits = ort_session.run(None, ort_inputs)\n",
|
||||
"masks = masks > predictor.model.mask_threshold"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "d778a8fb",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"masks.shape"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "badb1175",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"plt.figure(figsize=(10,10))\n",
|
||||
"plt.imshow(image)\n",
|
||||
"show_mask(masks, plt.gca())\n",
|
||||
"show_points(input_point, input_label, plt.gca())\n",
|
||||
"plt.axis('off')\n",
|
||||
"plt.show() "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "1f1d4d15",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Example mask input"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "b319da82",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"input_point = np.array([[500, 375], [1125, 625]])\n",
|
||||
"input_label = np.array([1, 1])\n",
|
||||
"\n",
|
||||
"# Use the mask output from the previous run. It is already in the correct form for input to the ONNX model.\n",
|
||||
"onnx_mask_input = low_res_logits"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "b1823b37",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Transform the points as in the previous example."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "8885130f",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"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",
|
||||
"\n",
|
||||
"onnx_coord = predictor.transform.apply_coords(onnx_coord, image.shape[:2]).astype(np.float32)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "28e47b69",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The `has_mask_input` indicator is now 1."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "3ab4483a",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"onnx_has_mask_input = np.ones(1, dtype=np.float32)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "d3781955",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Package inputs, then predict and threshold the mask."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "0c1ec096",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"ort_inputs = {\n",
|
||||
" \"image_embeddings\": image_embedding,\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.run(None, ort_inputs)\n",
|
||||
"masks = masks > predictor.model.mask_threshold"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "1e36554b",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"plt.figure(figsize=(10,10))\n",
|
||||
"plt.imshow(image)\n",
|
||||
"show_mask(masks, plt.gca())\n",
|
||||
"show_points(input_point, input_label, plt.gca())\n",
|
||||
"plt.axis('off')\n",
|
||||
"plt.show() "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "2ef211d0",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Example box and point input"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "51e58d2e",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"input_box = np.array([425, 600, 700, 875])\n",
|
||||
"input_point = np.array([[575, 750]])\n",
|
||||
"input_label = np.array([0])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "6e119dcb",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Add a batch index, concatenate a box and point inputs, add the appropriate labels for the box corners, and transform. There is no padding point since the input includes a box input."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "bfbe4911",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"onnx_box_coords = input_box.reshape(2, 2)\n",
|
||||
"onnx_box_labels = np.array([2,3])\n",
|
||||
"\n",
|
||||
"onnx_coord = np.concatenate([input_point, onnx_box_coords], axis=0)[None, :, :]\n",
|
||||
"onnx_label = np.concatenate([input_label, onnx_box_labels], axis=0)[None, :].astype(np.float32)\n",
|
||||
"\n",
|
||||
"onnx_coord = predictor.transform.apply_coords(onnx_coord, image.shape[:2]).astype(np.float32)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "65edabd2",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Package inputs, then predict and threshold the mask."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "2abfba56",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"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_embedding,\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.run(None, ort_inputs)\n",
|
||||
"masks = masks > predictor.model.mask_threshold"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "8301bf33",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"plt.figure(figsize=(10, 10))\n",
|
||||
"plt.imshow(image)\n",
|
||||
"show_mask(masks[0], plt.gca())\n",
|
||||
"show_box(input_box, plt.gca())\n",
|
||||
"show_points(input_point, input_label, plt.gca())\n",
|
||||
"plt.axis('off')\n",
|
||||
"plt.show()"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"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.10.10"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
Reference in new issue
Block a user