783 lines
32 KiB
Python
783 lines
32 KiB
Python
'''
|
|
AnyText2: Visual Text Generation and Editing With Customizable Attributes
|
|
Paper: https://arxiv.org/abs/2411.15245
|
|
Code: https://github.com/tyxsspa/AnyText2
|
|
Copyright (c) Alibaba, Inc. and its affiliates.
|
|
'''
|
|
import os
|
|
import numpy as np
|
|
import cv2
|
|
import random
|
|
import math
|
|
import time
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
from torch.utils.data import Dataset, DataLoader
|
|
from .AnyText2_dataset_util import load, show_bbox_on_image
|
|
from opencc import OpenCC
|
|
# from ....Log.logger import logger
|
|
from .... import comfy_temp_dir
|
|
import shutil
|
|
anytext_tmp_dir = os.path.join(comfy_temp_dir, "AnyText")
|
|
if os.path.exists(anytext_tmp_dir):
|
|
shutil.rmtree(anytext_tmp_dir)
|
|
os.makedirs(anytext_tmp_dir, exist_ok=True)
|
|
SHOW_GLYPH = False
|
|
|
|
cc = OpenCC('t2s')
|
|
|
|
phrase_list = [
|
|
'Text says ',
|
|
'Image with words ',
|
|
'The picture reads ',
|
|
'Captions are ',
|
|
'Texts are ',
|
|
'Text: '
|
|
' '
|
|
]
|
|
|
|
default_color = [500, 500, 500]
|
|
|
|
fix_masked_img_bug = True
|
|
|
|
|
|
def tra_chinese(text):
|
|
simplified_text = cc.convert(text)
|
|
return text != simplified_text
|
|
|
|
|
|
def copy_and_rename_dict_keys(original_dict, key_mapping):
|
|
new_dict = {}
|
|
for old_key, new_key in key_mapping.items():
|
|
if old_key in original_dict:
|
|
new_dict[new_key] = original_dict[old_key]
|
|
return new_dict
|
|
|
|
|
|
key_mapping = {
|
|
'en': 'Latin',
|
|
'ch_sim_char': 'Chinese_sim',
|
|
'ch_tra_char': 'Chinese_tra',
|
|
'hi': 'Hindi',
|
|
'ar': 'Arabic',
|
|
'ja': 'Japanese',
|
|
'ko': 'Korean',
|
|
'bn': 'Bangla',
|
|
}
|
|
|
|
|
|
def random_rotate(image, angle_range):
|
|
angle = random.uniform(angle_range[0], angle_range[1])
|
|
h, w = image.shape[:2]
|
|
M = cv2.getRotationMatrix2D((w/2, h/2), angle, 1)
|
|
rotated = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
|
|
return rotated
|
|
|
|
|
|
def random_translate(image, translate_range):
|
|
tx = random.uniform(translate_range[0], translate_range[1])
|
|
ty = random.uniform(translate_range[0], translate_range[1])
|
|
M = np.float32([[1, 0, tx], [0, 1, ty]])
|
|
h, w = image.shape[:2]
|
|
translated = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
|
|
return translated
|
|
|
|
|
|
def random_scale(image, scale_range):
|
|
scale = random.uniform(scale_range[0], scale_range[1])
|
|
h, w = image.shape[:2]
|
|
scaled = cv2.resize(image, (int(w*scale), int(h*scale)), interpolation=cv2.INTER_LINEAR)
|
|
if scale >= 1:
|
|
scaled = scaled[(scaled.shape[0]-h)//2: (scaled.shape[0]+h)//2, (scaled.shape[1]-w)//2: (scaled.shape[1]+w)//2]
|
|
else:
|
|
pad_h = (h - scaled.shape[0]) // 2
|
|
pad_w = (w - scaled.shape[1]) // 2
|
|
scaled = cv2.copyMakeBorder(scaled, pad_h, h - scaled.shape[0] - pad_h, pad_w, w - scaled.shape[1] - pad_w, cv2.BORDER_REPLICATE)
|
|
return scaled
|
|
|
|
|
|
def random_augment(image, rot=(-10, 10), trans=(-5, 5), scale=(0.9, 1.1)):
|
|
image = random_rotate(image, rot)
|
|
image = random_translate(image, trans)
|
|
image = random_scale(image, scale)
|
|
return image
|
|
|
|
|
|
def insert_spaces(text, num_spaces):
|
|
return (' ' * num_spaces).join(text)
|
|
|
|
|
|
def draw_glyph(font, text, color, show_glyph: bool=False):
|
|
if isinstance(font, str):
|
|
font = ImageFont.truetype(font, size=60)
|
|
g_size = 50
|
|
W, H = (512, 80)
|
|
new_font = font.font_variant(size=g_size)
|
|
img = Image.new(mode='1', size=(W, H), color=0)
|
|
draw = ImageDraw.Draw(img)
|
|
left, top, right, bottom = new_font.getbbox(text)
|
|
text_width = max(right-left, 5)
|
|
text_height = max(bottom - top, 5)
|
|
ratio = min(W*0.9/text_width, H*0.9/text_height)
|
|
new_font = font.font_variant(size=int(g_size*ratio))
|
|
|
|
try:
|
|
text_width, text_height = new_font.getsize(text)
|
|
offset_x, offset_y = new_font.getoffset(text)
|
|
except Exception as e:
|
|
#增加使用pillow>9.5
|
|
# print(f'\033[93mWarning: pillow>9.5, change getsize to get bbox, maybe result to lower quality,\033[0m')
|
|
x0, y0, x1, y1 = new_font.getbbox(text)
|
|
text_width, text_height = x1-x0, y1-y0
|
|
offset_x, offset_y = text_width, text_height
|
|
|
|
x = (img.width - text_width) // 2
|
|
y = (img.height - text_height) // 2 - offset_y//2
|
|
|
|
if color.mean() < 0:
|
|
color = np.array([255, 255, 255])
|
|
color = np.clip(color, 10, 255) # RGB >= 10
|
|
|
|
# draw.text((x, y), text, font=new_font, fill='white')
|
|
draw.text((x, y), text, font=new_font, fill='tomato')
|
|
|
|
img = np.expand_dims(np.array(img), axis=2).astype(np.float64)
|
|
# if SHOW_GLYPH:
|
|
if show_glyph:
|
|
i = 0
|
|
# file_path = f'tmp_{i}.jpg'
|
|
# while os.path.exists(file_path):
|
|
# i += 1
|
|
# file_path = f'tmp_{i}.jpg'
|
|
file_path = os.path.join(anytext_tmp_dir, f'glyph_{i}.png')
|
|
while os.path.exists(file_path):
|
|
i += 1
|
|
file_path = os.path.join(anytext_tmp_dir, f'glyph_{i}.png')
|
|
cv2.imwrite(file_path, img*255)
|
|
return img
|
|
|
|
|
|
def draw_glyph2(font, text, polygon, color, vertAng=10, scale=1, width=512, height=512, add_space=True, show_glyph=False):
|
|
def initialize_img(width, height, scale):
|
|
img = np.zeros((height * scale, width * scale, 3), np.uint8)
|
|
return Image.fromarray(img)
|
|
|
|
def prepare_image(img):
|
|
return np.array(img.convert('RGB')).astype(np.float64) / 255.0
|
|
|
|
try:
|
|
if color.mean() < 0:
|
|
color = np.array([255, 255, 255])
|
|
color = np.clip(color, 10, 255) # RGB >= 10
|
|
if isinstance(font, str):
|
|
if os.path.exists(font):
|
|
font = ImageFont.truetype(font, size=60)
|
|
else:
|
|
img = initialize_img(width, height, scale)
|
|
return prepare_image(img)
|
|
enlarge_polygon = np.array(polygon) * scale
|
|
rect = cv2.minAreaRect(enlarge_polygon)
|
|
box = cv2.boxPoints(rect)
|
|
box = np.int0(box)
|
|
w, h = rect[1]
|
|
angle = rect[2]
|
|
|
|
if angle < -45:
|
|
angle += 90
|
|
angle = -angle
|
|
if w < h:
|
|
angle += 90
|
|
|
|
vert = False
|
|
if (abs(angle) % 90 < vertAng or abs(90 - abs(angle) % 90) % 90 < vertAng):
|
|
_w = max(box[:, 0]) - min(box[:, 0])
|
|
_h = max(box[:, 1]) - min(box[:, 1])
|
|
if _h >= _w:
|
|
vert = True
|
|
angle = 0
|
|
|
|
img = initialize_img(width, height, scale)
|
|
image4ratio = Image.new("RGB", img.size, "white")
|
|
draw = ImageDraw.Draw(image4ratio)
|
|
min_dim = min(w, h)
|
|
max_dim = max(w, h)
|
|
|
|
# Binary search for optimal font size
|
|
def adjust_font_size(min_size, max_size, text):
|
|
while min_size < max_size:
|
|
mid_size = (min_size + max_size) // 2
|
|
new_font = font.font_variant(size=int(mid_size))
|
|
bbox = draw.textbbox((0, 0), text=text, font=new_font)
|
|
text_w = bbox[2] - bbox[0]
|
|
text_h = bbox[3] - bbox[1]
|
|
if text_w <= max_dim and text_h <= min_dim:
|
|
min_size = mid_size + 1
|
|
else:
|
|
max_size = mid_size
|
|
return max_size - 1
|
|
|
|
optimal_font_size = adjust_font_size(1, min_dim, text)
|
|
new_font = font.font_variant(size=int(optimal_font_size))
|
|
|
|
extra_space = 0
|
|
if add_space:
|
|
if vert:
|
|
# Calculate total height with added space
|
|
total_height = sum(draw.textbbox((0, 0), text=char, font=new_font)[3] -
|
|
draw.textbbox((0, 0), text=char, font=new_font)[1]
|
|
for char in text)
|
|
if total_height < max_dim and len(text) > 1:
|
|
extra_space = (max_dim - total_height) // (len(text) - 1)
|
|
else:
|
|
# Handle horizontal text space addition
|
|
for i in range(1, 100):
|
|
text_space = insert_spaces(text, i)
|
|
bbox2 = draw.textbbox((0, 0), text=text_space, font=new_font)
|
|
text_w, text_h = bbox2[2] - bbox2[0], bbox2[3] - bbox2[1]
|
|
if text_w > max_dim or text_h > min_dim:
|
|
text = insert_spaces(text, i - 1)
|
|
break
|
|
|
|
left, top, right, bottom = draw.textbbox((0, 0), text=text, font=new_font)
|
|
text_width = right - left
|
|
text_height = bottom - top
|
|
|
|
layer = Image.new('RGBA', img.size, (0, 0, 0, 0))
|
|
draw = ImageDraw.Draw(layer)
|
|
|
|
if not vert:
|
|
text_y_center = rect[0][1] - (text_height / 2)
|
|
draw.text((rect[0][0] - text_width / 2, text_y_center - top), text, font=new_font, fill=tuple(color)+(255,))
|
|
else:
|
|
x_s = min(box[:, 0]) + _w // 2 - text_height // 2
|
|
y_s = min(box[:, 1])
|
|
for c in text:
|
|
draw.text((x_s, y_s), c, font=new_font, fill=tuple(color)+(255,))
|
|
_, _t, _, _b = draw.textbbox((0, 0), text=c, font=new_font)
|
|
char_height = _b - _t
|
|
y_s += char_height + extra_space
|
|
|
|
rotated_layer = layer.rotate(angle, expand=1, center=(rect[0][0], rect[0][1]))
|
|
x_offset = int((img.width - rotated_layer.width) / 2)
|
|
y_offset = int((img.height - rotated_layer.height) / 2)
|
|
img.paste(rotated_layer, (x_offset, y_offset), rotated_layer)
|
|
|
|
if show_glyph:
|
|
i = 0
|
|
file_path = os.path.join(anytext_tmp_dir, f'glyph2_{i}.png')
|
|
while os.path.exists(file_path):
|
|
i += 1
|
|
file_path = os.path.join(anytext_tmp_dir, f'glyph2_{i}.png')
|
|
img.save(file_path)
|
|
|
|
return prepare_image(img)
|
|
|
|
except Exception as e:
|
|
print(f"An error occurred in draw_glyph2: {e}")
|
|
img = initialize_img(width, height, scale)
|
|
return prepare_image(img)
|
|
|
|
|
|
'''
|
|
target_img: (-1,1), hwc
|
|
return font_hint: (0,1), hw1
|
|
'''
|
|
|
|
|
|
def draw_font_hint(target_img, polygon, target_area_range=[1.0, 1.0], prob=1.0, randaug=False):
|
|
height, width, _ = target_img.shape
|
|
img = np.zeros((height, width), dtype=np.uint8)
|
|
|
|
if random.random() < (1 - prob): # Empty font hint
|
|
return img[..., None] / 255.0, img[..., None] / 255.0
|
|
|
|
polygon[:, 0] = np.clip(polygon[:, 0], 0, width - 1)
|
|
polygon[:, 1] = np.clip(polygon[:, 1], 0, height - 1)
|
|
pts = polygon.reshape((-1, 1, 2)).astype(np.int32)
|
|
cv2.fillPoly(img, [pts], color=255)
|
|
rect = cv2.minAreaRect(pts)
|
|
center, size, angle = rect
|
|
rect_width, rect_height = size
|
|
x, y, w, h = cv2.boundingRect(np.clip(polygon, 0, None))
|
|
target_img_scaled = (target_img + 1.0) / 2.0
|
|
cropped_ori_img = target_img_scaled[y:y+h, x:x+w]
|
|
if randaug:
|
|
augmented_cropped = random_augment(cropped_ori_img, rot=(-10, 10), trans=(-10, 10), scale=(0.9, 1.1))
|
|
else:
|
|
augmented_cropped = cropped_ori_img
|
|
augmented_cropped_gray = cv2.cvtColor((augmented_cropped * 255).astype(np.uint8), cv2.COLOR_RGB2GRAY)
|
|
thresholded = cv2.adaptiveThreshold(augmented_cropped_gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)
|
|
|
|
thresholded_resized = np.zeros_like(img.squeeze())
|
|
thresholded_resized[y:y+h, x:x+w] = (1 - thresholded / 255.0)
|
|
|
|
# gen a random mask
|
|
area_ratio = random.uniform(target_area_range[0], target_area_range[1])
|
|
long_side, short_side = max(rect_width, rect_height), min(rect_width, rect_height)
|
|
long_axis_mask_length = long_side * (1 - area_ratio)
|
|
angle_rad = np.radians(angle)
|
|
rect_vector = np.array([np.cos(angle_rad), np.sin(angle_rad)])
|
|
if rect_width < rect_height:
|
|
rect_vector = np.array([-rect_vector[1], rect_vector[0]])
|
|
start_offset = random.uniform(0, long_side - long_axis_mask_length)
|
|
start_point = center - rect_vector * (long_side / 2 - start_offset)
|
|
mask_center = start_point + rect_vector * (long_axis_mask_length / 2)
|
|
mask_vector = rect_vector * (long_axis_mask_length / 2)
|
|
short_axis_vector = np.array([-rect_vector[1], rect_vector[0]]) * (short_side / 2)
|
|
mask_corners = np.array([
|
|
mask_center - mask_vector - short_axis_vector,
|
|
mask_center + mask_vector - short_axis_vector,
|
|
mask_center + mask_vector + short_axis_vector,
|
|
mask_center - mask_vector + short_axis_vector
|
|
], dtype=np.int32)
|
|
cv2.fillPoly(img, [mask_corners], color=0)
|
|
img = img[..., None] / 255.0
|
|
|
|
# Compute font hint
|
|
font_hint = img.squeeze() * thresholded_resized
|
|
return font_hint[..., None], img
|
|
|
|
|
|
def get_text_caption(n_line, ori_caption, place_holder='*'):
|
|
return random.choice(phrase_list) + ' , '.join([f'{place_holder}']*n_line) + ' . '
|
|
|
|
|
|
def generate_random_rectangles(w, h, box_num):
|
|
rectangles = []
|
|
for i in range(box_num):
|
|
x = random.randint(0, w)
|
|
y = random.randint(0, h)
|
|
w = random.randint(16, 256)
|
|
h = random.randint(16, 96)
|
|
angle = random.randint(-45, 45)
|
|
p1 = (x, y)
|
|
p2 = (x + w, y)
|
|
p3 = (x + w, y + h)
|
|
p4 = (x, y + h)
|
|
center = ((x + x + w) / 2, (y + y + h) / 2)
|
|
p1 = rotate_point(p1, center, angle)
|
|
p2 = rotate_point(p2, center, angle)
|
|
p3 = rotate_point(p3, center, angle)
|
|
p4 = rotate_point(p4, center, angle)
|
|
rectangles.append((p1, p2, p3, p4))
|
|
return rectangles
|
|
|
|
|
|
def rotate_point(point, center, angle):
|
|
# rotation
|
|
angle = math.radians(angle)
|
|
x = point[0] - center[0]
|
|
y = point[1] - center[1]
|
|
x1 = x * math.cos(angle) - y * math.sin(angle)
|
|
y1 = x * math.sin(angle) + y * math.cos(angle)
|
|
x1 += center[0]
|
|
y1 += center[1]
|
|
return int(x1), int(y1)
|
|
|
|
|
|
def truncate_string(input_string, max_words=45):
|
|
input_string = input_string[:500] # be careful for repeat CN words, may cause OOM
|
|
words = input_string.split()
|
|
if len(words) <= max_words:
|
|
return input_string
|
|
# Split into sentences
|
|
sentences = input_string.split('.')
|
|
first_sentence = sentences[0].strip()
|
|
remaining_sentences = sentences[1:]
|
|
# Start with the first sentence
|
|
truncated_string = first_sentence
|
|
words_count = len(first_sentence.split())
|
|
if words_count >= max_words:
|
|
return ' '.join(first_sentence.split()[:max_words])
|
|
# Randomly add sentences until we reach or exceed the word limit
|
|
while words_count < max_words and remaining_sentences:
|
|
random_index = random.randint(0, len(remaining_sentences) - 1)
|
|
sentence = remaining_sentences.pop(random_index).strip()
|
|
sentence_words = sentence.split()
|
|
if words_count + len(sentence_words) > max_words:
|
|
# If adding the sentence exceeds word limit, truncate the sentence
|
|
sentence_words = sentence_words[:max_words - words_count] + ['.']
|
|
if sentence_words:
|
|
truncated_string += '. ' + ' '.join(sentence_words)
|
|
words_count += len(sentence_words)
|
|
return truncated_string
|
|
|
|
|
|
class T3DataSet(Dataset):
|
|
def __init__(
|
|
self,
|
|
json_path,
|
|
max_lines=5,
|
|
max_chars=20,
|
|
place_holder='*',
|
|
font_path='./font/Arial_Unicode.ttf',
|
|
mask_pos_prob=1.0,
|
|
mask_img_prob=0.5,
|
|
for_show=False,
|
|
using_dlc=False,
|
|
glyph_scale=1,
|
|
percent=1.0,
|
|
debug=False,
|
|
wm_thresh=1.0,
|
|
render_glyph=True,
|
|
trunc_cap=128, # caption truncation
|
|
rand_font=False,
|
|
lang_font_path='./font/lang_font_dict.npy',
|
|
font_hint_prob=0, # set 0 to disable font_hint
|
|
color_prob=1.0, # set 0 to disable color
|
|
font_hint_area=[1.0, 1.0], # reserved area on each font_hint line
|
|
font_hint_randaug=False,
|
|
cap_watermark=True,
|
|
img_wh=512,
|
|
):
|
|
assert isinstance(json_path, (str, list))
|
|
if isinstance(json_path, str):
|
|
json_path = [json_path]
|
|
data_list = []
|
|
self.using_dlc = using_dlc
|
|
self.max_lines = max_lines
|
|
self.max_chars = max_chars
|
|
self.place_holder = place_holder
|
|
self.font = ImageFont.truetype(font_path, size=60)
|
|
self.mask_pos_prob = mask_pos_prob
|
|
self.mask_img_prob = mask_img_prob
|
|
self.for_show = for_show
|
|
self.glyph_scale = glyph_scale
|
|
self.wm_thresh = wm_thresh
|
|
self.render_glyph = render_glyph
|
|
self.trunc_cap = trunc_cap
|
|
self.rand_font = rand_font
|
|
self.font_hint_prob = font_hint_prob
|
|
self.color_prob = color_prob
|
|
self.font_hint_area = font_hint_area
|
|
self.font_hint_randaug = font_hint_randaug
|
|
self.cap_watermark = cap_watermark
|
|
self.img_wh = img_wh
|
|
|
|
if self.rand_font:
|
|
lang_font_dict = np.load(lang_font_path, allow_pickle=True)[()]
|
|
self.lang_font = copy_and_rename_dict_keys(lang_font_dict, key_mapping)
|
|
for lang in self.lang_font:
|
|
self.lang_font[lang] = [ImageFont.truetype(p, size=60) for p in self.lang_font[lang]['fonts']]
|
|
print('rand_font=True, all fonts are loaded!')
|
|
|
|
for jp in json_path:
|
|
data_list += self.load_data(jp, percent)
|
|
self.data_list = data_list
|
|
print(f'All dataset loaded, imgs={len(self.data_list)}')
|
|
self.debug = debug
|
|
if self.debug:
|
|
self.tmp_items = [i for i in range(100)]
|
|
|
|
def load_data(self, json_path, percent):
|
|
tic = time.time()
|
|
content = load(json_path)
|
|
d = []
|
|
count = 0
|
|
wm_skip = 0
|
|
max_img = len(content['data_list']) * percent
|
|
for gt in content['data_list']:
|
|
if len(d) > max_img:
|
|
break
|
|
if 'wm_score' in gt and gt['wm_score'] > self.wm_thresh: # wm_score > thresh will be skiped as an img with watermark
|
|
wm_skip += 1
|
|
continue
|
|
data_root = content['data_root']
|
|
if self.using_dlc:
|
|
data_root = data_root.replace('/data/vdb', '/mnt/data', 1)
|
|
img_path = os.path.join(data_root, gt['img_name'])
|
|
info = {}
|
|
info['img_path'] = img_path
|
|
info['caption'] = gt['caption'] if 'caption' in gt else ''
|
|
if 'wm_score' in gt and self.cap_watermark:
|
|
if gt['wm_score'] > 0.5:
|
|
info['caption'] += ' with watermarks '
|
|
else:
|
|
info['caption'] += ' no watermarks '
|
|
if self.place_holder in info['caption']:
|
|
count += 1
|
|
info['caption'] = info['caption'].replace(self.place_holder, " ")
|
|
if 'annotations' in gt:
|
|
polygons = []
|
|
invalid_polygons = []
|
|
texts = []
|
|
languages = []
|
|
pos = []
|
|
color = []
|
|
for annotation in gt['annotations']:
|
|
if len(annotation['polygon']) == 0:
|
|
continue
|
|
if 'valid' in annotation and annotation['valid'] is False:
|
|
invalid_polygons.append(annotation['polygon'])
|
|
continue
|
|
polygons.append(annotation['polygon'])
|
|
texts.append(annotation['text'])
|
|
lang = annotation['language']
|
|
if lang == 'Chinese':
|
|
lang = 'Chinese_tra' if tra_chinese(annotation['text']) else 'Chinese_sim'
|
|
languages.append(lang)
|
|
if 'pos' in annotation:
|
|
pos.append(annotation['pos'])
|
|
if 'color' in annotation:
|
|
color.append(annotation['color'])
|
|
else:
|
|
color.append(default_color)
|
|
info['polygons'] = [np.array(i) for i in polygons]
|
|
info['invalid_polygons'] = [np.array(i) for i in invalid_polygons]
|
|
info['texts'] = texts
|
|
info['language'] = languages
|
|
info['pos'] = pos
|
|
info['color'] = [np.array(i) for i in color]
|
|
d.append(info)
|
|
print(f'{json_path} loaded, imgs={len(d)}, wm_skip={wm_skip}, time={(time.time()-tic):.2f}s')
|
|
if count > 0:
|
|
print(f"Found {count} image's caption contain placeholder: {self.place_holder}, change to ' '...")
|
|
return d
|
|
|
|
def __getitem__(self, item):
|
|
editing_mode = False
|
|
if random.random() < self.mask_img_prob:
|
|
editing_mode = True
|
|
item_dict = {}
|
|
if self.debug: # sample fixed items
|
|
item = self.tmp_items.pop()
|
|
print(f'item = {item}')
|
|
cur_item = self.data_list[item]
|
|
# img
|
|
target = np.array(Image.open(cur_item['img_path']).convert('RGB'))
|
|
if target.shape[0] != self.img_wh or target.shape[1] != self.img_wh:
|
|
target = cv2.resize(target, (self.img_wh, self.img_wh))
|
|
target = (target.astype(np.float32) / 127.5) - 1.0
|
|
item_dict['img'] = target
|
|
# caption
|
|
if self.trunc_cap > 0:
|
|
cur_item['caption'] = truncate_string(cur_item['caption']+'. ', max_words=self.trunc_cap)
|
|
item_dict['img_caption'] = cur_item['caption']
|
|
item_dict['text_caption'] = ''
|
|
item_dict['glyphs'] = []
|
|
item_dict['gly_line'] = []
|
|
item_dict['positions'] = []
|
|
font_hints = []
|
|
font_hints_mask = []
|
|
|
|
item_dict['texts'] = []
|
|
item_dict['language'] = []
|
|
item_dict['inv_mask'] = []
|
|
item_dict['color'] = []
|
|
texts = cur_item.get('texts', [])
|
|
if len(texts) == 0: # padding empty text on image, prevent mismatch between collectives on ranks
|
|
texts = [' ']
|
|
cur_item['color'] = [np.array(default_color)]
|
|
cur_item['polygons'] = [np.array([[10, 10], [100, 10], [100, 100], [10, 100]])]
|
|
cur_item['language'] = ['Latin']
|
|
cur_item['texts'] = texts
|
|
|
|
idxs = [i for i in range(len(texts))]
|
|
if len(texts) > self.max_lines:
|
|
sel_idxs = random.sample(idxs, self.max_lines)
|
|
unsel_idxs = [i for i in idxs if i not in sel_idxs]
|
|
else:
|
|
sel_idxs = idxs
|
|
unsel_idxs = []
|
|
item_dict['color'] = [cur_item['color'][i] for i in sel_idxs]
|
|
if self.color_prob < 1.0:
|
|
for i, c in enumerate(item_dict['color']):
|
|
if random.random() < (1 - self.color_prob):
|
|
item_dict['color'][i] = np.array(default_color)
|
|
item_dict['text_caption'] = get_text_caption(len(sel_idxs), self.place_holder)
|
|
item_dict['polygons'] = [cur_item['polygons'][i] for i in sel_idxs]
|
|
item_dict['texts'] = [cur_item['texts'][i][:self.max_chars] for i in sel_idxs]
|
|
item_dict['language'] = [cur_item['language'][i] for i in sel_idxs]
|
|
# glyphs
|
|
for idx, text in enumerate(item_dict['texts']):
|
|
if self.rand_font:
|
|
lang = item_dict['language'][idx]
|
|
assert lang in self.lang_font
|
|
use_font = random.choice(self.lang_font[lang]) # random font
|
|
else:
|
|
use_font = self.font # arial unicode
|
|
gly_line = draw_glyph(use_font, text)
|
|
if self.render_glyph:
|
|
glyphs = draw_glyph2(use_font, text, item_dict['polygons'][idx], item_dict['color'][idx], scale=self.glyph_scale, width=self.img_wh, height=self.img_wh)
|
|
else:
|
|
glyphs = np.zeros((self.img_wh*self.glyph_scale, self.img_wh*self.glyph_scale, 3), np.float64)
|
|
item_dict['glyphs'] += [glyphs]
|
|
item_dict['gly_line'] += [gly_line]
|
|
# mask_pos
|
|
for polygon in item_dict['polygons']:
|
|
target_area_ratio_pos = [1.0, 1.0] # 0.6--0.9
|
|
item_dict['positions'] += [self.draw_pos(polygon, self.mask_pos_prob, target_area_ratio_pos)]
|
|
if self.font_hint_prob > 0:
|
|
font_hint, font_hint_mask = draw_font_hint(target, polygon, target_area_range=self.font_hint_area, prob=self.font_hint_prob, randaug=self.font_hint_randaug)
|
|
font_hints += [font_hint]
|
|
font_hints_mask += [font_hint_mask]
|
|
|
|
# inv_mask
|
|
invalid_polygons = cur_item['invalid_polygons'] if 'invalid_polygons' in cur_item else []
|
|
if len(texts) > 0:
|
|
invalid_polygons += [cur_item['polygons'][i] for i in unsel_idxs]
|
|
if editing_mode:
|
|
# randomly generate 0 masks (disabled)
|
|
box_num = random.randint(0, 0)
|
|
boxes = generate_random_rectangles(self.img_wh, self.img_wh, box_num)
|
|
boxes = np.array(boxes)
|
|
pos_list = item_dict['positions'].copy()
|
|
for i in range(box_num):
|
|
pos_list += [self.draw_pos(boxes[i], self.mask_pos_prob)]
|
|
invalid_polygons = [] # clear invalid_polygons for editing mode
|
|
mask = self.get_hint(pos_list)
|
|
if fix_masked_img_bug:
|
|
masked_img = (target-mask*10).clip(-1, 1)
|
|
else:
|
|
masked_img = target*(1-mask)
|
|
else:
|
|
if fix_masked_img_bug:
|
|
masked_img = np.zeros_like(target)-1
|
|
else:
|
|
masked_img = np.zeros_like(target)
|
|
item_dict['masked_img'] = masked_img
|
|
item_dict['inv_mask'] = self.draw_inv_mask(invalid_polygons)
|
|
item_dict['hint'] = self.get_hint(item_dict['positions'])
|
|
item_dict['font_hint'] = self.get_hint(font_hints)
|
|
|
|
if self.for_show:
|
|
item_dict['img_name'] = os.path.split(cur_item['img_path'])[-1]
|
|
return item_dict
|
|
if len(texts) > 0:
|
|
del item_dict['polygons']
|
|
# padding
|
|
n_lines = min(len(texts), self.max_lines)
|
|
item_dict['n_lines'] = n_lines
|
|
n_pad = self.max_lines - n_lines
|
|
if n_pad > 0:
|
|
item_dict['glyphs'] += [np.zeros((self.img_wh*self.glyph_scale, self.img_wh*self.glyph_scale, 3))] * n_pad
|
|
item_dict['gly_line'] += [np.zeros((80, 512, 1))] * n_pad
|
|
item_dict['positions'] += [np.zeros((self.img_wh, self.img_wh, 1))] * n_pad
|
|
item_dict['texts'] += [' '] * n_pad
|
|
item_dict['language'] += [' '] * n_pad
|
|
item_dict['color'] += [np.array(default_color)] * n_pad
|
|
|
|
return item_dict
|
|
|
|
def __len__(self):
|
|
return len(self.data_list)
|
|
|
|
def draw_inv_mask(self, polygons):
|
|
img = np.zeros((self.img_wh, self.img_wh))
|
|
for p in polygons:
|
|
pts = p.reshape((-1, 1, 2))
|
|
cv2.fillPoly(img, [pts], color=255)
|
|
img = img[..., None]
|
|
return img/255.
|
|
|
|
def draw_pos(self, ploygon, prob=1.0, target_area_range=[1.0, 1.0]):
|
|
img = np.zeros((self.img_wh, self.img_wh))
|
|
rect = cv2.minAreaRect(ploygon)
|
|
center, size, angle = rect
|
|
w, h = size
|
|
small = False
|
|
min_wh = 20*self.img_wh/512
|
|
if w < min_wh or h < min_wh:
|
|
small = True
|
|
if random.random() < prob:
|
|
pts = ploygon.reshape((-1, 1, 2))
|
|
cv2.fillPoly(img, [pts], color=255)
|
|
# 10% dilate / 10% erode / 5% dilatex2 5% erodex2
|
|
random_value = random.random()
|
|
kernel = np.ones((3, 3), dtype=np.uint8)
|
|
if random_value < 0.7:
|
|
pass
|
|
elif random_value < 0.8:
|
|
img = cv2.dilate(img.astype(np.uint8), kernel, iterations=1)
|
|
elif random_value < 0.9 and not small:
|
|
img = cv2.erode(img.astype(np.uint8), kernel, iterations=1)
|
|
elif random_value < 0.95:
|
|
img = cv2.dilate(img.astype(np.uint8), kernel, iterations=2)
|
|
elif random_value < 1.0 and not small:
|
|
img = cv2.erode(img.astype(np.uint8), kernel, iterations=2)
|
|
# gen a random mask(for editing mode)
|
|
if target_area_range[0] < 1.0 or target_area_range[1] < 1.0:
|
|
area_ratio = random.uniform(target_area_range[0], target_area_range[1])
|
|
long_side, short_side = max(w, h), min(w, h)
|
|
long_axis_mask_length = long_side * (1 - area_ratio)
|
|
angle_rad = np.radians(angle)
|
|
rect_vector = np.array([np.cos(angle_rad), np.sin(angle_rad)])
|
|
if w < h:
|
|
rect_vector = np.array([-rect_vector[1], rect_vector[0]])
|
|
start_offset = random.uniform(0, long_side - long_axis_mask_length)
|
|
start_point = center - rect_vector * (long_side / 2 - start_offset)
|
|
mask_center = start_point + rect_vector * (long_axis_mask_length / 2)
|
|
mask_vector = rect_vector * (long_axis_mask_length / 2)
|
|
short_axis_vector = np.array([-rect_vector[1], rect_vector[0]]) * (short_side / 2)
|
|
mask_corners = np.array([
|
|
mask_center - mask_vector - short_axis_vector,
|
|
mask_center + mask_vector - short_axis_vector,
|
|
mask_center + mask_vector + short_axis_vector,
|
|
mask_center - mask_vector + short_axis_vector
|
|
], dtype=np.int32)
|
|
cv2.fillPoly(img, [mask_corners], color=0)
|
|
img = img[..., None]
|
|
return img/255.
|
|
|
|
def get_hint(self, positions):
|
|
if len(positions) == 0:
|
|
return np.zeros((self.img_wh, self.img_wh, 1))
|
|
return np.sum(positions, axis=0).clip(0, 1)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
'''
|
|
Run this script to show details of your dataset, such as ocr annotations, glyphs, prompts, etc.
|
|
'''
|
|
from tqdm import tqdm
|
|
from matplotlib import pyplot as plt
|
|
import shutil
|
|
|
|
show_imgs_dir = 'show_results'
|
|
show_count = 100
|
|
if os.path.exists(show_imgs_dir):
|
|
shutil.rmtree(show_imgs_dir)
|
|
os.makedirs(show_imgs_dir)
|
|
plt.rcParams['axes.unicode_minus'] = False
|
|
json_paths = [
|
|
'/path/of/your/dataset/data1.json',
|
|
'/path/of/your/dataset/data2.json',
|
|
# ...
|
|
]
|
|
dataset = T3DataSet(json_paths, for_show=True, max_lines=20, glyph_scale=1, mask_img_prob=0.5,
|
|
render_glyph=True, rand_font=True, font_hint_prob=1, font_hint_area=[0.7, 1],
|
|
font_hint_randaug=True, color_prob=1, img_wh=512)
|
|
train_loader = DataLoader(dataset=dataset, batch_size=1, shuffle=False, num_workers=0)
|
|
pbar = tqdm(total=show_count)
|
|
for i, data in enumerate(train_loader):
|
|
if i == show_count:
|
|
break
|
|
img = ((data['img'][0].numpy() + 1.0) / 2.0 * 255).astype(np.uint8)
|
|
masked_img = ((data['masked_img'][0].numpy() + 1.0) / 2.0 * 255)[..., ::-1].astype(np.uint8)
|
|
cv2.imwrite(os.path.join(show_imgs_dir, f'plots_{i}_masked.jpg'), masked_img)
|
|
if 'texts' in data and len(data['texts']) > 0:
|
|
texts = [x[0] for x in data['texts']]
|
|
img = show_bbox_on_image(Image.fromarray(img), data['polygons'], texts)
|
|
cv2.imwrite(os.path.join(show_imgs_dir, f'plots_{i}.jpg'), np.array(img)[..., ::-1])
|
|
with open(os.path.join(show_imgs_dir, f'plots_{i}.txt'), 'w') as fin:
|
|
fin.writelines([data['img_caption'][0], data['text_caption'][0]])
|
|
all_glyphs = []
|
|
for k, glyphs in enumerate(data['glyphs']):
|
|
cv2.imwrite(os.path.join(show_imgs_dir, f'plots_{i}_glyph_{k}.jpg'), (glyphs[0].numpy()*255).astype(np.int32)[..., ::-1])
|
|
all_glyphs += [(glyphs[0].numpy()*255).astype(np.int32)[..., ::-1]]
|
|
if len(all_glyphs) > 0:
|
|
img_allglyph = np.sum(all_glyphs, axis=0)
|
|
black_pixels = (img_allglyph[:, :, 0] < 5) & (img_allglyph[:, :, 1] < 5) & (img_allglyph[:, :, 2] < 5)
|
|
img_allglyph[black_pixels] = [215, 215, 215]
|
|
cv2.imwrite(os.path.join(show_imgs_dir, f'plots_{i}_allglyphs.jpg'), img_allglyph)
|
|
for k, gly_line in enumerate(data['gly_line']):
|
|
cv2.imwrite(os.path.join(show_imgs_dir, f'plots_{i}_gly_line_{k}.jpg'), gly_line[0].numpy().astype(np.int32)*255)
|
|
for k, position in enumerate(data['positions']):
|
|
if position is not None:
|
|
cv2.imwrite(os.path.join(show_imgs_dir, f'plots_{i}_pos_{k}.jpg'), position[0].numpy().astype(np.int32)*255)
|
|
cv2.imwrite(os.path.join(show_imgs_dir, f'plots_{i}_hint.jpg'), data['hint'][0].numpy().astype(np.int32)*255)
|
|
img_font_hint = (1-data['font_hint'][0].numpy()).astype(np.int32)*215
|
|
cv2.imwrite(os.path.join(show_imgs_dir, f'plots_{i}_font_hint.jpg'), img_font_hint)
|
|
cv2.imwrite(os.path.join(show_imgs_dir, f'plots_{i}_inv_mask.jpg'), np.array(img)[..., ::-1]*(1-data['inv_mask'][0].numpy().astype(np.int32)))
|
|
pbar.update(1)
|
|
pbar.close()
|