From 33ebacb186c8111549dff31379931b82cbc9d286 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fr=C3=A9d=C3=A9ric?= Date: Sat, 11 May 2024 23:30:41 +0200 Subject: [PATCH] initial commit --- __init__.py | 7 ++ grid_node.py | 118 ++++++++++++++++++++++++++ map_node.py | 232 +++++++++++++++++++++++++++++++++++++++++++++++++++ utils.py | 15 ++++ 4 files changed, 372 insertions(+) create mode 100644 __init__.py create mode 100644 grid_node.py create mode 100644 map_node.py create mode 100644 utils.py diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..e007555 --- /dev/null +++ b/__init__.py @@ -0,0 +1,7 @@ +from .grid_node import BattlemapGrid +from .map_node import BattlemapMapGenerator + +NODE_CLASS_MAPPINGS = { + "Battlemap Grid": BattlemapGrid, + "Map Generator": BattlemapMapGenerator, +} diff --git a/grid_node.py b/grid_node.py new file mode 100644 index 0000000..5ae7e34 --- /dev/null +++ b/grid_node.py @@ -0,0 +1,118 @@ +import torch +from aggdraw import Draw, Brush, Pen +import math +from .utils import tensor_to_pil, pil_to_tensor + + +class BaseHexagonGenerator(object): + """ + Abstract classe for hexagon generators for hexagons of the specified size. + thanks to Elmer de Looff for the code. + https://variable-scope.com/posts/hexagon-tilings-with-pytho + """ + + def __init__(self, edge_length, center_x=0, center_y=0): + self.edge_length = edge_length + self.center_x = center_x + self.center_y = center_y + + @property + def col_width(self): + return self.edge_length * 3 + + @property + def row_height(self): + return math.sin(math.pi / 3) * self.edge_length + + +class HorizontalHexagonGenerator(BaseHexagonGenerator): + def __call__(self, row, col): + x = self.center_x + self.col_width / 3 + ( + col + 0.5 * (row % 2)) * self.col_width + y = self.center_y + row * self.row_height + for angle in range(0, 360, 60): + x += math.cos(math.radians(angle)) * self.edge_length + y += math.sin(math.radians(angle)) * self.edge_length + yield x + yield y + + +class VerticalHexagonGenerator(BaseHexagonGenerator): + def __call__(self, row, col): + x = self.center_x + row * self.row_height + y = self.center_y + self.col_width / 3 + ( + col + 0.5 * (row % 2)) * self.col_width + for angle in range(0, 360, 60): + x += math.sin(math.radians(angle)) * self.edge_length + y += math.cos(math.radians(angle)) * self.edge_length + yield x + yield y + + +class BattlemapGrid: + _grid_type = ["square", "vertical hexagon", "horizontal hexagon"] + + @classmethod + def INPUT_TYPES(cls): + return {"required": { + "image": ("IMAGE",), + "grid_type": (cls._grid_type, {"default": "square"}), + "grid_side": ("INT", {"default": 64, "min": 10, "max": 2048}), + "line_width": ("INT", {"default": 1, "min": 1, "max": 20}), + "red": ("INT", {"default": 255, "min": 0, "max": 255}), + "green": ("INT", {"default": 255, "min": 0, "max": 255}), + "blue": ("INT", {"default": 255, "min": 0, "max": 255}), + "alpha": ("INT", {"default": 255, "min": 0, "max": 255}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "grid_overlay" + CATEGORY = "Battlemaps" + + def square_grid(self, image, draw, center_x, center_y, + pen, grid_side): + for i in range(0, max(center_x, center_y), grid_side): + draw.line((center_x - i, 0, center_x - i, image.height), pen) + draw.line((center_x + i, 0, center_x + i, image.height), pen) + draw.line((0, center_y - i, image.width, center_y - i), pen) + draw.line((0, center_y + i, image.width, center_y + i), pen) + + def hexagon_grid(self, HexagonGenerator: BaseHexagonGenerator, + image, draw, center_x, center_y, + pen, grid_side): + hexagon_generator = HexagonGenerator(grid_side, center_x, center_y) + nb_row = math.ceil( + (image.height / hexagon_generator.row_height) / 2) + 2 + nb_col = math.ceil((image.width / hexagon_generator.col_width) / 2) + 2 + for row in range(-nb_row, nb_row): + for col in range(-nb_col, nb_col): + hexagon = hexagon_generator(row, col) + draw.polygon(list(hexagon), pen) + + def grid_overlay(self, image: torch.Tensor, + grid_type: str, grid_side: int, line_width: int, + red: int, green: int, blue: int, alpha: int): + image_pil = tensor_to_pil(image) + # Create Grid + + draw = Draw(image_pil) + pen = Pen((red, green, blue, alpha), line_width) + center_x, center_y = int(image_pil.width / 2), int(image_pil.height / 2) + if grid_type == "square": + self.square_grid(image_pil, draw, center_x, center_y, pen, + grid_side) + elif grid_type == "vertical hexagon": + self.hexagon_grid(VerticalHexagonGenerator, + image_pil, draw, center_x, center_y, + pen, grid_side) + elif grid_type == "horizontal hexagon": + self.hexagon_grid(HorizontalHexagonGenerator, + image_pil, draw, center_x, center_y, + pen, grid_side) + else: + raise Exception + draw.flush() + image_tensor_out = pil_to_tensor(image_pil) + return (image_tensor_out,) diff --git a/map_node.py b/map_node.py new file mode 100644 index 0000000..f4be9f4 --- /dev/null +++ b/map_node.py @@ -0,0 +1,232 @@ +from PIL import Image, ImageDraw, ImageColor, ImageFilter +import math +from colorsys import rgb_to_hsv, hsv_to_rgb +from random import random, randint, seed, choice +from .utils import pil_to_tensor + + +class BattlemapMapGenerator: + @classmethod + def INPUT_TYPES(cls): + return {"required": { + "node_seed": ( + "INT", {"default": 0, "min": -1, "max": 0xffffffffffffffff, + "step": 1, "label": "seed"}), + "grid_width": ("INT", {"default": 24, "min": 10, "max": 128}), + "grid_height": ("INT", {"default": 32, "min": 10, "max": 128}), + "grid_side": ("INT", {"default": 32, "min": 10, "max": 2048}), + "bg_color": ("STRING", {"default": "#FFFFFF"}), + "river": ( + "BOOLEAN", + {"default": True, "label_off": "OFF", "label_on": "ON"}), + "road": ( + "BOOLEAN", + {"default": True, "label_off": "OFF", "label_on": "ON"}), + "trees": ( + "BOOLEAN", + {"default": True, "label_off": "OFF", "label_on": "ON"}), + "rocks": ( + "BOOLEAN", + {"default": True, "label_off": "OFF", "label_on": "ON"}), + + } + } + + RETURN_TYPES = ("IMAGE", "STRING", "STRING", + "INT", "INT", + "INT", "INT", "INT") + RETURN_NAMES = ("image", "positive prompt", "negative prompt", + "image width", "image height", + "grid width", "grid height", "grid side") + FUNCTION = "map_generator" + CATEGORY = "Battlemaps" + + def generate_noise(self, image, x1, y1, x2, y2, padding=0): + def noise_color(i, j, color_index): + noise = randint(-128, 128) + return max(0, min(pixels[i, j][color_index] + noise, 255)) + + pixels = image.load() + for i in range(x1 - padding, x2 + padding): + for j in range(y1 - padding, y2 + padding): + try: + pixels[i, j] = ( + noise_color(i, j, 0), + noise_color(i, j, 1), + noise_color(i, j, 2) + ) + except IndexError: + pass + + def generate_image(self, width, height, bg_color): + image = Image.new("RGBA", (width, height), bg_color) + return image + + def generate_bg(self, image, draw, bg_color): + r, g, b = ImageColor.getcolor(bg_color, "RGB") + + for x in range(0, image.width + 10, 10): + for y in range(0, image.height + 10, 10): + color = (min(abs(r + randint(-100, 100)), 255), + min(abs(g + randint(-100, 100)), 255), + min(abs(b + randint(-100, 100)), 255)) + draw.ellipse([(x - 4, y - 4), (x + 4, y + 4)], fill=color) + self.generate_noise(image, 0, 0, image.width, image.height) + + def generator_flowers(self, image, draw): + pass + + def start_point(self, image): + match randint(0, 3): + case 0: + point = (randint(0, image.width), 0) + angle = randint(200, 340) + case 1: + point = (randint(0, image.width), image.height) + angle = randint(2, 160) + case 2: + point = (0, randint(0, image.height)) + angle = 70 - randint(0, 140) + case 3: + point = (image.width, randint(0, image.height)) + angle = randint(110, 250) + return point, angle + + def generate_path(self, image, draw, color, + start_point, start_angle, width=20, + depth=0, max_depth=3): + if depth > max_depth: + return + point1, angle, length = start_point, start_angle, randint(100, 150) + point_list = [point1] + while 0 <= point1[0] <= image.width or 0 <= point1[ + 1] <= image.height: + point2 = ( + int(point1[0] + length * math.cos(angle * math.pi / 180)), + int(point1[0] + length * math.sin(angle * math.pi / 180)) + ) + point_list.append(point2) + match randint(0, 50): + case 0: + break + case 1: + self.generate_path(image, draw, color, + point1, angle + randint(-10, 10), + max(5, width + randint(-5, 5)), + depth + 1, max_depth) + case 2: + draw.ellipse([(point2[0] - width, point2[1] - width), + (point2[0] + width, point2[1] + width)], + fill=color) + case _: + pass + point1, angle, length = (point2, angle + randint(-10, 10), + randint(20, 200)) + draw.line(point_list, fill=color, width=width, joint="curve") + + def generate_rivers(self, image, draw): + point, angle = self.start_point(image) + color = (0, 0, 255, 255) + self.generate_path(image, draw, color, point, angle) + + def generate_roads(self, image, draw): + point, angle = self.start_point(image) + color = "saddlebrown" + self.generate_path(image, draw, color, point, angle) + + def generate_stars(self, image, draw, size_color): + point = randint(0, image.width), randint(0, image.height) + size_multiplicator = random() + for size, color in size_color: + for angle in range(0, 360, 5): + l = max(size / 5, size * random() * size_multiplicator) + x = point[0] + math.ceil(l * math.cos(angle * math.pi / 180)) + y = point[1] + math.ceil(l * math.sin(angle * math.pi / 180)) + draw.line([point, (x, y)], + fill=color, width=7 + randint(-2, 2)) + self.generate_noise(image, point[0] - size, point[1] - size, + point[0] + size, point[1] + size, 10) + + def generate_polygon(self, image, draw, size_color, nb_point): + point = randint(0, image.width), randint(0, image.height) + size_multiplicator = max(0.5, random()) + for size, color in size_color: + point_list = list() + for angle in range(0, 360, int(360 / nb_point)): + l = max(size / 5, size * random() * size_multiplicator) + x = point[0] + math.ceil(l * math.cos(angle * math.pi / 180)) + y = point[1] + math.ceil(l * math.sin(angle * math.pi / 180)) + point_list.append((x, y)) + draw.polygon(point_list, fill=color, outline="black", width=2) + self.generate_noise(image, point[0] - size, point[1] - size, + point[0] + size, point[1] + size, 10) + + def generate_ellipses(self, image, draw, size_color): + point = randint(0, image.width), randint(0, image.height) + x_min, y_min = image.width, image.height + x_max, y_max = 0, 0 + size_multiplicator = random() + for size, color in size_color: + _size = math.ceil(size * size_multiplicator) + x1, y1 = (point[0] - _size + randint(-5, 5), + point[1] - _size + randint(-5, 5)) + x2, y2 = (point[0] + _size + randint(-5, 5), + point[1] + _size + randint(-5, 5)) + x_min, y_min = min(x_min, x1), min(y_min, y1) + x_max, y_max = max(x_max, x2), max(y_max, y2) + draw.ellipse( + [(min(x1, x2), min(y1, y2)), (max(x1, x2), max(y1, y2))], + fill=color, outline="black", width=2) + self.generate_noise(image, x_min, y_min, x_max, y_max, 10) + + def generate_rocks(self, image, draw): + for i in range(10): + self.generate_polygon(image, draw, [(60, "#111111"), + (50, "darkgray"), + (40, "gray")], + choice([3, 4, 5, 6, 8, 9, 10])) + + def generate_trees(self, image, draw): + for i in range(30): + self.generate_stars(image, draw, [(45, "darkgray"), + (40, "darkgreen"), + (30, "green"), + (20, "lightgreen")]) + + def map_generator(self, node_seed, grid_width, grid_height, + grid_side, bg_color, river, road, trees, rocks): + seed(node_seed) + positive, negative = list(), list() + width, height = grid_width * grid_side, grid_height * grid_side + image_pil = self.generate_image(width, height, bg_color) + draw = ImageDraw.Draw(image_pil) + self.generate_bg(image_pil, draw, bg_color) + + positive.append("BattleMap, outdoor, old medieval.") + positive.append("lightgreen background grass with small flowers") + if river: + self.generate_rivers(image_pil, draw) + positive.append("blue river, water") + else: + negative.append("blue river, water") + if road: + self.generate_roads(image_pil, draw) + positive.append("saddlebrown road") + else: + negative.append("road") + + if trees: + self.generate_rocks(image_pil, draw) + positive.append("gray rocks") + else: + negative.append("rocks") + + if rocks: + self.generate_trees(image_pil, draw) + positive.append("green trees") + else: + negative.append("trees") + + return (pil_to_tensor(image_pil), + ".\n".join(positive), ".\n".join(negative), + width, height, grid_width, grid_height, grid_side) diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..aed0f18 --- /dev/null +++ b/utils.py @@ -0,0 +1,15 @@ +import numpy as np +import torch +from PIL import Image + + +def tensor_to_pil(image_tensor: torch.tensor): + image_np = image_tensor.cpu().numpy() + image_pil = Image.fromarray((image_np.squeeze(0) * 255).astype(np.uint8)) + return image_pil.convert("RGBA") + + +def pil_to_tensor(image_pil: Image): + image_tensor_out = torch.tensor( + np.array(image_pil).astype(np.float32) / 255.0) + return torch.unsqueeze(image_tensor_out, 0)