133 lines
5.1 KiB
Python
133 lines
5.1 KiB
Python
import torch
|
|
from aggdraw import Draw, Brush, Pen
|
|
import math
|
|
from .utils import tensor_to_pil, pil_to_tensor
|
|
from .point import Point
|
|
from dataclasses import dataclass
|
|
|
|
|
|
@dataclass
|
|
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
|
|
"""
|
|
|
|
edge_length: int = 0
|
|
center: Point = None
|
|
|
|
@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})
|
|
},
|
|
"optional": {
|
|
"orig_grid_width": ("INT",
|
|
{"default": None, "min": 0, "max": 128,
|
|
"forceInput": True}),
|
|
"orig_grid_height": ("INT",
|
|
{"default": None, "min": 00, "max": 128,
|
|
"forceInput": True}),
|
|
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "grid_overlay"
|
|
CATEGORY = "Battlemaps"
|
|
|
|
def square_grid(self, image, draw, center,
|
|
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, hexagon_generator: BaseHexagonGenerator,
|
|
image, draw, center, pen, grid_side):
|
|
hexagon_generator = hexagon_generator(grid_side, center)
|
|
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,
|
|
orig_grid_width=None, orig_grid_height=None):
|
|
image_pil = tensor_to_pil(image)
|
|
# Create Grid
|
|
|
|
draw = Draw(image_pil)
|
|
pen = Pen((red, green, blue, alpha), line_width)
|
|
center = Point(int(image_pil.width / 2), int(image_pil.height / 2))
|
|
if (orig_grid_width and orig_grid_height):
|
|
width_rate = (image_pil.width / orig_grid_width)
|
|
height_rate = (image_pil.height / orig_grid_height)
|
|
if 0.99 <= (width_rate / height_rate) <= 1.01:
|
|
grid_side = round(image_pil.width / orig_grid_width)
|
|
if grid_type == "square":
|
|
self.square_grid(image_pil, draw, center, pen,
|
|
grid_side)
|
|
elif grid_type == "vertical hexagon":
|
|
self.hexagon_grid(VerticalHexagonGenerator,
|
|
image_pil, draw, center, pen, grid_side)
|
|
elif grid_type == "horizontal hexagon":
|
|
self.hexagon_grid(HorizontalHexagonGenerator,
|
|
image_pil, draw, center, pen, grid_side)
|
|
else:
|
|
raise Exception
|
|
draw.flush()
|
|
image_tensor_out = pil_to_tensor(image_pil)
|
|
return (image_tensor_out,)
|