feature: 🎉 outpainting
Bryce
1 year ago
committed by
Bryce Drennan
@ -0,0 +1,118 @@
|
||||
import re
|
||||
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
|
||||
def prepare_image_for_outpaint(
|
||||
img, mask=None, up=None, down=None, left=None, right=None, _all=0, snap_multiple=64
|
||||
):
|
||||
up = up if up is not None else _all
|
||||
down = down if down is not None else _all
|
||||
left = left if left is not None else _all
|
||||
right = right if right is not None else _all
|
||||
|
||||
lft_pct = left / (left + right)
|
||||
rgt_pct = right / (left + right)
|
||||
up_pct = up / (up + down)
|
||||
dwn_pct = down / (up + down)
|
||||
|
||||
new_width = round((img.width + left + right) / snap_multiple) * snap_multiple
|
||||
new_height = round((img.height + up + down) / snap_multiple) * snap_multiple
|
||||
height_addition = max(new_height - img.height, 0)
|
||||
width_addition = max(new_width - img.width, 0)
|
||||
up = int(round(height_addition * up_pct))
|
||||
down = int(round(height_addition * dwn_pct))
|
||||
left = int(round(width_addition * lft_pct))
|
||||
right = int(round(width_addition * rgt_pct))
|
||||
|
||||
expanded_image = Image.new(
|
||||
"RGB", (img.width + left + right, img.height + up + down), (0, 0, 0)
|
||||
)
|
||||
expanded_image.paste(img, (left, up))
|
||||
|
||||
# extend border pixels outward, this helps prevents lines at the boundary because masks getting reduced to
|
||||
# 64x64 latent space can cause som inaccuracies
|
||||
|
||||
if up > 0:
|
||||
expanded_image.paste(
|
||||
img.crop((0, 0, img.width, 1)).resize((expanded_image.width, up)),
|
||||
(0, 0),
|
||||
)
|
||||
expanded_image.paste(
|
||||
img.crop((0, 0, img.width, 1)).resize((img.width, up)),
|
||||
(left, 0),
|
||||
)
|
||||
if down > 0:
|
||||
expanded_image.paste(
|
||||
img.crop((0, img.height - 1, img.width, img.height)).resize(
|
||||
(expanded_image.width, down)
|
||||
),
|
||||
(0, expanded_image.height - down),
|
||||
)
|
||||
expanded_image.paste(
|
||||
img.crop((0, img.height - 1, img.width, img.height)).resize(
|
||||
(img.width, down)
|
||||
),
|
||||
(left, expanded_image.height - down),
|
||||
)
|
||||
if left > 0:
|
||||
expanded_image.paste(
|
||||
img.crop((0, 0, 1, img.height)).resize((left, expanded_image.height)),
|
||||
(0, 0),
|
||||
)
|
||||
expanded_image.paste(
|
||||
img.crop((0, 0, 1, img.height)).resize((left, img.height)),
|
||||
(0, up),
|
||||
)
|
||||
if right > 0:
|
||||
expanded_image.paste(
|
||||
img.crop((img.width - 1, 0, img.width, img.height)).resize(
|
||||
(right, expanded_image.height)
|
||||
),
|
||||
(expanded_image.width - right, 0),
|
||||
)
|
||||
expanded_image.paste(
|
||||
img.crop((img.width - 1, 0, img.width, img.height)).resize(
|
||||
(right, img.height)
|
||||
),
|
||||
(expanded_image.width - right, up),
|
||||
)
|
||||
|
||||
# create a mask for the new boundaries
|
||||
expanded_mask = Image.new("L", (expanded_image.width, expanded_image.height), 255)
|
||||
if mask is None:
|
||||
draw = ImageDraw.Draw(expanded_mask)
|
||||
draw.rectangle(
|
||||
(left, up, left + img.width, up + img.height), fill="black", outline="black"
|
||||
)
|
||||
else:
|
||||
expanded_mask.paste(mask, (left, up))
|
||||
|
||||
return expanded_image, expanded_mask
|
||||
|
||||
|
||||
def outpaint_arg_str_parse(arg_str):
|
||||
arg_pattern = re.compile(r"([A-Z]+)(\d+)")
|
||||
|
||||
args = arg_str.upper().split(",")
|
||||
valid_directions = ["up", "down", "left", "right", "all"]
|
||||
valid_direction_chars = {c[0]: c for c in valid_directions}
|
||||
kwargs = {}
|
||||
for arg in args:
|
||||
match = arg_pattern.match(arg)
|
||||
if not match:
|
||||
raise ValueError(f"Invalid outpaint argument '{arg}'")
|
||||
direction, amount = match.groups()
|
||||
direction = direction.lower()
|
||||
if len(direction) == 1:
|
||||
if direction not in valid_direction_chars:
|
||||
raise ValueError(f"Invalid outpaint direction '{direction}'")
|
||||
direction = valid_direction_chars[direction]
|
||||
elif direction not in valid_directions:
|
||||
raise ValueError(f"Invalid outpaint direction '{direction}'")
|
||||
kwargs[direction] = int(amount)
|
||||
|
||||
if "all" in kwargs:
|
||||
kwargs["_all"] = kwargs.pop("all")
|
||||
|
||||
return kwargs
|
||||
|
Before Width: | Height: | Size: 323 KiB After Width: | Height: | Size: 325 KiB |
|
Before Width: | Height: | Size: 320 KiB After Width: | Height: | Size: 318 KiB |
|
Before Width: | Height: | Size: 257 KiB After Width: | Height: | Size: 256 KiB |
|
Before Width: | Height: | Size: 254 KiB After Width: | Height: | Size: 252 KiB |
|
Before Width: | Height: | Size: 251 KiB After Width: | Height: | Size: 248 KiB |
|
Before Width: | Height: | Size: 245 KiB After Width: | Height: | Size: 244 KiB |
|
Before Width: | Height: | Size: 259 KiB After Width: | Height: | Size: 259 KiB |
|
After Width: | Height: | Size: 326 KiB |
@ -0,0 +1,40 @@
|
||||
import pytest
|
||||
|
||||
from imaginairy import ImaginePrompt, LazyLoadingImage, imagine
|
||||
from imaginairy.outpaint import outpaint_arg_str_parse
|
||||
from tests import TESTS_FOLDER
|
||||
from tests.utils import assert_image_similar_to_expectation
|
||||
|
||||
|
||||
def test_outpainting_outpaint(filename_base_for_outputs):
|
||||
img = LazyLoadingImage(
|
||||
filepath=f"{TESTS_FOLDER}/data/girl_with_a_pearl_earring.jpg"
|
||||
)
|
||||
prompt = ImaginePrompt(
|
||||
prompt="woman standing",
|
||||
init_image=img,
|
||||
init_image_strength=0,
|
||||
mask_prompt="background",
|
||||
outpaint="all250,up0,down600",
|
||||
mask_mode="replace",
|
||||
negative_prompt="picture frame, borders, framing, text, writing, watermarks, indoors, advertisement, paper, canvas, stock photo",
|
||||
steps=20,
|
||||
seed=542906833,
|
||||
)
|
||||
result = list(imagine([prompt]))[0]
|
||||
img_path = f"{filename_base_for_outputs}.png"
|
||||
assert_image_similar_to_expectation(result.img, img_path=img_path, threshold=2800)
|
||||
|
||||
|
||||
outpaint_test_params = [
|
||||
("A132", {"_all": 132}),
|
||||
("A132,U50", {"_all": 132, "up": 50}),
|
||||
("A132,U50,D50", {"_all": 132, "up": 50, "down": 50}),
|
||||
("a132,u50,d50", {"_all": 132, "up": 50, "down": 50}),
|
||||
("all50,up20,down600", {"_all": 50, "up": 20, "down": 600}),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("arg_str, expected_kwargs", outpaint_test_params)
|
||||
def test_outpaint_parse_kwargs(arg_str, expected_kwargs):
|
||||
assert outpaint_arg_str_parse(arg_str) == expected_kwargs
|
||||