use margin in draw_grid_annotations

This commit is contained in:
EllangoK 2023-02-05 03:04:18 -05:00
parent 95fc45b7e0
commit 5fdbc06fb4
3 changed files with 17 additions and 33 deletions

View File

@ -130,7 +130,7 @@ class GridAnnotation:
self.size = None
def draw_grid_annotations(im, width, height, hor_texts, ver_texts):
def draw_grid_annotations(im, width, height, hor_texts, ver_texts, margin=0):
def wrap(drawing, text, font, line_length):
lines = ['']
for word in text.split():
@ -194,48 +194,35 @@ def draw_grid_annotations(im, width, height, hor_texts, ver_texts):
line.allowed_width = allowed_width
hor_text_heights = [sum([line.size[1] + line_spacing for line in lines]) - line_spacing for lines in hor_texts]
ver_text_heights = [sum([line.size[1] + line_spacing for line in lines]) - line_spacing * len(lines) for lines in
ver_texts]
ver_text_heights = [sum([line.size[1] + line_spacing for line in lines]) - line_spacing * len(lines) for lines in ver_texts]
pad_top = 0 if sum(hor_text_heights) == 0 else max(hor_text_heights) + line_spacing * 2
result = Image.new("RGB", (im.width + pad_left, im.height + pad_top), "white")
result.paste(im, (pad_left, pad_top))
result = Image.new("RGB", (im.width + pad_left + margin * (cols-1), im.height + pad_top + margin * (rows-1)), "white")
for row in range(rows):
for col in range(cols):
cell = im.crop((width * col, height * row, width * (col+1), height * (row+1)))
result.paste(cell, (pad_left + (width + margin) * col, pad_top + (height + margin) * row))
d = ImageDraw.Draw(result)
for col in range(cols):
x = pad_left + width * col + width / 2
x = pad_left + (width + margin) * col + width / 2
y = pad_top / 2 - hor_text_heights[col] / 2
draw_texts(d, x, y, hor_texts[col], fnt, fontsize)
for row in range(rows):
x = pad_left / 2
y = pad_top + height * row + height / 2 - ver_text_heights[row] / 2
y = pad_top + (height + margin) * row + height / 2 - ver_text_heights[row] / 2
draw_texts(d, x, y, ver_texts[row], fnt, fontsize)
return result, pad_left, pad_top
return result
def draw_inner_margins(im, margin, rows, cols, pad_left, pad_top, margin_color=(255, 255, 255)):
if margin == 0: # Early exit, not strictly necessary
return im
cell_width = (im.width - pad_left) // cols
cell_height = (im.height - pad_top) // rows
padded_im = Image.new("RGB", (im.width + (cols - 1) * margin, im.height + (rows - 1) * margin), margin_color)
for r in range(rows):
for c in range(cols):
cell_location = (c * cell_width + pad_left * (c != 0),
r * cell_height + pad_top * (r != 0),
(c + 1) * cell_width + pad_left,
(r + 1) * cell_height + pad_top)
paste_location = (c * (cell_width + margin) + pad_left * (c != 0),
r * (cell_height + margin) + pad_top * (r != 0))
padded_im.paste(im.crop(cell_location), paste_location)
return padded_im
def draw_prompt_matrix(im, width, height, all_prompts):
def draw_prompt_matrix(im, width, height, all_prompts, margin=0):
prompts = all_prompts[1:]
boundary = math.ceil(len(prompts) / 2)
@ -245,7 +232,7 @@ def draw_prompt_matrix(im, width, height, all_prompts):
hor_texts = [[GridAnnotation(x, is_active=pos & (1 << i) != 0) for i, x in enumerate(prompts_horiz)] for pos in range(1 << len(prompts_horiz))]
ver_texts = [[GridAnnotation(x, is_active=pos & (1 << i) != 0) for i, x in enumerate(prompts_vert)] for pos in range(1 << len(prompts_vert))]
return draw_grid_annotations(im, width, height, hor_texts, ver_texts)
return draw_grid_annotations(im, width, height, hor_texts, ver_texts, margin)
def resize_image(resize_mode, im, width, height, upscaler_name=None):

View File

@ -33,7 +33,7 @@ def draw_xy_grid(xs, ys, x_label, y_label, cell):
res.append(processed.images[0])
grid = images.image_grid(res, rows=len(ys))
grid, _, _ = images.draw_grid_annotations(grid, res[0].width, res[0].height, hor_texts, ver_texts)
grid = images.draw_grid_annotations(grid, res[0].width, res[0].height, hor_texts, ver_texts)
first_processed.images = [grid]
@ -100,8 +100,7 @@ class Script(scripts.Script):
processed = process_images(p)
grid = images.image_grid(processed.images, p.batch_size, rows=1 << ((len(prompt_matrix_parts) - 1) // 2))
grid, pad_left, pad_top = images.draw_prompt_matrix(grid, p.width, p.height, prompt_matrix_parts)
grid = images.draw_inner_margins(grid, margin_size, grid.height // p.height, grid.width // p.width, pad_left, pad_top)
grid = images.draw_prompt_matrix(grid, p.width, p.height, prompt_matrix_parts, margin_size)
processed.images.insert(0, grid)
processed.index_of_first_image = 1
processed.infotexts.insert(0, processed.infotexts[0])

View File

@ -291,10 +291,8 @@ def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend
start_index = i * len(xs) * len(ys)
end_index = start_index + len(xs) * len(ys)
grid = images.image_grid(image_cache[start_index:end_index], rows=len(ys))
pad_left, pad_top = 0, 0
if draw_legend:
grid, pad_left, pad_top = images.draw_grid_annotations(grid, cell_size[0], cell_size[1], hor_texts, ver_texts)
grid = images.draw_inner_margins(grid, margin_size, len(ys), len(xs), pad_left, pad_top)
grid = images.draw_grid_annotations(grid, cell_size[0], cell_size[1], hor_texts, ver_texts, margin_size)
sub_grids[i] = grid
if include_sub_grids and len(zs) > 1:
processed_result.images.insert(i+1, grid)
@ -302,7 +300,7 @@ def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend
sub_grid_size = sub_grids[0].size
z_grid = images.image_grid(sub_grids, rows=1)
if draw_legend:
z_grid, _, _ = images.draw_grid_annotations(z_grid, sub_grid_size[0], sub_grid_size[1], title_texts, [[images.GridAnnotation()]])
z_grid = images.draw_grid_annotations(z_grid, sub_grid_size[0], sub_grid_size[1], title_texts, [[images.GridAnnotation()]])
processed_result.images[0] = z_grid
return processed_result, sub_grids