Skip to content

Commit f3bd3bf

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
feat: LVM - Added support for Videos from GCS uri for multimodal embeddings
PiperOrigin-RevId: 606757856
1 parent e35ab64 commit f3bd3bf

5 files changed

Lines changed: 411 additions & 9 deletions

File tree

‎tests/system/aiplatform/test_vision_models.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,12 @@ def _load_image_from_gcs(
4444
return vision_models.Image.load_from_file(gcs_uri)
4545

4646

47+
def _load_video_from_gcs(
48+
gcs_uri: str = "gs://cloud-samples-data/vertex-ai-vision/highway_vehicles.mp4",
49+
) -> vision_models.Video:
50+
return vision_models.Video.load_from_file(gcs_uri)
51+
52+
4753
class VisionModelTestSuite(e2e_base.TestEndToEnd):
4854
"""System tests for vision models."""
4955

@@ -98,13 +104,18 @@ def test_multi_modal_embedding_model_with_gcs_uri(self):
98104
"multimodalembedding@001"
99105
)
100106
image = _load_image_from_gcs()
107+
video = _load_video_from_gcs()
108+
video_segment_config = vision_models.VideoSegmentConfig()
101109
embeddings = model.get_embeddings(
102110
image=image,
111+
video=video,
103112
# Optional:
104113
contextual_text="this is a car",
114+
video_segment_config=video_segment_config,
105115
)
106116
# The service is expected to return the embeddings of size 1408
107117
assert len(embeddings.image_embedding) == 1408
118+
assert len(embeddings.video_embeddings[0].embedding) == 1408
108119
assert len(embeddings.text_embedding) == 1408
109120

110121
def test_image_generation_model_generate_images(self):

‎tests/unit/aiplatform/test_vision_models.py‎

