Skip to content

Commit f42890d

Browse files
katiemnkatiemn
authored
feat: imagen editing with masks (#13449)
1 parent 61db94c commit f42890d

6 files changed

+204
-0
lines changed

‎genai/image_generation/imggen_inpainting_insert_mask_with_txt_img.py

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
# Copyright 2025 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+
from google.genai.types import Image
16+
17+
18+
def edit_inpainting_insert_mask(output_file: str) -> Image:
19+
# [START googlegenaisdk_imggen_inpainting_insert_mask_with_txt_img]
20+
from google import genai
21+
from google.genai.types import RawReferenceImage, MaskReferenceImage, MaskReferenceConfig, EditImageConfig
22+
23+
client = genai.Client()
24+
25+
# TODO(developer): Update and un-comment below line
26+
# output_file = "output-image.png"
27+
28+
raw_ref = RawReferenceImage(
29+
reference_image=Image.from_file(location='test_resources/fruit.png'), reference_id=0)
30+
mask_ref = MaskReferenceImage(
31+
reference_id=1,
32+
reference_image=Image.from_file(location='test_resources/fruit_mask.png'),
33+
config=MaskReferenceConfig(
34+
mask_mode="MASK_MODE_USER_PROVIDED",
35+
mask_dilation=0.01,
36+
),
37+
)
38+
39+
image = client.models.edit_image(
40+
model="imagen-3.0-capability-001",
41+
prompt="A plate of cookies",
42+
reference_images=[raw_ref, mask_ref],
43+
config=EditImageConfig(
44+
edit_mode="EDIT_MODE_INPAINT_INSERTION",
45+
),
46+
)
47+
48+
image.generated_images[0].image.save(output_file)
49+
50+
print(f"Created output image using {len(image.generated_images[0].image.image_bytes)} bytes")
51+
# Example response:
52+
# Created output image using 1234567 bytes
53+
54+
# [END googlegenaisdk_imggen_inpainting_insert_mask_with_txt_img]
55+
return image.generated_images[0].image
56+
57+
58+
if __name__ == "__main__":
59+
edit_inpainting_insert_mask(output_file="test_resources/fruit_edit.png")

‎genai/image_generation/imggen_inpainting_removal_mask_with_txt_img.py

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
# Copyright 2025 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+
from google.genai.types import Image
16+
17+
18+
def edit_inpainting_removal_mask(output_file: str) -> Image:
19+
# [START googlegenaisdk_imggen_inpainting_removal_mask_with_txt_img]
20+
from google import genai
21+
from google.genai.types import RawReferenceImage, MaskReferenceImage, MaskReferenceConfig, EditImageConfig
22+
23+
client = genai.Client()
24+
25+
# TODO(developer): Update and un-comment below line
26+
# output_file = "output-image.png"
27+
28+
raw_ref = RawReferenceImage(
29+
reference_image=Image.from_file(location='test_resources/fruit.png'), reference_id=0)
30+
mask_ref = MaskReferenceImage(
31+
reference_id=1,
32+
reference_image=Image.from_file(location='test_resources/fruit_mask.png'),
33+
config=MaskReferenceConfig(
34+
mask_mode="MASK_MODE_USER_PROVIDED",
35+
mask_dilation=0.01,
36+
),
37+
)
38+
39+
image = client.models.edit_image(
40+
model="imagen-3.0-capability-001",
41+
prompt="",
42+
reference_images=[raw_ref, mask_ref],
43+
config=EditImageConfig(
44+
edit_mode="EDIT_MODE_INPAINT_REMOVAL",
45+
),
46+
)
47+
48+
image.generated_images[0].image.save(output_file)
49+
50+
print(f"Created output image using {len(image.generated_images[0].image.image_bytes)} bytes")
51+
# Example response:
52+
# Created output image using 1234567 bytes
53+
54+
# [END googlegenaisdk_imggen_inpainting_removal_mask_with_txt_img]
55+
return image.generated_images[0].image
56+
57+
58+
if __name__ == "__main__":
59+
edit_inpainting_removal_mask(output_file="test_resources/fruit_edit.png")

‎genai/image_generation/imggen_product_background_mask_with_txt_img.py

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
# Copyright 2025 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+
from google.genai.types import Image
16+
17+
18+
def edit_product_background_mask(output_file: str) -> Image:
19+
# [START googlegenaisdk_imggen_product_background_mask_with_txt_img]
20+
from google import genai
21+
from google.genai.types import RawReferenceImage, MaskReferenceImage, MaskReferenceConfig, EditImageConfig
22+
23+
client = genai.Client()
24+
25+
# TODO(developer): Update and un-comment below line
26+
# output_file = "output-image.png"
27+
28+
raw_ref = RawReferenceImage(
29+
reference_image=Image.from_file(location='test_resources/suitcase.png'), reference_id=0)
30+
mask_ref = MaskReferenceImage(
31+
reference_id=1,
32+
reference_image=Image.from_file(location='test_resources/suitcase_mask.png'),
33+
config=MaskReferenceConfig(
34+
mask_mode="MASK_MODE_USER_PROVIDED",
35+
mask_dilation=0.0,
36+
),
37+
)
38+
39+
image = client.models.edit_image(
40+
model="imagen-3.0-capability-001",
41+
prompt="A light blue suitcase in an airport",
42+
reference_images=[raw_ref, mask_ref],
43+
config=EditImageConfig(
44+
edit_mode="EDIT_MODE_BGSWAP",
45+
),
46+
)
47+
48+
image.generated_images[0].image.save(output_file)
49+
50+
print(f"Created output image using {len(image.generated_images[0].image.image_bytes)} bytes")
51+
# Example response:
52+
# Created output image using 1234567 bytes
53+
54+
# [END googlegenaisdk_imggen_product_background_mask_with_txt_img]
55+
return image.generated_images[0].image
56+
57+
58+
if __name__ == "__main__":
59+
edit_product_background_mask(output_file="test_resources/suitcase_edit.png")

‎genai/image_generation/test_image_generation.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,13 +25,16 @@
2525
import pytest
2626

2727
import imggen_canny_ctrl_type_with_txt_img
28+
import imggen_inpainting_insert_mask_with_txt_img
2829
import imggen_inpainting_insert_with_txt_img
30+
import imggen_inpainting_removal_mask_with_txt_img
2931
import imggen_inpainting_removal_with_txt_img
3032
import imggen_mask_free_edit_with_txt_img
3133
import imggen_mmflash_edit_img_with_txt_img
3234
import imggen_mmflash_txt_and_img_with_txt
3335
import imggen_mmflash_with_txt
3436
import imggen_outpainting_with_txt_img
37+
import imggen_product_background_mask_with_txt_img
3538
import imggen_product_background_with_txt_img
3639
import imggen_raw_reference_with_txt_img
3740
import imggen_scribble_ctrl_type_with_txt_img
@@ -70,6 +73,14 @@ def test_img_generation() -> None:
7073
assert response
7174

7275

76+
def test_img_edit_inpainting_insert_with_mask() -> None:
77+
OUTPUT_FILE = os.path.join(RESOURCES, "fruit_edit.png")
78+
response = imggen_inpainting_insert_mask_with_txt_img.edit_inpainting_insert_mask(
79+
OUTPUT_FILE
80+
)
81+
assert response
82+
83+
7384
def test_img_edit_inpainting_insert() -> None:
7485
OUTPUT_FILE = os.path.join(RESOURCES, "fruit_edit.png")
7586
response = imggen_inpainting_insert_with_txt_img.edit_inpainting_insert(
@@ -78,6 +89,14 @@ def test_img_edit_inpainting_insert() -> None:
7889
assert response
7990

8091

92+
def test_img_edit_inpainting_removal_mask() -> None:
93+
OUTPUT_FILE = os.path.join(RESOURCES, "fruit_edit.png")
94+
response = imggen_inpainting_removal_mask_with_txt_img.edit_inpainting_removal_mask(
95+
OUTPUT_FILE
96+
)
97+
assert response
98+
99+
81100
def test_img_edit_inpainting_removal() -> None:
82101
OUTPUT_FILE = os.path.join(RESOURCES, "fruit_edit.png")
83102
response = imggen_inpainting_removal_with_txt_img.edit_inpainting_removal(
@@ -86,6 +105,14 @@ def test_img_edit_inpainting_removal() -> None:
86105
assert response
87106

88107

108+
def test_img_edit_product_background_mask() -> None:
109+
OUTPUT_FILE = os.path.join(RESOURCES, "suitcase_edit.png")
110+
response = imggen_product_background_mask_with_txt_img.edit_product_background_mask(
111+
OUTPUT_FILE
112+
)
113+
assert response
114+
115+
89116
def test_img_edit_product_background() -> None:
90117
OUTPUT_FILE = os.path.join(RESOURCES, "suitcase_edit.png")
91118
response = imggen_product_background_with_txt_img.edit_product_background(

‎genai/image_generation/test_resources/fruit_mask.png

12.7 KB
Loading

‎genai/image_generation/test_resources/suitcase_mask.png

13 KB
Loading

0 commit comments

Comments
 (0)