65 lines
2.5 KiB
Python
65 lines
2.5 KiB
Python
import torch
|
|
from PIL import ImageDraw, ImageFont
|
|
from .utils import tensor_to_pil, pil_to_tensor
|
|
from .point import Point
|
|
|
|
BLACK = (0, 0, 0, 128)
|
|
WHITE = (255, 255, 255, 128)
|
|
|
|
|
|
class CompassGrid:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"font": ("STRING", {
|
|
"default": "/System/Library/Fonts/Supplemental/Arial.ttf"}),
|
|
"cardinals": (["NWSE", "NOSE"], {"default": "NWSE"}),
|
|
"size": ("INT", {"default": 64, "min": 10, "max": 2048}),
|
|
"rotation": ("INT", {"default": 0, "min": 0, "max": 360}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "compass_overlay"
|
|
CATEGORY = "Battlemaps"
|
|
|
|
def draw_arrow(self, draw: ImageDraw.Draw, font, center: Point,
|
|
angle: int, size: int, cardinal=None):
|
|
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)
|
|
|
|
def compass_overlay(self, image: torch.Tensor, font: str, cardinals: str,
|
|
size: int, rotation: int):
|
|
image_pil = tensor_to_pil(image)
|
|
# Create Gri
|
|
draw = ImageDraw.Draw(image_pil)
|
|
font = ImageFont.truetype(font, int(size / 4))
|
|
point = Point(image_pil.width / 2, image_pil.height / 2)
|
|
draw.ellipse(
|
|
[(point - size * 0.8).coord(), (point + size * 0.8).coord()],
|
|
outline=BLACK, width=2)
|
|
draw.ellipse(
|
|
[(point - size * 0.7).coord(), (point + size * 0.7).coord()],
|
|
outline=BLACK, width=2)
|
|
for i in range(0, 360, 90):
|
|
self.draw_arrow(draw, font, point, i + 45 + rotation, size)
|
|
for i, card in zip(range(0, 360, 90), cardinals):
|
|
self.draw_arrow(draw, font, point, i + rotation, size,
|
|
cardinal=card)
|
|
image_tensor_out = pil_to_tensor(image_pil)
|
|
return (image_tensor_out,)
|