Skip to content

Commit 07b1b8b

Browse files
refactor: (GenAI) Reorganized Image Generation Samples (Group B) (#12595)
* New Image Generation folder with tests
1 parent ea03651 commit 07b1b8b

File tree

54 files changed

+1333
-0
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

54 files changed

+1333
-0
lines changed
Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
# Copyright 2024 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Google Cloud Vertex AI sample for editing an image using a mask file.
16+
Inpainting can insert the object designated by the prompt into the masked
17+
area.
18+
"""
19+
import os
20+
21+
from vertexai.preview import vision_models
22+
23+
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
24+
25+
26+
def edit_image_inpainting_insert_mask(
27+
input_file: str,
28+
mask_file: str,
29+
output_file: str,
30+
prompt: str,
31+
) -> vision_models.ImageGenerationResponse:
32+
# [START generativeaionvertexai_imagen_edit_image_inpainting_insert_mask]
33+
34+
import vertexai
35+
from vertexai.preview.vision_models import Image, ImageGenerationModel
36+
37+
# TODO(developer): Update and un-comment below lines
38+
# PROJECT_ID = "your-project-id"
39+
# input_file = "input-image.png"
40+
# mask_file = "mask-image.png"
41+
# output_file = "output-image.png"
42+
# prompt = "red hat" # The text prompt describing what you want to see inserted.
43+
44+
vertexai.init(project=PROJECT_ID, location="us-central1")
45+
46+
model = ImageGenerationModel.from_pretrained("imagegeneration@006")
47+
base_img = Image.load_from_file(location=input_file)
48+
mask_img = Image.load_from_file(location=mask_file)
49+
50+
images = model.edit_image(
51+
base_image=base_img,
52+
mask=mask_img,
53+
prompt=prompt,
54+
edit_mode="inpainting-insert",
55+
)
56+
57+
images[0].save(location=output_file, include_generation_parameters=False)
58+
59+
# Optional. View the edited image in a notebook.
60+
# images[0].show()
61+
62+
print(f"Created output image using {len(images[0]._image_bytes)} bytes")
63+
# Example response:
64+
# Created output image using 1400814 bytes
65+
66+
# [END generativeaionvertexai_imagen_edit_image_inpainting_insert_mask]
67+
68+
return images
69+
70+
71+
if __name__ == "__main__":
72+
edit_image_inpainting_insert_mask(
73+
input_file="test_resources/woman.png",
74+
mask_file="test_resources/woman_inpainting_insert_mask.png",
75+
output_file="test_resources/woman_with_hat.png",
76+
prompt="red hat",
77+
)
Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
# Copyright 2024 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Google Cloud Vertex AI sample for editing an image using a mask mode. The
16+
mask mode is used to automatically select the background, foreground (i.e.,
17+
the primary subject of the image), or an object based on segmentation class.
18+
Inpainting can insert the object or background designated by the prompt.
19+
"""
20+
import os
21+
22+
from vertexai.preview import vision_models
23+
24+
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
25+
26+
27+
def edit_image_inpainting_insert_mask_mode(
28+
input_file: str,
29+
mask_mode: str,
30+
output_file: str,
31+
prompt: str,
32+
) -> vision_models.ImageGenerationResponse:
33+
# [START generativeaionvertexai_imagen_edit_image_inpainting_insert_mask_mode]
34+
35+
import vertexai
36+
from vertexai.preview.vision_models import Image, ImageGenerationModel
37+
38+
# TODO(developer): Update and un-comment below lines
39+
# PROJECT_ID = "your-project-id"
40+
# input_file = "input-image.png"
41+
# mask_mode = "background" # 'background', 'foreground', or 'semantic'
42+
# output_file = "output-image.png"
43+
# prompt = "beach" # The text prompt describing what you want to see inserted.
44+
45+
vertexai.init(project=PROJECT_ID, location="us-central1")
46+
47+
model = ImageGenerationModel.from_pretrained("imagegeneration@006")
48+
base_img = Image.load_from_file(location=input_file)
49+
50+
images = model.edit_image(
51+
base_image=base_img,
52+
mask_mode=mask_mode,
53+
prompt=prompt,
54+
edit_mode="inpainting-insert",
55+
)
56+
57+
images[0].save(location=output_file, include_generation_parameters=False)
58+
59+
# Optional. View the edited image in a notebook.
60+
# images[0].show()
61+
62+
print(f"Created output image using {len(images[0]._image_bytes)} bytes")
63+
# Example response:
64+
# Created output image using 1234567 bytes
65+
66+
# [END generativeaionvertexai_imagen_edit_image_inpainting_insert_mask_mode]
67+
68+
return images
69+
70+
71+
if __name__ == "__main__":
72+
edit_image_inpainting_insert_mask_mode(
73+
input_file="test_resources/woman.png",
74+
mask_mode="background",
75+
output_file="test_resources/woman_at_beach.png",
76+
prompt="beach",
77+
)
Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
# Copyright 2024 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import os
16+
17+
import backoff
18+
19+
import edit_image_inpainting_insert_mask_mode
20+
21+
from google.api_core.exceptions import ResourceExhausted
22+
23+
24+
_RESOURCES = os.path.join(os.path.dirname(__file__), "test_resources")
25+
_INPUT_FILE = os.path.join(_RESOURCES, "woman.png")
26+
_MASK_MODE = "background"
27+
_OUTPUT_FILE = os.path.join(_RESOURCES, "woman_at_beach.png")
28+
_PROMPT = "beach"
29+
30+
31+
@backoff.on_exception(backoff.expo, ResourceExhausted, max_time=60)
32+
def test_edit_image_inpainting_insert_mask_mode() -> None:
33+
response = (
34+
edit_image_inpainting_insert_mask_mode.edit_image_inpainting_insert_mask_mode(
35+
_INPUT_FILE,
36+
_MASK_MODE,
37+
_OUTPUT_FILE,
38+
_PROMPT,
39+
)
40+
)
41+
42+
assert len(response[0]._image_bytes) > 1000
Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
# Copyright 2024 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
import os
15+
16+
import backoff
17+
18+
import edit_image_inpainting_insert_mask
19+
20+
from google.api_core.exceptions import ResourceExhausted
21+
22+
23+
_RESOURCES = os.path.join(os.path.dirname(__file__), "test_resources")
24+
_INPUT_FILE = os.path.join(_RESOURCES, "woman.png")
25+
_MASK_FILE = os.path.join(_RESOURCES, "woman_inpainting_insert_mask.png")
26+
_OUTPUT_FILE = os.path.join(_RESOURCES, "woman_with_hat.png")
27+
_PROMPT = "hat"
28+
29+
30+
@backoff.on_exception(backoff.expo, ResourceExhausted, max_time=60)
31+
def test_edit_image_inpainting_insert_mask() -> None:
32+
response = edit_image_inpainting_insert_mask.edit_image_inpainting_insert_mask(
33+
_INPUT_FILE,
34+
_MASK_FILE,
35+
_OUTPUT_FILE,
36+
_PROMPT,
37+
)
38+
39+
assert len(response[0]._image_bytes) > 1000
Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
# Copyright 2024 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Google Cloud Vertex AI sample for editing an image using a mask file.
16+
Inpainting can remove an object from the masked area.
17+
"""
18+
import os
19+
20+
from vertexai.preview import vision_models
21+
22+
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
23+
24+
25+
def edit_image_inpainting_remove_mask(
26+
input_file: str,
27+
mask_file: str,
28+
output_file: str,
29+
prompt: str,
30+
) -> vision_models.ImageGenerationResponse:
31+
# [START generativeaionvertexai_imagen_edit_image_inpainting_remove_mask]
32+
33+
import vertexai
34+
from vertexai.preview.vision_models import Image, ImageGenerationModel
35+
36+
# TODO(developer): Update and un-comment below lines
37+
# PROJECT_ID = "your-project-id"
38+
# input_file = "input-image.png"
39+
# mask_file = "mask-image.png"
40+
# output_file = "outpur-image.png"
41+
# prompt = "" # The text prompt describing the entire image.
42+
43+
vertexai.init(project=PROJECT_ID, location="us-central1")
44+
45+
model = ImageGenerationModel.from_pretrained("imagegeneration@006")
46+
base_img = Image.load_from_file(location=input_file)
47+
mask_img = Image.load_from_file(location=mask_file)
48+
49+
images = model.edit_image(
50+
base_image=base_img,
51+
mask=mask_img,
52+
prompt=prompt,
53+
edit_mode="inpainting-remove",
54+
# Optional parameters
55+
# negative_prompt="", # Describes the object being removed (i.e., "person")
56+
)
57+
58+
images[0].save(location=output_file, include_generation_parameters=False)
59+
60+
# Optional. View the edited image in a notebook.
61+
# images[0].show()
62+
63+
print(f"Created output image using {len(images[0]._image_bytes)} bytes")
64+
# Example response:
65+
# Created output image using 12345678 bytes
66+
67+
# [END generativeaionvertexai_imagen_edit_image_inpainting_remove_mask]
68+
69+
return images
70+
71+
72+
if __name__ == "__main__":
73+
edit_image_inpainting_remove_mask(
74+
input_file="test_resources/volleyball_game.png",
75+
mask_file="test_resources/volleyball_game_inpainting_remove_mask.png",
76+
output_file="test_resources/volleyball_game_single_blue_player.png",
77+
prompt="volleyball game",
78+
)
Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
# Copyright 2024 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Google Cloud Vertex AI sample for editing an image using a mask mode. The
16+
mask mode is used to automatically select the background, foreground (i.e.,
17+
the primary subject of the image), or an object based on segmentation class.
18+
Inpainting can remove the object or background designated by the prompt.
19+
"""
20+
import os
21+
22+
from vertexai.preview import vision_models
23+
24+
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
25+
26+
27+
def edit_image_inpainting_remove_mask_mode(
28+
input_file: str,
29+
mask_mode: str,
30+
output_file: str,
31+
prompt: str,
32+
) -> vision_models.ImageGenerationResponse:
33+
# [START generativeaionvertexai_imagen_edit_image_inpainting_remove_mask_mode]
34+
35+
import vertexai
36+
from vertexai.preview.vision_models import Image, ImageGenerationModel
37+
38+
# TODO(developer): Update and un-comment below lines
39+
# PROJECT_ID = "your-project-id"
40+
# input_file = "input-image.png"
41+
# mask_mode = "foreground" # 'background', 'foreground', or 'semantic'
42+
# output_file = "output-image.png"
43+
# prompt = "sports car" # The text prompt describing what you want to see in the edited image.
44+
45+
vertexai.init(project=PROJECT_ID, location="us-central1")
46+
47+
model = ImageGenerationModel.from_pretrained("imagegeneration@006")
48+
base_img = Image.load_from_file(location=input_file)
49+
50+
images = model.edit_image(
51+
base_image=base_img,
52+
mask_mode=mask_mode,
53+
prompt=prompt,
54+
edit_mode="inpainting-remove",
55+
)
56+
57+
images[0].save(location=output_file, include_generation_parameters=False)
58+
59+
# Optional. View the edited image in a notebook.
60+
# images[0].show()
61+
62+
print(f"Created output image using {len(images[0]._image_bytes)} bytes")
63+
# Example response:
64+
# Created output image using 1279948 bytes
65+
66+
# [END generativeaionvertexai_imagen_edit_image_inpainting_remove_mask_mode]
67+
68+
return images
69+
70+
71+
if __name__ == "__main__":
72+
edit_image_inpainting_remove_mask_mode(
73+
input_file="test_resources/woman.png",
74+
mask_mode="foreground",
75+
output_file="test_resources/sports_car.png",
76+
prompt="sports car",
77+
)

0 commit comments

Comments
 (0)