Lines changed: 212 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,12 @@ def generate_image_from_gcs_uri(
158158
return ga_vision_models.Image.load_from_file(gcs_uri)
159159

160160

161+
def generate_video_from_gcs_uri(
162+
gcs_uri: str = "gs://cloud-samples-data/vertex-ai-vision/highway_vehicles.mp4",
163+
) -> ga_vision_models.Video:
164+
return ga_vision_models.Video.load_from_file(gcs_uri)
165+
166+
161167
@pytest.mark.usefixtures("google_auth_mock")
162168
class TestImageGenerationModels:
163169
"""Unit tests for the image generation models."""
@@ -888,6 +894,212 @@ def test_image_embedding_model_with_gcs_uri(self):
888894
assert embedding_response.image_embedding == test_embeddings
889895
assert embedding_response.text_embedding == test_embeddings
890896

897+
def test_video_embedding_model_with_only_video(self):
898+
aiplatform.init(
899+
project=_TEST_PROJECT,
900+
location=_TEST_LOCATION,
901+
)
902+
with mock.patch.object(
903+
target=model_garden_service_client.ModelGardenServiceClient,
904+
attribute="get_publisher_model",
905+
return_value=gca_publisher_model.PublisherModel(
906+
_IMAGE_EMBEDDING_PUBLISHER_MODEL_DICT
907+
),
908+
) as mock_get_publisher_model:
909+
model = preview_vision_models.MultiModalEmbeddingModel.from_pretrained(
910+
"multimodalembedding@001"
911+
)
912+
913+
mock_get_publisher_model.assert_called_once_with(
914+
name="publishers/google/models/multimodalembedding@001",
915+
retry=base._DEFAULT_RETRY,
916+
)
917+
918+
test_video_embeddings = [
919+
ga_vision_models.VideoEmbedding(
920+
start_offset_sec=0,
921+
end_offset_sec=7,
922+
embedding=[0, 7],
923+
)
924+
]
925+
926+
gca_predict_response = gca_prediction_service.PredictResponse()
927+
gca_predict_response.predictions.append(
928+
{
929+
"videoEmbeddings": [
930+
{
931+
"startOffsetSec": test_video_embeddings[0].start_offset_sec,
932+
"endOffsetSec": test_video_embeddings[0].end_offset_sec,
933+
"embedding": test_video_embeddings[0].embedding,
934+
}
935+
]
936+
}
937+
)
938+
939+
video = generate_video_from_gcs_uri()
940+
941+
with mock.patch.object(
942+
target=prediction_service_client.PredictionServiceClient,
943+
attribute="predict",
944+
return_value=gca_predict_response,
945+
):
946+
embedding_response = model.get_embeddings(video=video)
947+
948+
assert (
949+
embedding_response.video_embeddings[0].embedding
950+
== test_video_embeddings[0].embedding
951+
)
952+
assert (
953+
embedding_response.video_embeddings[0].start_offset_sec
954+
== test_video_embeddings[0].start_offset_sec
955+
)
956+
assert (
957+
embedding_response.video_embeddings[0].end_offset_sec
958+
== test_video_embeddings[0].end_offset_sec
959+
)
960+
assert not embedding_response.text_embedding
961+
assert not embedding_response.image_embedding
962+
963+
def test_video_embedding_model_with_video_and_text(self):
964+
aiplatform.init(
965+
project=_TEST_PROJECT,
966+
location=_TEST_LOCATION,
967+
)
968+
with mock.patch.object(
969+
target=model_garden_service_client.ModelGardenServiceClient,
970+
attribute="get_publisher_model",
971+
return_value=gca_publisher_model.PublisherModel(
972+
_IMAGE_EMBEDDING_PUBLISHER_MODEL_DICT
973+
),
974+
) as mock_get_publisher_model:
975+
model = preview_vision_models.MultiModalEmbeddingModel.from_pretrained(
976+
"multimodalembedding@001"
977+
)
978+
979+
mock_get_publisher_model.assert_called_once_with(
980+
name="publishers/google/models/multimodalembedding@001",
981+
retry=base._DEFAULT_RETRY,
982+
)
983+
984+
test_text_embedding = [0, 0]
985+
test_video_embeddings = [
986+
ga_vision_models.VideoEmbedding(
987+
start_offset_sec=0,
988+
end_offset_sec=7,
989+
embedding=test_text_embedding,
990+
)
991+
]
992+
gca_predict_response = gca_prediction_service.PredictResponse()
993+
gca_predict_response.predictions.append(
994+
{
995+
"textEmbedding": test_text_embedding,
996+
"videoEmbeddings": [
997+
{
998+
"startOffsetSec": test_video_embeddings[0].start_offset_sec,
999+
"endOffsetSec": test_video_embeddings[0].end_offset_sec,
1000+
"embedding": test_video_embeddings[0].embedding,
1001+
}
1002+
],
1003+
}
1004+
)
1005+
1006+
video = generate_video_from_gcs_uri()
1007+
1008+
with mock.patch.object(
1009+
target=prediction_service_client.PredictionServiceClient,
1010+
attribute="predict",
1011+
return_value=gca_predict_response,
1012+
):
1013+
embedding_response = model.get_embeddings(
1014+
video=video, contextual_text="hello world"
1015+
)
1016+
1017+
assert (
1018+
embedding_response.video_embeddings[0].embedding
1019+
== test_video_embeddings[0].embedding
1020+
)
1021+
assert (
1022+
embedding_response.video_embeddings[0].start_offset_sec
1023+
== test_video_embeddings[0].start_offset_sec
1024+
)
1025+
assert (
1026+
embedding_response.video_embeddings[0].end_offset_sec
1027+
== test_video_embeddings[0].end_offset_sec
1028+
)
1029+
assert embedding_response.text_embedding == test_text_embedding
1030+
assert not embedding_response.image_embedding
1031+
1032+
def test_multimodal_embedding_model_with_image_video_and_text(self):
1033+
aiplatform.init(
1034+
project=_TEST_PROJECT,
1035+
location=_TEST_LOCATION,
1036+
)
1037+
with mock.patch.object(
1038+
target=model_garden_service_client.ModelGardenServiceClient,
1039+
attribute="get_publisher_model",
1040+
return_value=gca_publisher_model.PublisherModel(
1041+
_IMAGE_EMBEDDING_PUBLISHER_MODEL_DICT
1042+
),
1043+
) as mock_get_publisher_model:
1044+
model = preview_vision_models.MultiModalEmbeddingModel.from_pretrained(
1045+
"multimodalembedding@001"
1046+
)
1047+
1048+
mock_get_publisher_model.assert_called_once_with(
1049+
name="publishers/google/models/multimodalembedding@001",
1050+
retry=base._DEFAULT_RETRY,
1051+
)
1052+
1053+
test_embedding = [0, 0]
1054+
test_video_embeddings = [
1055+
ga_vision_models.VideoEmbedding(
1056+
start_offset_sec=0,
1057+
end_offset_sec=7,
1058+
embedding=test_embedding,
1059+
)
1060+
]
1061+
gca_predict_response = gca_prediction_service.PredictResponse()
1062+
gca_predict_response.predictions.append(
1063+
{
1064+
"textEmbedding": test_embedding,
1065+
"imageEmbedding": test_embedding,
1066+
"videoEmbeddings": [
1067+
{
1068+
"startOffsetSec": test_video_embeddings[0].start_offset_sec,
1069+
"endOffsetSec": test_video_embeddings[0].end_offset_sec,
1070+
"embedding": test_video_embeddings[0].embedding,
1071+
}
1072+
],
1073+
}
1074+
)
1075+
1076+
image = generate_image_from_file()
1077+
video = generate_video_from_gcs_uri()
1078+
1079+
with mock.patch.object(
1080+
target=prediction_service_client.PredictionServiceClient,
1081+
attribute="predict",
1082+
return_value=gca_predict_response,
1083+
):
1084+
embedding_response = model.get_embeddings(
1085+
video=video, image=image, contextual_text="hello world"
1086+
)
1087+
1088+
assert (
1089+
embedding_response.video_embeddings[0].embedding
1090+
== test_video_embeddings[0].embedding
1091+
)
1092+
assert (
1093+
embedding_response.video_embeddings[0].start_offset_sec
1094+
== test_video_embeddings[0].start_offset_sec
1095+
)
1096+
assert (
1097+
embedding_response.video_embeddings[0].end_offset_sec
1098+
== test_video_embeddings[0].end_offset_sec
1099+
)
1100+
assert embedding_response.text_embedding == test_embedding
1101+
assert embedding_response.image_embedding == test_embedding
1102+
8911103

8921104
@pytest.mark.usefixtures("google_auth_mock")
8931105
class ImageTextModelTests:

‎vertexai/preview/vision_models.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,9 @@
2424
GeneratedImage,
2525
MultiModalEmbeddingModel,
2626
MultiModalEmbeddingResponse,
27+
Video,
28+
VideoEmbedding,
29+
VideoSegmentConfig,
2730
)
2831

2932
__all__ = [
@@ -36,4 +39,7 @@
3639
"GeneratedImage",
3740
"MultiModalEmbeddingModel",
3841
"MultiModalEmbeddingResponse",
42+
"Video",
43+
"VideoEmbedding",
44+
"VideoSegmentConfig",
3945
]

‎vertexai/vision_models/__init__.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,9 @@
2121
ImageTextModel,
2222
MultiModalEmbeddingModel,
2323
MultiModalEmbeddingResponse,
24+
Video,
25+
VideoEmbedding,
26+
VideoSegmentConfig,
2427
)
2528

2629
__all__ = [
@@ -30,4 +33,7 @@
3033
"ImageTextModel",
3134
"MultiModalEmbeddingModel",
3235
"MultiModalEmbeddingResponse",
36+
"Video",
37+
"VideoEmbedding",
38+
"VideoSegmentConfig",
3339
]

0 commit comments

Comments
 (0)