Skip to content

Commit 2990c53

Browse files
vertex-mg-botcopybara-github
authored andcommitted
Added notebook sample for batch inference using the remote sensing VMG models
PiperOrigin-RevId: 866061602
1 parent 531d9cf commit 2990c53

1 file changed

Lines changed: 320 additions & 0 deletions

File tree

Lines changed: 320 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,320 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "code",
5+
"execution_count": null,
6+
"metadata": {
7+
"cellView": "form",
8+
"id": "rPH7NVRWhCGT"
9+
},
10+
"outputs": [],
11+
"source": [
12+
"# Copyright 2025 Google LLC\n",
13+
"#\n",
14+
"# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
15+
"# you may not use this file except in compliance with the License.\n",
16+
"# You may obtain a copy of the License at\n",
17+
"#\n",
18+
"# https://www.apache.org/licenses/LICENSE-2.0\n",
19+
"#\n",
20+
"# Unless required by applicable law or agreed to in writing, software\n",
21+
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
22+
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
23+
"# See the License for the specific language governing permissions and\n",
24+
"# limitations under the License."
25+
]
26+
},
27+
{
28+
"cell_type": "markdown",
29+
"metadata": {
30+
"id": "OOTGI5tjjWqL"
31+
},
32+
"source": [
33+
"# RS Imagery Batch Inference on VertexAI\n",
34+
"\n",
35+
"This notebook shows how to run a Batch Prediction Job deployed VLMs (image and text) on Vertex AI.\n",
36+
"\n",
37+
"**Prepare the environment for interacting with Vertex AI:**\n",
38+
"\n",
39+
"Initialize the Vertex AI SDK using the aiplatform.init() function.\n",
40+
"\n",
41+
"Configure the SDK to work with your specific Google Cloud project (PROJECT_ID) and region (REGION) that were defined in the previous configuration cell. This step is necessary before using other SDK functions to manage Vertex AI resources."
42+
]
43+
},
44+
{
45+
"cell_type": "code",
46+
"execution_count": null,
47+
"metadata": {
48+
"cellView": "form",
49+
"id": "cH6TEJQ-F_ne"
50+
},
51+
"outputs": [],
52+
"source": [
53+
"# @title Setup Notebook\n",
54+
"\n",
55+
"# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
56+
"\n",
57+
"# @markdown 2. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
58+
"\n",
59+
"REGION = \"\" # @param {type:\"string\"}\n",
60+
"\n",
61+
"# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
62+
"\n",
63+
"# @markdown | Machine Type | Accelerator Type | Recommended Regions |\n",
64+
"# @markdown | ----------- | ----------- | ----------- |\n",
65+
"# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
66+
"# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
67+
"# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
68+
"# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
69+
"\n",
70+
"import base64\n",
71+
"import io\n",
72+
"import json\n",
73+
"from typing import Any\n",
74+
"\n",
75+
"from google.cloud import aiplatform, storage\n",
76+
"from PIL import Image\n",
77+
"\n",
78+
"# Import common utils\n",
79+
"if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
80+
" ! pip install --upgrade tensorflow\n",
81+
"! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
82+
"\n",
83+
"common_util = importlib.import_module(\n",
84+
" \"vertex-ai-samples.notebooks.community.model_garden.docker_source_codes.notebook_util.common_util\"\n",
85+
")\n",
86+
"\n",
87+
"# Setup GCP & VertexAI\n",
88+
"\n",
89+
"# Get the default cloud project id.\n",
90+
"PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
91+
"\n",
92+
"# Get the default region for launching jobs.\n",
93+
"if not REGION:\n",
94+
" REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
95+
"\n",
96+
"# Enable the Vertex AI API and Compute Engine API, if not already.\n",
97+
"print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
98+
"! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
99+
"\n",
100+
"# Initialize Vertex AI API.\n",
101+
"print(\"Initializing Vertex AI API.\")\n",
102+
"aiplatform.init(project=PROJECT_ID, location=REGION)\n",
103+
"\n",
104+
"# Gets the default SERVICE_ACCOUNT.\n",
105+
"shell_output = ! gcloud projects describe $PROJECT_ID\n",
106+
"project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
107+
"SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
108+
"print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
109+
"\n",
110+
"! gcloud config set project $PROJECT_ID\n",
111+
"import vertexai\n",
112+
"\n",
113+
"vertexai.init(\n",
114+
" project=PROJECT_ID,\n",
115+
" location=REGION,\n",
116+
")\n",
117+
"\n",
118+
"\n",
119+
"def to_png_bytes(img: Image.Image) -> bytes:\n",
120+
" \"\"\"Encodes the `img` as PNG bytes.\"\"\"\n",
121+
" buf = io.BytesIO()\n",
122+
" img.save(buf, format=\"PNG\")\n",
123+
" return buf.getvalue()\n",
124+
"\n",
125+
"\n",
126+
"def to_b64_png(img: Image.Image) -> str:\n",
127+
" \"\"\"Converts the `img` to a b64 encoded PNG bytes.\"\"\"\n",
128+
" return base64.b64encode(to_png_bytes(img)).decode()\n",
129+
"\n",
130+
"\n",
131+
"def write_jsonl_instances(\n",
132+
" bucket: storage.Bucket, path: str, instances: list[dict[str, Any]]\n",
133+
"):\n",
134+
" \"\"\"Writes the list of instances (dicts) as a JSONL serialized file.\n",
135+
"\n",
136+
" Each dict is an inference instance, matching one of the following structures:\n",
137+
"\n",
138+
" {'image': <b64 image>} - Image inference only.\n",
139+
" {'text': <str>} - Text inference only.\n",
140+
" {'image': <b64 image, 'texts': list<str>} - Image & text inference.\n",
141+
" \"\"\"\n",
142+
" with bucket.blob(path).open(\"wt\") as f:\n",
143+
" f.writelines(json.dumps(instance) for instance in instances)"
144+
]
145+
},
146+
{
147+
"cell_type": "code",
148+
"execution_count": null,
149+
"metadata": {
150+
"cellView": "form",
151+
"id": "qVHaf1GmKaT6"
152+
},
153+
"outputs": [],
154+
"source": [
155+
"# @title Initialize data bucket\n",
156+
"# @markdown ### Enter a GCS bucket name and a path within the bucket:\n",
157+
"\n",
158+
"BUCKET_NAME = \"\" # @param { type : 'string' }\n",
159+
"OUTPUT_PATH = \"batch_inference/inputs\" # @param { type : 'string' }\n",
160+
"\n",
161+
"storage_client = storage.Client()\n",
162+
"bucket = storage_client.bucket(BUCKET_NAME)\n",
163+
"\n",
164+
"# Download sample image\n",
165+
"!wget -O harbor.jpg https://mrsg.aegean.gr/images/uploads/it2zi0eidej4ql33llj.jpg\n",
166+
"harbor_img = Image.open(\"harbor.jpg\")"
167+
]
168+
},
169+
{
170+
"cell_type": "code",
171+
"execution_count": null,
172+
"metadata": {
173+
"cellView": "form",
174+
"id": "0nOBY2kjGHob"
175+
},
176+
"outputs": [],
177+
"source": [
178+
"# @title Prepare data: Image list file\n",
179+
"\n",
180+
"input_uris = []\n",
181+
"\n",
182+
"# The sample image (harbor) is replicated 10 times as an example.\n",
183+
"for i in range(10):\n",
184+
" img_path = f\"{OUTPUT_PATH}/images/img{i}.png\"\n",
185+
" input_uris.append(f\"gs://{BUCKET_NAME}/{img_path}\")\n",
186+
" bucket.blob(img_path).upload_from_string(\n",
187+
" to_png_bytes(harbor_img), content_type=\"image/png\"\n",
188+
" )\n",
189+
"\n",
190+
"with bucket.blob(f\"{OUTPUT_PATH}/input_uris.txt\").open(\"wt\") as f:\n",
191+
" f.writelines([f\"{i}\\n\" for i in input_uris])"
192+
]
193+
},
194+
{
195+
"cell_type": "code",
196+
"execution_count": null,
197+
"metadata": {
198+
"cellView": "form",
199+
"id": "ap9IjkphGHbZ"
200+
},
201+
"outputs": [],
202+
"source": [
203+
"# @title Prepare JSONL input\n",
204+
"\n",
205+
"# This cell generates sample inputs (JSONL files) in 3 formats:\n",
206+
"# Image, text and image_text, the input files are sharded to optionally reduce\n",
207+
"# the size of each file. The file pattern (with wildcards) is used as input for\n",
208+
"# the batch pipeline, e.g. \"gs://bucket_path/image*.jsonl\"\n",
209+
"\n",
210+
"# Write 3 shards of text input instances (10 instances each).\n",
211+
"instances = [{\"text\": \"test string\"}] * 10\n",
212+
"for i in range(3):\n",
213+
" write_jsonl_instances(bucket, f\"{OUTPUT_PATH}/text{i}.jsonl\", instances)\n",
214+
"\n",
215+
"# Write 10 image input instances into 3 shards.\n",
216+
"instances = [{\"image\": to_b64_png(harbor_img)}] * 10\n",
217+
"for p in range(3):\n",
218+
" write_jsonl_instances(bucket, f\"{OUTPUT_PATH}/image{p}.jsonl\", instances)\n",
219+
"\n",
220+
"# Write 10 image & texts input instances into a single file.\n",
221+
"instances = [\n",
222+
" {\"image\": to_b64_png(harbor_img), \"texts\": [\"text1\", \"text2\"]},\n",
223+
"] * 10\n",
224+
"write_jsonl_instances(bucket, f\"{OUTPUT_PATH}/combined.jsonl\", instances)"
225+
]
226+
},
227+
{
228+
"cell_type": "code",
229+
"execution_count": null,
230+
"metadata": {
231+
"cellView": "form",
232+
"id": "a2j33P9hGN9H"
233+
},
234+
"outputs": [],
235+
"source": [
236+
"# @title Run batch inference\n",
237+
"\n",
238+
"JOB_DISPLAY_NAME = \"batch-inference-vlm-test\" # @param { type: 'string' }\n",
239+
"# @markdown Enter the project number and model id, the project number is not the\n",
240+
"# @markdown same as project id (it can be found in the project settings).\n",
241+
"PROJECT_NUMBER = \"<project_number_here>\" # @param { type: 'string' }\n",
242+
"MODEL_ID = \"<model_id_here>\" # @param { type: 'string' }\n",
243+
"MODEL_RESOURCE_NAME = f\"projects/{PROJECT_NUMBER}/locations/{REGION}/models/{MODEL_ID}\"\n",
244+
"\n",
245+
"# @markdown Choose the input (instances) format, either a file-list of images or\n",
246+
"# @markdown a JSONL file pattern of JSON formatted inputs.\n",
247+
"INPUT_SOURCE_FORMAT = \"file-list\" # @param[\"file-list\", \"jsonl\"]\n",
248+
"# @markdown Configure batch input source This can use string wildcards such as\n",
249+
"# @markdown '*' and '?' to support sharded inputs.\n",
250+
"INPUT_SOURCE_PATTERN = \"gs://<bucket>/batch_inference/inputs/inputlist.txt\" # @param { type: 'string' }\n",
251+
"# @markdown Configure the output folder path, predictions will be written here.\n",
252+
"GCS_OUTPUT_PATH = \"gs://<bucket>/batch_inference/outputs\" # @param { type: 'string' }\n",
253+
"# @markdown Configure the batch runtime setup\n",
254+
"USE_GPU = True # @param { type: 'boolean' }\n",
255+
"BATCH_SIZE = 16 # @param { type: 'number' }\n",
256+
"REPLICA_COUNT = 1 # @param { type: 'number' }\n",
257+
"MAX_REPLICA_COUNT = 4 # @param { type: 'number' }\n",
258+
"\n",
259+
"machine_type = \"g2-standard-8\"\n",
260+
"if USE_GPU:\n",
261+
" accelerator_type = \"NVIDIA_L4\"\n",
262+
" accelerator_count = 1\n",
263+
"else:\n",
264+
" accelerator_type = None\n",
265+
" accelerator_count = None\n",
266+
"\n",
267+
"model = aiplatform.Model(MODEL_RESOURCE_NAME)\n",
268+
"\n",
269+
"job = model.batch_predict(\n",
270+
" job_display_name=JOB_DISPLAY_NAME,\n",
271+
" gcs_source=INPUT_SOURCE_PATTERN,\n",
272+
" gcs_destination_prefix=GCS_OUTPUT_PATH,\n",
273+
" instances_format=INPUT_SOURCE_FORMAT,\n",
274+
" machine_type=machine_type,\n",
275+
" accelerator_count=accelerator_count,\n",
276+
" accelerator_type=accelerator_type,\n",
277+
" starting_replica_count=REPLICA_COUNT,\n",
278+
" max_replica_count=MAX_REPLICA_COUNT,\n",
279+
" labels={\n",
280+
" \"task\": \"batch-inference\",\n",
281+
" \"vertex-ai-pipelines-run-billing-id\": JOB_DISPLAY_NAME,\n",
282+
" },\n",
283+
" batch_size=BATCH_SIZE,\n",
284+
" sync=False,\n",
285+
")\n",
286+
"\n",
287+
"print(f\"Batch prediction job started: {job}\")"
288+
]
289+
},
290+
{
291+
"cell_type": "code",
292+
"execution_count": null,
293+
"metadata": {
294+
"cellView": "form",
295+
"id": "KF9Lwsb5GOIS"
296+
},
297+
"outputs": [],
298+
"source": [
299+
"# Monitor the job status\n",
300+
"\n",
301+
"print(\n",
302+
" f\"Running batch prediction {job.display_name}, resource:\"\n",
303+
" f\" {job.resource_name}. State: {job.state}\"\n",
304+
")"
305+
]
306+
}
307+
],
308+
"metadata": {
309+
"colab": {
310+
"name": "model_garden_remote_sensing_batch_prediction.ipynb",
311+
"toc_visible": true
312+
},
313+
"kernelspec": {
314+
"display_name": "Python 3",
315+
"name": "python3"
316+
}
317+
},
318+
"nbformat": 4,
319+
"nbformat_minor": 0
320+
}

0 commit comments

Comments
 (0)