diff --git a/__init__.py b/__init__.py index d84184d..ef94565 100644 --- a/__init__.py +++ b/__init__.py @@ -1,4 +1,10 @@ -from .src import LatentCombineNode, ImagesGridByColumnsNode, ImagesGridByRowsNode, ImageCombineNode +from .src import ( + LatentCombineNode, + ImagesGridByColumnsNode, + ImagesGridByRowsNode, + ImageCombineNode, + GridAnnotationNode, +) NODE_CLASS_MAPPINGS = { @@ -6,4 +12,5 @@ NODE_CLASS_MAPPINGS = { "ImagesGridByColumns": ImagesGridByColumnsNode, "ImagesGridByRows": ImagesGridByRowsNode, "ImageCombine": ImageCombineNode, + "GridAnnotation": GridAnnotationNode, } diff --git a/src/__init__.py b/src/__init__.py index dbab28c..225673c 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -1,3 +1,4 @@ from .nodes.images_grid import ImagesGridByColumnsNode, ImagesGridByRowsNode from .nodes.latent_combine import LatentCombineNode from .nodes.image_combine import ImageCombineNode +from .nodes.grid_annotation import GridAnnotationNode diff --git a/src/base.py b/src/base.py index c06dc3f..54e149a 100644 --- a/src/base.py +++ b/src/base.py @@ -1,9 +1,12 @@ import typing as t +from pathlib import Path + + +STATIC_PATH = Path(__file__).parent.parent / "static" + +Image = t.Any class BaseNode(): CATEGORY: str = "ImagesGrid" FUNCTION: str = "execute" - - -Image = t.Any diff --git a/src/nodes/grid_annotation.py b/src/nodes/grid_annotation.py new file mode 100644 index 0000000..bc1851a --- /dev/null +++ b/src/nodes/grid_annotation.py @@ -0,0 +1,49 @@ +import typing as t + +from PIL import ImageFont + +from ..base import BaseNode, STATIC_PATH +from ..utils import Annotation + + +class GridAnnotationNode(BaseNode): + RETURN_TYPES: tuple[str] = ("GRID_ANNOTATION",) + + @classmethod + def INPUT_TYPES(cls) -> dict[str, t.Any]: + return { + "required": { + "column_texts": ("STRING", {"multiline": False}), + "row_texts": ("STRING", {"multiline": False}), + "font_size": ("INT", {"default": 30, "min": 1}), + }, + } + + def execute( + self, + column_texts: str, + row_texts: str, + font_size: int, + ) -> tuple[Annotation]: + font = ImageFont.truetype(str(STATIC_PATH / "Roboto-Regular.ttf"), size=font_size) + column_texts_list = self._set_value_to_texts_list( + self._get_texts_from_string(column_texts), + ) + row_texts_list = self._set_value_to_texts_list( + self._get_texts_from_string(row_texts), + ) + + result = Annotation(column_texts=column_texts_list, row_texts=row_texts_list, font=font) + return (result,) + + def _get_texts_from_string(self, string: str) -> list[str]: + return [ + result + for i in string.split(";") + if (result := i.strip()) != "" + ] + + def _set_value_to_texts_list(self, texts_list: list[str]) -> list[str]: + if not texts_list: + return ["None"] + return texts_list diff --git a/src/nodes/image_combine.py b/src/nodes/image_combine.py index ce7fde0..f7a9d46 100644 --- a/src/nodes/image_combine.py +++ b/src/nodes/image_combine.py @@ -6,10 +6,10 @@ from ..base import BaseNode, Image class ImageCombineNode(BaseNode): - RETURN_TYPES: t.Tuple[str] = ("IMAGE",) + RETURN_TYPES: tuple[str] = ("IMAGE",) @classmethod - def INPUT_TYPES(cls) -> t.Dict[str, t.Any]: + def INPUT_TYPES(cls) -> dict[str, t.Any]: return { "required": { "image_1": ("IMAGE",), @@ -21,12 +21,7 @@ class ImageCombineNode(BaseNode): self, image_1: Image, image_2: Image, - ) -> t.Tuple[Image]: - print(image_1.size()) - print(image_2.size()) - print(image_1) - + ) -> tuple[Image]: result = torch.cat((image_1, image_2), 0) - print(result.size()) return (result,) diff --git a/src/nodes/images_grid.py b/src/nodes/images_grid.py index 4bf1cfc..40783c6 100644 --- a/src/nodes/images_grid.py +++ b/src/nodes/images_grid.py @@ -6,25 +6,42 @@ from ..utils import ( pillow_to_tensor, create_images_grid_by_columns, create_images_grid_by_rows, + Annotation, ) class BaseImagesGridNode(BaseNode): - RETURN_TYPES: t.Tuple[str] = ("IMAGE",) + RETURN_TYPES: tuple[str] = ("IMAGE",) @classmethod - def _create_input_types(cls, coordinate_name: str) -> t.Dict[str, t.Any]: + def _create_input_types(cls, coordinate_name: str) -> dict[str, t.Any]: return { "required": { "images": ("IMAGE",), "gap": ("INT", {"default": 0, "min": 0}), coordinate_name: ("INT", {"default": 1, "min": 1}), + }, + "optional": { + "annotation": ("GRID_ANNOTATION",), } } - def _create_execute(self, images, function, function_kw) -> t.Tuple[Image]: + def _create_execute( + self, + function: t.Callable, + \ + images: Image, + gap: int, + annotation: Annotation | None = None, + **kw, + ) -> tuple[Image]: pillow_images = [tensor_to_pillow(i) for i in images] - pillow_grid = function(images=pillow_images, **function_kw) + pillow_grid = function( + images=pillow_images, + gap=gap, + annotation=annotation, + **kw, + ) tensor_grid = pillow_to_tensor(pillow_grid) return (tensor_grid,) @@ -32,17 +49,17 @@ class BaseImagesGridNode(BaseNode): class ImagesGridByColumnsNode(BaseImagesGridNode): @classmethod - def INPUT_TYPES(cls) -> t.Dict[str, t.Any]: + def INPUT_TYPES(cls) -> dict[str, t.Any]: return cls._create_input_types("max_columns") - def execute(self, images: Image, **kw) -> tuple[Image]: - return self._create_execute(images, create_images_grid_by_columns, kw) + def execute(self, **kw) -> tuple[Image]: + return self._create_execute(create_images_grid_by_columns, **kw) class ImagesGridByRowsNode(BaseImagesGridNode): @classmethod - def INPUT_TYPES(cls) -> t.Dict[str, t.Any]: + def INPUT_TYPES(cls) -> dict[str, t.Any]: return cls._create_input_types("max_rows") - def execute(self, images: Image, **kw) -> tuple[Image]: - return self._create_execute(images, create_images_grid_by_rows, kw) + def execute(self, **kw) -> tuple[Image]: + return self._create_execute(create_images_grid_by_rows, **kw) diff --git a/src/nodes/latent_combine.py b/src/nodes/latent_combine.py index 96e0259..69bb399 100644 --- a/src/nodes/latent_combine.py +++ b/src/nodes/latent_combine.py @@ -6,10 +6,10 @@ from ..base import BaseNode, Image class LatentCombineNode(BaseNode): - RETURN_TYPES: t.Tuple[str] = ("LATENT",) + RETURN_TYPES: tuple[str] = ("LATENT",) @classmethod - def INPUT_TYPES(cls) -> t.Dict[str, t.Any]: + def INPUT_TYPES(cls) -> dict[str, t.Any]: return { "required": { "latent_1": ("LATENT",), @@ -19,9 +19,9 @@ class LatentCombineNode(BaseNode): def execute( self, - latent_1: t.Dict[str, t.Any], - latent_2: t.Dict[str, t.Any], - ) -> t.Tuple[t.Dict[str, t.Any]]: + latent_1: dict[str, t.Any], + latent_2: dict[str, t.Any], + ) -> tuple[dict[str, t.Any]]: samples = torch.cat((latent_1["samples"], latent_2["samples"]), 0) return ({"samples": samples},) diff --git a/src/utils/__init__.py b/src/utils/__init__.py index b8706a0..dc39679 100644 --- a/src/utils/__init__.py +++ b/src/utils/__init__.py @@ -1,6 +1,6 @@ from .images_grid import ( create_images_grid_by_columns, create_images_grid_by_rows, + Annotation, ) -from grid_annotations import create_grid_annotations from .tensor_convert import tensor_to_pillow, pillow_to_tensor diff --git a/src/utils/grid_annotations.py b/src/utils/grid_annotations.py deleted file mode 100644 index a7c24f8..0000000 --- a/src/utils/grid_annotations.py +++ /dev/null @@ -1,104 +0,0 @@ -import typing as t -from dataclasses import dataclass - -from PIL import Image, ImageDraw, ImageFont - - -@dataclass -class GridImagesInfo(): - image: Image.Image - gap: int - one_image_size: t.Tuple[int, int] - - -def create_grid_annotations( - grid_info: GridImagesInfo, - column_texts, - row_texts, - font: ImageFont.ImageFont, - margin: int = 5, -) -> Image.Image: - grid = grid_info.image - left_padding = int(max(map(font.getlength, row_texts))) + 2*margin # type: ignore - top_padding = font.size + 2*margin # type: ignore - - image = Image.new( - "RGB", - (grid.size[0] + left_padding, grid.size[1] + top_padding), - color="white", - ) - draw = ImageDraw.Draw(image) - draw.font = font # type: ignore - - image.paste(grid, (image.size[0] - grid.size[0], image.size[1] - grid.size[1])) - _draw_column_text( - draw=draw, - texts=column_texts, - grid_info=grid_info, - left_padding=left_padding, - top_padding=top_padding, - ) - _draw_row_text( - draw=draw, - texts=row_texts, - grid_info=grid_info, - left_padding=left_padding, - top_padding=top_padding, - ) - - return image - - -def _draw_column_text( - draw: ImageDraw.ImageDraw, - texts: t.List[str], - grid_info: GridImagesInfo, - left_padding: int, - top_padding: int, -) -> None: - i = 0 - x0 = left_padding - y0 = 0 - x1 = left_padding + grid_info.one_image_size[0] - y1 = top_padding - while x0 != grid_info.image.size[0] + left_padding: - _draw_center_text(draw, (x0, y0, x1, y1), texts[i]) - x0 += grid_info.one_image_size[0] + grid_info.gap - x1 += grid_info.one_image_size[0] + grid_info.gap - i += 1 - - -def _draw_row_text( - draw: ImageDraw.ImageDraw, - texts: t.List[str], - grid_info: GridImagesInfo, - left_padding: int, - top_padding: int, -) -> None: - i = 0 - x0 = 0 - y0 = top_padding - x1 = left_padding - y1 = top_padding + grid_info.one_image_size[1] - while y0 != grid_info.image.size[1] + top_padding: - _draw_center_text(draw, (x0, y0, x1, y1), texts[i]) - y0 += grid_info.one_image_size[1] + grid_info.gap - y1 += grid_info.one_image_size[1] + grid_info.gap - i += 1 - - -def _draw_center_text( - draw: ImageDraw.ImageDraw, - xy: t.Tuple[int, int, int, int], - text: str, - fill: t.Any = "black", -) -> None: - _, _, *text_size = draw.textbbox((0, 0), text) - draw.text( - ( - (xy[2] - text_size[0] + xy[0]) / 2, - (xy[3] - text_size[1] + xy[1]) / 2, - ), - text, - fill=fill, - ) diff --git a/src/utils/images_grid.py b/src/utils/images_grid.py index 9b2e8e3..9a186b8 100644 --- a/src/utils/images_grid.py +++ b/src/utils/images_grid.py @@ -1,39 +1,81 @@ import typing as t +from dataclasses import dataclass +from contextlib import suppress from PIL import Image, ImageDraw, ImageFont +@dataclass +class Annotation(): + column_texts: list[str] + row_texts: list[str] + font: ImageFont.FreeTypeFont + + +@dataclass +class _GridInfo(): + image: Image.Image + gap: int + one_image_size: tuple[int, int] + + def create_images_grid_by_columns( - images: t.List[Image.Image], + images: list[Image.Image], gap: int, max_columns: int, + annotation: Annotation | None = None, ) -> Image.Image: max_rows = (len(images) + max_columns - 1) // max_columns - return _create_images_grid(images=images, gap=gap, max_columns=max_columns, max_rows=max_rows) + return _create_images_grid(images, gap, max_columns, max_rows, annotation) def create_images_grid_by_rows( - images: t.List[Image.Image], + images: list[Image.Image], gap: int, max_rows: int, + annotation: Annotation | None = None, ) -> Image.Image: max_columns = (len(images) + max_rows - 1) // max_rows - return _create_images_grid(images=images, gap=gap, max_columns=max_columns, max_rows=max_rows) + return _create_images_grid(images, gap, max_columns, max_rows, annotation) def _create_images_grid( - images: t.List[Image.Image], + images: list[Image.Image], gap: int, max_columns: int, max_rows: int, + annotation: Annotation | None, ) -> Image.Image: size = images[0].size - grid_width = size[0] * max_columns + (max_columns - 1) * gap grid_height = size[1] * max_rows + (max_rows - 1) * gap grid_image = Image.new("RGB", (grid_width, grid_height), color="white") + _arrange_images_on_grid(grid_image, images=images, size=size, max_columns=max_columns, gap=gap) + + if annotation is None: + return grid_image + return _create_grid_annotations( + grid_info=_GridInfo( + image=grid_image, + gap=gap, + one_image_size=size, + ), + column_texts=annotation.column_texts, + row_texts=annotation.row_texts, + font=annotation.font, + ) + + +def _arrange_images_on_grid( + grid_image: Image.Image, + /, + images: list[Image.Image], + size: tuple[int, int], + max_columns: int, + gap: int, +): for i, image in enumerate(images): image = image.crop((0, 0, size[0], size[1])) x = (i % max_columns) * (size[0] + gap) @@ -41,24 +83,112 @@ def _create_images_grid( grid_image.paste(image, (x, y)) - return grid_image + +def _create_grid_annotations( + grid_info: _GridInfo, + column_texts, + row_texts, + font: ImageFont.FreeTypeFont, +) -> Image.Image: + if not column_texts or not row_texts: + raise ValueError("Column text or row text is empty") + + grid = grid_info.image + margin = font.size // 2 + left_padding = int(max(map(font.getlength, row_texts))) + 2*margin + top_padding = font.size + 2*margin + + image = Image.new( + "RGB", + (grid.size[0] + left_padding, grid.size[1] + top_padding), + color="white", + ) + draw = ImageDraw.Draw(image) + draw.font = font # type: ignore + + _paste_image_to_lower_left_corner(image, grid) + _draw_column_text( + draw=draw, + texts=column_texts, + grid_info=grid_info, + left_padding=left_padding, + top_padding=top_padding, + ) + _draw_row_text( + draw=draw, + texts=row_texts, + grid_info=grid_info, + left_padding=left_padding, + top_padding=top_padding, + ) + + return image + + +def _draw_column_text( + draw: ImageDraw.ImageDraw, + texts: list[str], + grid_info: _GridInfo, + left_padding: int, + top_padding: int, +) -> None: + i = 0 + x0 = left_padding + y0 = 0 + x1 = left_padding + grid_info.one_image_size[0] + y1 = top_padding + while x0 != grid_info.image.size[0] + left_padding + grid_info.gap: + i = _draw_text_by_xy((x0, y0, x1, y1), i, draw=draw, texts=texts) + x0 += grid_info.one_image_size[0] + grid_info.gap + x1 += grid_info.one_image_size[0] + grid_info.gap + + +def _draw_row_text( + draw: ImageDraw.ImageDraw, + texts: list[str], + grid_info: _GridInfo, + left_padding: int, + top_padding: int, +) -> None: + i = 0 + x0 = 0 + y0 = top_padding + x1 = left_padding + y1 = top_padding + grid_info.one_image_size[1] + while y0 != grid_info.image.size[1] + top_padding + grid_info.gap: + i = _draw_text_by_xy((x0, y0, x1, y1), i, draw=draw, texts=texts) + y0 += grid_info.one_image_size[1] + grid_info.gap + y1 += grid_info.one_image_size[1] + grid_info.gap + + +def _draw_text_by_xy( + xy: tuple[int, int, int, int], + index: int, + \ + draw: ImageDraw.ImageDraw, + texts: list[str], +) -> int: + with suppress(IndexError): + _draw_center_text(draw, xy, texts[index]) + return index + 1 def _draw_center_text( - text: str, - draw: ImageDraw.ImageDraw, - font: ImageFont.ImageFont, - fill: int = 128, -): - image = draw.im # type: ignore - _, _, *text_size = draw.textbbox((0, 0), text, font=font) - draw.text( - ( - (image.size[0]-text_size[0])/2, - (image.size[1]-text_size[1])/2, - ), - text, - font=font, - fill=fill, - ) - return image + draw: ImageDraw.ImageDraw, + xy: tuple[int, int, int, int], + text: str, + fill: t.Any = "black", +) -> None: + _, _, *text_size = draw.textbbox((0, 0), text) + draw.text( + ( + (xy[2] - text_size[0] + xy[0]) / 2, + (xy[3] - text_size[1] + xy[1]) / 2, + ), + text, + fill=fill, + ) + + +def _paste_image_to_lower_left_corner(base: Image.Image, image: Image.Image) -> None: + base.paste(image, (base.size[0] - image.size[0], base.size[1] - image.size[1])) diff --git a/static/Roboto-Regular.ttf b/static/Roboto-Regular.ttf new file mode 100644 index 0000000..3033308 Binary files /dev/null and b/static/Roboto-Regular.ttf differ