from PIL import Image, ImageDraw, ImageColor, ImageFilter import math from dataclasses import dataclass from colorsys import rgb_to_hsv, hsv_to_rgb from random import random, randint, choice from random import seed as rseed from .utils import pil_to_tensor, generate_noise from .point import Point from .path import Path @dataclass class BattlemapMapGenerator: @classmethod def INPUT_TYPES(cls) -> dict: return {"required": { "seed": ( "INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), "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": "#00FF00"}), } } @classmethod @property def RETURN_TYPES(cls) -> tuple: return ("IMAGE", "INT", "INT", "INT", "INT", "INT") @classmethod @property def RETURN_NAMES(cls) -> tuple: return ("image", "image width", "image height", "grid width", "grid height", "grid side") FUNCTION = "map_generator" CATEGORY = "Battlemaps" def generate_image(self, width: int, height: int, bg_color: str) -> Image: image = Image.new("RGBA", (width, height), bg_color) return image def generate_bg(self, image: Image, draw: ImageDraw.Draw, bg_color: str): r, g, b = ImageColor.getcolor(bg_color, "RGB") steps = 5 for x in range(0, image.width + steps, steps): for y in range(0, image.height + steps, steps): color = (min(abs(r + randint(-100, 100)), 255), min(abs(g + randint(-100, 100)), 255), min(abs(b + randint(-100, 100)), 255)) point1 = Point(x - int(steps / 2), y - int(steps / 2)) point2 = Point(x + int(steps / 2), y + int(steps / 2)) draw.ellipse([point1.coord(), point2.coord()], fill=color) generate_noise(image, Point(0, 0), Point(image.width, image.height)) def _map_generator(self, seed: int, grid_width: int, grid_height: int, grid_side: int, bg_color: str) -> tuple: rseed(seed) 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) return (image_pil, width, height, grid_width, grid_height, grid_side) def map_generator(self, seed: int, grid_width: int, grid_height: int, grid_side: int, bg_color: str) -> tuple: image_pil, width, height, grid_width, grid_height, grid_side = self._map_generator( seed, grid_width, grid_height, grid_side, bg_color) return (pil_to_tensor(image_pil), width, height, grid_width, grid_height, grid_side) class BattlemapMapGeneratorOutdoors(BattlemapMapGenerator): @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() inputs['required'].update({ "river": ( "BOOLEAN", {"default": True, "label_off": "OFF", "label_on": "ON"}), "road": ( "BOOLEAN", {"default": True, "label_off": "OFF", "label_on": "ON"}), "rocks": ( "BOOLEAN", {"default": True, "label_off": "OFF", "label_on": "ON"}), "trees": ( "BOOLEAN", {"default": True, "label_off": "OFF", "label_on": "ON"}), "positive": ("STRING", {"default": '', "multiline": True}), "negative": ("STRING", {"default": '', "multiline": True}), }) return inputs @classmethod @property def RETURN_TYPES(cls) -> tuple: return_types = list(super().RETURN_TYPES) return_types.extend(["STRING", "STRING"]) return tuple(return_types) @classmethod @property def RETURN_NAMES(cls) -> tuple: return_names = list(super().RETURN_NAMES) return_names.extend(["positive prompt", "negative prompt"]) return tuple(return_names) def generator_flowers(self, image: Image, draw: ImageDraw.Draw): for _ in range(100, 500): point = Point(randint(0, image.width), randint(0, image.height)) size = randint(2, 6) draw.ellipse([(point - size).coord(), (point + size).coord()], fill=choice(list(ImageColor.colormap.keys())), outline="black", width=2) def start_point(self, image: Image) -> tuple[Point, int]: match randint(0, 3): case 0: point = Point(randint(0, image.width), -10) case 1: point = Point(randint(0, image.width), image.height + 10) case 2: point = Point(-10, randint(0, image.height)) case 3: point = Point(image.width + 10, randint(0, image.height)) angle = point.angle_between(Point(image.width / 2, image.height / 2)) return point, angle + randint(-45, 45) def generate_path(self, image: Image, paths: list[Path], start_point: Point, start_angle: int, width: int = 20, depth: int = 0, depth_max: int = 5): if depth > depth_max or width <= 5: return point1, angle, length = start_point, start_angle, randint(100, 150) path = Path(width=width) path.add_point(point1) while -10 <= point1.x <= image.width + 10 and -10 <= point1.y <= image.height + 10: point2 = point1.add_polar(length, angle) path.add_point(point2) match randint(0, 20): case 0: break case 1 | 2: new_angle = angle + randint(10, 30) * choice([-1, 1]) self.generate_path(image, paths, point1, new_angle, width=math.ceil(width * 2 / 3), depth=depth + 1, depth_max=depth_max) point1, angle, length = (point2, angle + randint(-20, 20), randint(20, 200)) paths.append(path) def draw_path(self, draw: ImageDraw.Draw, color: str, paths: list): width_max = 0 for path in paths: width_max = max(width_max, path.width) draw.line(path.coord(), fill="black", width=path.width + 2, joint="curve") for i in range(width_max, 0, -1): h, s, v = rgb_to_hsv(*ImageColor.getcolor(color, "RGB")) r, g, b = map(int, hsv_to_rgb(h, s, v / i ** 0.3)) for path in filter(lambda x: x.width <= width_max, paths): draw.line(path.coord(), fill=(r, g, b), width=i, joint="curve") def generate_rivers(self, image: Image, draw: ImageDraw.Draw): point, angle = self.start_point(image) paths = list() self.generate_path(image, paths, point, angle) self.draw_path(draw, "blue", paths) def generate_roads(self, image: Image, draw: ImageDraw.Draw): point, angle = self.start_point(image) paths = list() self.generate_path(image, paths, point, angle) self.draw_path(draw, "saddlebrown", paths) def generate_stars(self, image: Image, draw: ImageDraw.Draw, size_color: list): center = 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): length = max(size / 5, size * random() * size_multiplicator) point = center.add_polar(length, angle) witdh = randint(5, 10) draw.line([center.coord(), point.coord()], fill="black", width=witdh + 2) draw.line([center.coord(), point.coord()], fill=color, width=witdh) generate_noise(image, center - size, center + size, 10) def generate_polygon(self, image: Image, draw: ImageDraw.Draw, size_color: list, nb_point: int): center = 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)): length = max(size / 5, size * random() * size_multiplicator) point = center.add_polar(length, angle) point_list.append(point.coord()) draw.polygon(point_list, fill=color, outline="black", width=2) generate_noise(image, center - size, center + size, 10) def generate_ellipses(self, image: Image, draw: ImageDraw.Draw, size_color: list): center = Point(randint(0, image.width), randint(0, image.height)) size_multiplicator = random() size_max = 0 for size, color in size_color: size_max = max(size_max, math.ceil(size * size_multiplicator)) point1 = center - math.ceil(size * size_multiplicator) - randint(0, 10) point2 = center + math.ceil(size * size_multiplicator) + randint(0, 10) draw.ellipse([point1.coord(), point2.coord()], fill=color, outline="black", width=2) generate_noise(image, center - size_max, center + size_max, 10) def generate_rocks(self, image: Image, draw: ImageDraw.Draw): for _ in range(10): self.generate_polygon(image, draw, [(50, "#111111"), (35, "darkgray"), (25, "gray")], choice([3, 4, 5, 6, 8, 9])) def generate_trees(self, image: Image, draw: ImageDraw.Draw): for _ in range(100): self.generate_stars(image, draw, [(45, "darkgray"), (40, "darkgreen"), (30, "green"), (20, "lightgreen")]) def map_generator(self, seed: int, grid_width: int, grid_height: int, grid_side: int, bg_color: str, river: bool=False, road: bool=False, rocks: bool=False, trees: bool=False, positive: str = "", negative: str = ""): rseed(seed) image, width, height, grid_width, grid_height, grid_side = ( self._map_generator(seed, grid_width, grid_height, grid_side, bg_color)) draw = ImageDraw.Draw(image) self.generator_flowers(image, draw) if river: self.generate_rivers(image, draw) positive += "blue river, water." else: negative += "blue river, water." if road: self.generate_roads(image, draw) positive += "saddlebrown road." else: negative += "road." if rocks: self.generate_rocks(image, draw) positive += "green trees." else: negative += "trees." if trees: self.generate_trees(image, draw) positive += "gray rocks." else: negative += "rocks." return (pil_to_tensor(image), width, height, grid_width, grid_height, grid_side, positive, negative)