-
Notifications
You must be signed in to change notification settings - Fork 4
/
Copy pathtextMaskOverlay.py
208 lines (177 loc) · 7.97 KB
/
textMaskOverlay.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
"""
@author: Manny Gonzalez
@title: 🐯 YFG Comical Nodes
@nickname: 🐯 YFG Comical Nodes
@description: Utility custom nodes for special effects, image manipulation and quality of life tools.
"""
import cv2
import numpy as np
import math
import torch
from PIL import Image, ImageDraw, ImageFont
import os
import platform
class TextMaskOverlay:
@classmethod
def INPUT_TYPES(cls):
font_sources = ["System", "Node Folder", "Directory"]
# Default user directory to fonts
user_directory = os.path.expandvars("%LocalAppData%\\Microsoft\\Windows\\Fonts")
# Define font lists for different sources
system_fonts = cls.get_font_list("System")
node_fonts = cls.get_font_list("Node Folder")
directory_fonts = cls.get_font_list("Directory", user_directory)
if not system_fonts:
print("Error: No system fonts found.")
system_fonts = ["No system fonts available"]
if not node_fonts:
print("Error: No fonts found in the YFG node folder.")
node_fonts = ["No fonts available in YFG folder"]
if not directory_fonts:
print(f"Error: No fonts found in directory '{user_directory}' or the directory does not exist.")
directory_fonts = ["No fonts available in specified directory"]
return {
"required": {
"font_source": (font_sources,),
"system_font": (system_fonts, {"default": "impact.ttf"}),
"node_font": (node_fonts, {"default": "malgunbd.ttf"}),
"usr_font_dir": ("STRING", {"default": user_directory, "multiline": False}),
"directory_font": (directory_fonts, {"default": directory_fonts[0] if directory_fonts else ""}),
"text": ("STRING", {"default": "TEXT", "multiline": False}),
'mask': ("MASK",),
},
}
CATEGORY = "🐯 YFG"
RETURN_TYPES = ("MASK", "IMAGE")
RETURN_NAMES = ("mask", "image")
FUNCTION = "process_mask"
OUTPUT_NODE = True
def process_mask(self, mask, font_source, system_font, node_font, usr_font_dir, directory_font, text):
# Process the mask with text overlay
processed_mask = self.apply_text_mask(mask, font_source, system_font, node_font, usr_font_dir, directory_font, text)
# Convert the modified mask to an image
image = self.mask_to_image(processed_mask)
return processed_mask, image
def apply_text_mask(self, mask, font_source, system_font, node_font, usr_font_dir, directory_font, text):
if len(mask.shape) > 2:
mask = mask[0]
mask = (mask.numpy() * 255).astype('uint8')
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if contours:
largest_contour = max(contours, key=cv2.contourArea)
x, y, w, h = cv2.boundingRect(largest_contour)
tmp = np.zeros_like(mask)
rows = min(max(round(h / w), 1), 3)
row_num = math.ceil(len(text) / rows)
font_size = math.floor(w / row_num)
dx = round(w - font_size * row_num, 1)
dy = round((h - font_size * rows) / 2 / rows, 1)
print(rows, row_num, font_size)
# Select font based on source
if font_source == "System":
font = self.get_font("System", "", system_font, font_size)
elif font_source == "Node Folder":
font = self.get_font("Node Folder", "", node_font, font_size)
elif font_source == "Directory":
font = self.get_font("Directory", usr_font_dir, directory_font, font_size)
image2 = Image.fromarray(tmp)
draw = ImageDraw.Draw(image2)
for row in range(rows):
draw.text((x + dx, y + (font_size + dy * 2) * row + dy - font_size / 5),
text[row_num * row:row_num * (row + 1)], font=font, fill='white')
else:
image2 = Image.fromarray(np.zeros_like(mask))
return torch.from_numpy((np.array(image2).astype('float32') / 255))
def mask_to_image(self, mask):
"""
Converts a modified mask to an RGB image format.
"""
# Using the provided method for mask conversion
result = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
return result
def get_font(self, font_source, usr_font_dir, font_name, size):
"""
Loads a font based on the selected source and specified filename.
"""
font_path = None
if font_source == "System":
font_path = self.find_system_font(font_name)
elif font_source == "Node Folder":
font_path = self.find_node_font(font_name)
elif font_source == "Directory":
font_path = os.path.join(usr_font_dir, font_name)
if not os.path.exists(usr_font_dir):
print(f"Error: The directory '{usr_font_dir}' does not exist.")
return self.get_default_system_font(size)
if font_path and os.path.isfile(font_path):
try:
font = ImageFont.truetype(font_path, size)
except IOError:
print(f"Font '{font_name}' at '{font_path}' not found. Using default system font 'impact.ttf'.")
font = self.get_default_system_font(size)
else:
print(f"Font path '{font_path}' not valid or does not exist. Using default system font 'impact.ttf'.")
font = self.get_default_system_font(size)
return font
def find_system_font(self, font_name):
"""
Finds the specified font in the system's font directory.
"""
font_dir = self.get_system_font_dir()
return os.path.join(font_dir, font_name) if font_dir else None
def find_node_font(self, font_name):
"""
Finds the specified font in the node's predefined font directory.
"""
yfg_font_dir = self.get_node_font_dir()
return os.path.join(yfg_font_dir, font_name)
def get_default_system_font(self, size):
"""
Loads a common system font if available. Defaults to a basic bitmap font if not found.
"""
try:
font = ImageFont.truetype("impact.ttf", size)
except IOError:
print("System font 'impact.ttf' not found. Using default bitmap font.")
font = ImageFont.load_default()
return font
@classmethod
def get_font_list(cls, source, directory=""):
"""
Gets the list of fonts based on the selected source.
"""
font_list = []
if source == "System":
font_dir = cls.get_system_font_dir()
elif source == "Node Folder":
font_dir = cls.get_node_font_dir()
elif source == "Directory":
font_dir = directory
if font_dir and os.path.exists(font_dir):
font_list = [f for f in os.listdir(font_dir) if os.path.isfile(os.path.join(font_dir, f)) and f.lower().endswith(".ttf")]
if not font_list:
print(f"Error: No fonts found in '{font_dir}'.")
else:
print(f"Error: The directory '{font_dir}' does not exist.")
return font_list
@classmethod
def get_system_font_dir(cls):
"""
Returns the system font directory based on the OS.
"""
if platform.system() == "Windows":
system_root = os.environ.get("SystemRoot")
return os.path.join(system_root, "Fonts") if system_root else None
elif platform.system() == "Linux":
return "/usr/share/fonts/truetype"
elif platform.system() == "Darwin":
return "/System/Library/Fonts"
else:
return None
@classmethod
def get_node_font_dir(cls):
"""
Returns the node's font directory relative to the current file.
"""
node_dir = os.path.dirname(os.path.abspath(__file__))
return os.path.join(node_dir, "fonts")