diff --git a/compass_node.py b/compass_node.py index b888da0..ce30019 100644 --- a/compass_node.py +++ b/compass_node.py @@ -1,13 +1,18 @@ import torch -from PIL import ImageDraw, ImageFont +from PIL.ImageDraw import Draw +from PIL import ImageFont from .utils import tensor_to_pil, pil_to_tensor from .point import Point - -BLACK = (0, 0, 0, 128) -WHITE = (255, 255, 255, 128) +import random class CompassGrid: + + def __init__(self): + self.black_pen = (0, 0, 0) + self.white_pen = (255, 255, 255) + self.gray_pen = (80, 80, 80) + @classmethod def INPUT_TYPES(cls): return { @@ -16,8 +21,16 @@ class CompassGrid: "font": ("STRING", { "default": "/System/Library/Fonts/Supplemental/Arial.ttf"}), "cardinals": (["NWSE", "NOSE"], {"default": "NWSE"}), + "position": ( + ["top left", "top right", "bottom left", "bottom right"], + {"default": "bottom right"}), "size": ("INT", {"default": 64, "min": 10, "max": 2048}), "rotation": ("INT", {"default": 0, "min": 0, "max": 360}), + "seed": ( + "INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "random_rotation": ("BOOLEAN", + {"default": True, "label_off": "OFF", + "label_on": "ON"}), }, } @@ -26,39 +39,67 @@ class CompassGrid: FUNCTION = "compass_overlay" CATEGORY = "Battlemaps" - def draw_arrow(self, draw: ImageDraw.Draw, font, center: Point, + def get_center_position(self, image_pil, position: str, + size: int) -> Point: + point = Point(0, 0) + padding = int(size / 2) + if "top" in position: + point.y = size + padding + elif "bottom" in position: + point.y = image_pil.height - size - padding + else: + raise ValueError("Cannot position the compass") + + if "left" in position: + point.x = size + padding + elif "right" in position: + point.x = image_pil.width - size - padding + else: + raise ValueError("Cannot position the compass") + + return point + + def draw_arrow(self, draw: Draw, font, center: Point, angle: int, size: int, cardinal=None): + + if cardinal: + card_point = center.add_polar(size * 1.2, 90 + angle) + draw.text(card_point.coord(), cardinal, font=font, anchor="mm", + fill=self.black_pen, + stroke_width=1, stroke_fill=self.gray_pen) + else: + size = size * 0.75 + point1 = center.add_polar(size / 4, 45 + angle) point2 = center.add_polar(size, 90 + angle) point3 = center.add_polar(size / 4, 135 + angle) - color1, color2 = BLACK, WHITE - if cardinal: - card_point = center.add_polar(size + size / 8, 90 + angle) - draw.text(card_point.coord(), cardinal, font=font, fill=BLACK, - anchor="mm") - draw.polygon([center.coord(), point1.coord(), point2.coord()], - fill=color1, width=2) - - draw.polygon([center.coord(), point2.coord(), point3.coord()], - fill=color2, width=2) + draw.polygon( + [center.coord(), point1.coord(), point2.coord()], self.black_pen) + draw.polygon( + [center.coord(), point2.coord(), point3.coord()], self.white_pen) def compass_overlay(self, image: torch.Tensor, font: str, cardinals: str, - size: int, rotation: int): + position: str, size: int, + seed: int, + rotation: int, random_rotation: bool): image_pil = tensor_to_pil(image) - # Create Gri - draw = ImageDraw.Draw(image_pil) + draw = Draw(image_pil, "RGBA") + if random_rotation: + random.seed(seed) + rotation = random.randint(0, 360) + font = ImageFont.truetype(font, int(size / 4)) - point = Point(image_pil.width / 2, image_pil.height / 2) + center = self.get_center_position(image_pil, position, size) draw.ellipse( - [(point - size * 0.8).coord(), (point + size * 0.8).coord()], - outline=BLACK, width=2) + [(center - size * 0.8).coord(), (center + size * 0.8).coord()], + outline=self.black_pen, width=min(6, round(size / 12))) draw.ellipse( - [(point - size * 0.7).coord(), (point + size * 0.7).coord()], - outline=BLACK, width=2) + [(center - size * 0.7).coord(), (center + size * 0.7).coord()], + outline=self.black_pen, width=min(4, round(size / 16))) for i in range(0, 360, 90): - self.draw_arrow(draw, font, point, i + 45 + rotation, size) + self.draw_arrow(draw, font, center, i + 45 + rotation, size) for i, card in zip(range(0, 360, 90), cardinals): - self.draw_arrow(draw, font, point, i + rotation, size, + self.draw_arrow(draw, font, center, i + rotation, size, cardinal=card) image_tensor_out = pil_to_tensor(image_pil) return (image_tensor_out,) diff --git a/map_node.py b/map_node.py index 01bd31d..8359c82 100644 --- a/map_node.py +++ b/map_node.py @@ -2,7 +2,8 @@ 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, seed, choice +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 @@ -11,11 +12,10 @@ from .path import Path @dataclass class BattlemapMapGenerator: @classmethod - def INPUT_TYPES(cls): + def INPUT_TYPES(cls) -> dict: return {"required": { - "node_seed": ( - "INT", {"default": 0, "min": -1, "max": 0xffffffffffffffff, - "step": 1, "label": "seed", "forceInput": True}), + "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}), @@ -25,24 +25,24 @@ class BattlemapMapGenerator: @classmethod @property - def RETURN_TYPES(cls): + def RETURN_TYPES(cls) -> tuple: return ("IMAGE", "INT", "INT", "INT", "INT", "INT") @classmethod @property - def RETURN_NAMES(cls): + 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, height, bg_color): + 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, draw, bg_color): + 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): @@ -55,9 +55,9 @@ class BattlemapMapGenerator: draw.ellipse([point1.coord(), point2.coord()], fill=color) generate_noise(image, Point(0, 0), Point(image.width, image.height)) - def _map_generator(self, node_seed, grid_width, grid_height, - grid_side, bg_color): - seed(node_seed) + 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) @@ -65,10 +65,10 @@ class BattlemapMapGenerator: return (image_pil, width, height, grid_width, grid_height, grid_side) - def map_generator(self, node_seed, grid_width, grid_height, - grid_side, bg_color): + 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( - node_seed, grid_width, grid_height, grid_side, bg_color) + seed, grid_width, grid_height, grid_side, bg_color) return (pil_to_tensor(image_pil), width, height, grid_width, grid_height, grid_side) @@ -98,27 +98,27 @@ class BattlemapMapGeneratorOutdoors(BattlemapMapGenerator): @classmethod @property - def RETURN_TYPES(cls): + 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): + 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, draw): - for i in range(100, 500): + 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): + def start_point(self, image: Image) -> tuple[Point, int]: match randint(0, 3): case 0: point = Point(randint(0, image.width), -10) @@ -131,9 +131,10 @@ class BattlemapMapGeneratorOutdoors(BattlemapMapGenerator): angle = point.angle_between(Point(image.width / 2, image.height / 2)) return point, angle + randint(-45, 45) - def generate_path(self, image, paths: list[Path], start_point, start_angle, - width=20, - depth=0, depth_max=5): + 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) @@ -150,13 +151,11 @@ class BattlemapMapGeneratorOutdoors(BattlemapMapGenerator): self.generate_path(image, paths, point1, new_angle, width=math.ceil(width * 2 / 3), depth=depth + 1, depth_max=depth_max) - case _: - pass point1, angle, length = (point2, angle + randint(-20, 20), randint(20, 200)) paths.append(path) - def draw_path(self, draw, color, paths: list): + 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) @@ -168,19 +167,20 @@ class BattlemapMapGeneratorOutdoors(BattlemapMapGenerator): 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, draw): + 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, draw): + 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, draw, size_color): + 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: @@ -194,7 +194,8 @@ class BattlemapMapGeneratorOutdoors(BattlemapMapGenerator): fill=color, width=witdh) generate_noise(image, center - size, center + size, 10) - def generate_polygon(self, image, draw, size_color, nb_point): + 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: @@ -206,7 +207,8 @@ class BattlemapMapGeneratorOutdoors(BattlemapMapGenerator): draw.polygon(point_list, fill=color, outline="black", width=2) generate_noise(image, center - size, center + size, 10) - def generate_ellipses(self, image, draw, size_color): + 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 @@ -220,26 +222,28 @@ class BattlemapMapGeneratorOutdoors(BattlemapMapGenerator): fill=color, outline="black", width=2) generate_noise(image, center - size_max, center + size_max, 10) - def generate_rocks(self, image, draw): - for i in range(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, draw): - for i in range(100): + 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, node_seed, grid_width, grid_height, - grid_side, bg_color, river, road, rocks, trees, - positive, negative): - seed(node_seed) + 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(node_seed, 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)