233 lines
9.6 KiB
Python
233 lines
9.6 KiB
Python
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)
|