Update SAM notebook to use new container that is FedRamp compliant.

PiperOrigin-RevId: 782652376
This commit is contained in:
Dustin Luong
2025-07-13 14:06:26 -07:00
committed by Copybara-Service
parent ca19b8f8e7
commit 0ff91f926d
@@ -116,14 +116,12 @@
"! pip3 install --upgrade --quiet 'google-cloud-aiplatform==1.103.0'\n",
"! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
"\n",
"! pip install --quiet pycocotools\n",
"\n",
"import gc\n",
"import importlib\n",
"import os\n",
"\n",
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"import pycocotools.mask as mask_util\n",
"from google.cloud import aiplatform\n",
"\n",
"if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
@@ -174,7 +172,7 @@
"\n",
"# The pre-built serving docker image.\n",
"# The model artifacts are embedded within the container, except for model weights which will be downloaded during deployment.\n",
"SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/sam-serve:public-image-20240121\"\n",
"SERVE_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/pytorch-inference.cu125.0-4.ubuntu2204.py310\"\n",
"\n",
"# @markdown Set the accelerator type.\n",
"accelerator_type = \"NVIDIA_L4\" # @param[\"NVIDIA_TESLA_V100\", \"NVIDIA_L4\"]\n",
@@ -211,7 +209,7 @@
" display_name=task,\n",
" serving_container_image_uri=SERVE_DOCKER_URI,\n",
" serving_container_ports=[7080],\n",
" serving_container_predict_route=\"/predictions/sam_serving\",\n",
" serving_container_predict_route=\"/predict\",\n",
" serving_container_health_route=\"/ping\",\n",
" serving_container_environment_variables=serving_env,\n",
" model_garden_source_model_name=\"publishers/meta/models/segment-anything\",\n",
@@ -309,11 +307,8 @@
"input_image2 = \"http://images.cocodataset.org/val2017/000000000285.jpg\" # @param {type:\"string\"}\n",
"\n",
"\n",
"def decode_rle_masks(pred_masks_rle):\n",
" return np.stack([mask_util.decode(rle) for rle in pred_masks_rle])\n",
"\n",
"\n",
"def show_mask(mask, ax, random_color=False):\n",
" mask = np.array(mask)\n",
" if random_color:\n",
" color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)\n",
" else:\n",
@@ -321,27 +316,20 @@
" h, w = mask.shape[-2:]\n",
" mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)\n",
" ax.imshow(mask_image)\n",
" del mask\n",
" gc.collect()\n",
"\n",
"\n",
"def show_predictions(preds):\n",
" fig = plt.figure(figsize=(10, 7))\n",
"\n",
" fig.add_subplot(1, 2, 1)\n",
" plt.imshow(np.array(image1))\n",
"def show_masks_on_image(raw_image, masks):\n",
" plt.imshow(np.array(raw_image))\n",
" ax = plt.gca()\n",
" masks = decode_rle_masks(preds[0][\"masks_rle\"])\n",
" for mask in masks:\n",
" show_mask(mask, ax=ax, random_color=True)\n",
" plt.axis(\"off\")\n",
"\n",
" fig.add_subplot(1, 2, 2)\n",
" plt.imshow(np.array(image2))\n",
" ax = plt.gca()\n",
" masks = decode_rle_masks(preds[1][\"masks_rle\"])\n",
" ax.set_autoscale_on(False)\n",
" for mask in masks:\n",
" show_mask(mask, ax=ax, random_color=True)\n",
" plt.axis(\"off\")\n",
" plt.show()\n",
" del mask\n",
" gc.collect()\n",
"\n",
"\n",
"image1 = common_util.download_image(input_image1)\n",
@@ -357,7 +345,9 @@
"preds = endpoint.predict(\n",
" instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
").predictions\n",
"show_predictions(preds)"
"\n",
"show_masks_on_image(image1, preds[0][\"masks\"])\n",
"show_masks_on_image(image2, preds[1][\"masks\"])"
]
},
{