265 lines
13 KiB
Python
265 lines
13 KiB
Python
import cv2
|
|
import torch
|
|
import numpy as np
|
|
from .AnyText2_utils import check_channels, resize_image, separate_pos_imgs, find_polygon, arr2tensor
|
|
from .AnyText2_t3_dataset import draw_glyph, draw_glyph2, draw_font_hint
|
|
from AnyText2ControlDiffusion.cldm.recognizer import crop_image
|
|
from ....UL_common.common import tensor2pil
|
|
import os
|
|
from .AnyText_Infer import generate_rectangles
|
|
|
|
mimic_list = font_hint_image = font_hint_mask =[]
|
|
|
|
def encode(
|
|
model, font, fonts, mask, prompt, texts, latent, mode, sort_radio, a_prompt, n_prompt, revise_pos, random_mask, Random_Gen_Mask_path, image, font_apply, show_glyph, **extra_kwargs
|
|
):
|
|
max_chars = 50
|
|
batch_size, height, width = latent["samples"].shape[0], latent["samples"].shape[2] * 8, latent["samples"].shape[3] * 8 # B, C, H, W
|
|
|
|
sort_radio = '↔' if sort_radio else '↕'
|
|
|
|
img_count = batch_size
|
|
h, w = height, width
|
|
glyline_font_path, text_colors = fonts.glyline_font_path, fonts.text_colors
|
|
n_lines = len(texts)
|
|
|
|
if not random_mask:
|
|
pos_image = np.array(tensor2pil(mask))
|
|
pos_image = 255 - resize_image(pos_image, max(height, width))#[..., :3]
|
|
else:
|
|
generate_rectangles(width, height, n_lines, 500, Random_Gen_Mask_path)
|
|
pos_image = cv2.imread(Random_Gen_Mask_path)[..., ::-1]
|
|
if random_mask:
|
|
pos_image = pos_image[..., 0:1]
|
|
draw_pos = pos_image
|
|
|
|
h, w = draw_pos.shape[:2]
|
|
|
|
if mode:
|
|
mode = 'text-generation'
|
|
revise_pos = revise_pos
|
|
else:
|
|
if image == None:
|
|
raise ValueError('Edit mode need a image.')
|
|
ori_image = np.array(tensor2pil(image))
|
|
mode = 'text-editing'
|
|
revise_pos = False
|
|
|
|
if mode in ['text-generation', 'gen']:
|
|
edit_image = np.zeros((h, w, 3))
|
|
elif mode in ['text-editing', 'edit']:
|
|
if draw_pos is None or ori_image is None:
|
|
return None, -1, "Reference image and position image are needed for text editing!", ""
|
|
if isinstance(ori_image, str):
|
|
ori_image = cv2.imread(ori_image)[..., ::-1]
|
|
assert ori_image is not None, f"Can't read ori_image image from{ori_image}!"
|
|
elif isinstance(ori_image, torch.Tensor):
|
|
ori_image = ori_image.cpu().numpy()
|
|
else:
|
|
assert isinstance(ori_image, np.ndarray), f'Unknown format of ori_image: {type(ori_image)}'
|
|
# edit_image = ori_image[..., ::-1] #bgr2rgb
|
|
edit_image = ori_image.clip(1, 255) # for mask reason
|
|
edit_image = check_channels(edit_image)
|
|
edit_image = resize_image(edit_image, max_length=max(height, width)) # make w h multiple of 64, resize if w or h > max_length
|
|
h, w = edit_image.shape[:2] # change h, w by input ref_img
|
|
|
|
# preprocess pos_imgs(if numpy, make sure it's white pos in black bg)
|
|
if draw_pos is None:
|
|
pos_imgs = np.zeros((w, h, 1))
|
|
if isinstance(draw_pos, str):
|
|
draw_pos = cv2.imread(draw_pos)[..., ::-1]
|
|
assert draw_pos is not None, f"Can't read draw_pos image from{draw_pos}!"
|
|
pos_imgs = 255-draw_pos
|
|
elif isinstance(draw_pos, torch.Tensor):
|
|
pos_imgs = draw_pos.cpu().numpy()
|
|
else:
|
|
# assert isinstance(draw_pos, np.ndarray), f'Unknown format of draw_pos: {type(draw_pos)}'
|
|
pos_imgs = cv2.resize(draw_pos, (w, h))
|
|
if mode in ['text-editing', 'edit']:
|
|
pos_imgs = cv2.resize(pos_imgs, (w, h))
|
|
|
|
pos_imgs = 255 - pos_imgs#[..., 0:1]
|
|
pos_imgs = cv2.convertScaleAbs(pos_imgs)
|
|
_, pos_imgs = cv2.threshold(pos_imgs, 254, 255, cv2.THRESH_BINARY)
|
|
# seprate pos_imgs
|
|
pos_imgs = separate_pos_imgs(pos_imgs, sort_radio)
|
|
|
|
if len(pos_imgs) == 0:
|
|
pos_imgs = [np.zeros((h, w, 1))]
|
|
if len(pos_imgs) < n_lines:
|
|
if n_lines == 1 and texts[0] == ' ':
|
|
# pass # text-to-image without text
|
|
print('\033[93m', f'Warning: text-to-image without text.', '\033[0m')
|
|
else:
|
|
raise ValueError(f'Found {len(pos_imgs)} positions that < needed {n_lines} from prompt, check and try again(手绘遮罩数少于要绘制的文本数，检查修改再重试)!')
|
|
elif len(pos_imgs) > n_lines:
|
|
print('\033[93m', f'Warning: found {len(pos_imgs)} positions that > needed {n_lines} from prompt.', '\033[0m')
|
|
|
|
# get pre_pos, poly_list, hint that needed for anytext
|
|
pre_pos = []
|
|
poly_list = []
|
|
for input_pos in pos_imgs:
|
|
if input_pos.mean() != 0:
|
|
input_pos = input_pos[..., np.newaxis] if len(input_pos.shape) == 2 else input_pos
|
|
poly, pos_img = find_polygon(input_pos)
|
|
pre_pos += [pos_img/255.]
|
|
poly_list += [poly]
|
|
else:
|
|
pre_pos += [np.zeros((h, w, 1))]
|
|
poly_list += [None]
|
|
np_hint = np.sum(pre_pos, axis=0).clip(0, 1)
|
|
# prepare info dict
|
|
info = {}
|
|
info['glyphs'] = []
|
|
info['gly_line'] = []
|
|
info['positions'] = []
|
|
info['n_lines'] = [len(texts)]*img_count
|
|
font_hint = []
|
|
font_paths = ['None' for i in range(len(texts))]
|
|
if glyline_font_path:
|
|
glyline_font_path = glyline_font_path[:len(texts)]
|
|
font_paths[:len(glyline_font_path)] = glyline_font_path
|
|
info['colors'] = [np.array([500, 500, 500]) for i in range(len(texts))]
|
|
if text_colors:
|
|
text_colors = text_colors.strip().split()[:len(texts)]
|
|
info['colors'][:len(text_colors)] = [np.array([int(p) for p in s.split(',')]) for s in text_colors]
|
|
|
|
gly_pos_imgs = []
|
|
font_hint_mimic_imgs = []
|
|
|
|
gly_line_list = []
|
|
glyphs_list = []
|
|
|
|
for i in range(len(texts)):
|
|
text = texts[i]
|
|
if len(text) > max_chars:
|
|
# str_warning = f'"{text}" length > max_chars: {max_chars}, will be cut off...'
|
|
print(f'"{text}" length > max_chars: {max_chars}, will be cut off...')
|
|
text = text[:max_chars]
|
|
|
|
gly_scale = model.model.control_model.glyph_scale
|
|
|
|
if pre_pos[i].mean() != 0:
|
|
if font_apply:
|
|
gly_line = draw_glyph(font, text, info['colors'][i], show_glyph)
|
|
glyphs = draw_glyph2(font, text, poly_list[i], info['colors'][i], scale=gly_scale, width=w, height=h, add_space=True, show_glyph=show_glyph)
|
|
else:
|
|
gly_line = draw_glyph(font_paths[i], text, info['colors'][i], show_glyph)
|
|
glyphs = draw_glyph2(font_paths[i], text, poly_list[i], info['colors'][i], scale=gly_scale, width=w, height=h, add_space=True, show_glyph=show_glyph)
|
|
|
|
gly_line_list+=[gly_line]
|
|
glyphs_list+=[glyphs]
|
|
|
|
if i < len(font_hint_image) and font_hint_image[i] is not None:
|
|
hint_poly = font_hint_mask[i]
|
|
poly, _ = find_polygon(hint_poly)
|
|
font_hint_mimic_img, _ = draw_font_hint((font_hint_image[i]/127.5-1), poly)
|
|
font_hint_mimic_img = torch.from_numpy(font_hint_mimic_img*255).permute(2, 0, 1).repeat(3, 1, 1)
|
|
font_hint_mimic_imgs += [crop_image(font_hint_mimic_img, hint_poly)] # chw, tensor, 0-255
|
|
font_paths[i] = 'None' # not render
|
|
else:
|
|
font_hint_mimic_imgs += [None]
|
|
if font_apply:
|
|
font_hint_line = draw_glyph2(font_paths[i], text, poly_list[i], np.array([255, 255, 255]), scale=1, width=w, height=h, add_space=True, show_glyph=show_glyph)
|
|
else:
|
|
font_hint_line = draw_glyph2(font_paths[i], text, poly_list[i], info['colors'][i], scale=1, width=w, height=h, add_space=True, show_glyph=show_glyph)
|
|
gly_pos_img = cv2.drawContours(glyphs*255, [poly_list[i]*gly_scale], 0, (255, 255, 255), 1)
|
|
|
|
if revise_pos:
|
|
resize_gly = cv2.resize(glyphs, (pre_pos[i].shape[1], pre_pos[i].shape[0]))
|
|
new_pos = cv2.morphologyEx((resize_gly*255).astype(np.uint8), cv2.MORPH_CLOSE, kernel=np.ones((resize_gly.shape[0]//10, resize_gly.shape[1]//10), dtype=np.uint8), iterations=1)
|
|
new_pos = new_pos[..., np.newaxis] if len(new_pos.shape) == 2 else new_pos
|
|
contours, _ = cv2.findContours(new_pos[..., 0:1], cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
|
|
if len(contours) != 1:
|
|
str_warning = f'Fail to revise position {i} to bounding rect, remain position unchanged...'
|
|
else:
|
|
rect = cv2.minAreaRect(contours[0])
|
|
poly = np.int0(cv2.boxPoints(rect))
|
|
pre_pos[i] = cv2.drawContours(new_pos, [poly], -1, 255, -1) / 255.
|
|
gly_pos_img = cv2.drawContours(glyphs*255, [poly*gly_scale], 0, (255, 255, 255), 1)
|
|
|
|
gly_pos_imgs += [gly_pos_img] # for show
|
|
else:
|
|
glyphs = np.zeros((h*gly_scale, w*gly_scale, 3))
|
|
gly_line = np.zeros((80, 512, 1))
|
|
gly_pos_imgs += [np.zeros((h*gly_scale, w*gly_scale, 1))] # for show
|
|
font_hint_line = np.zeros((h, w, 3))
|
|
pos = pre_pos[i][..., 0:1]
|
|
info['glyphs'] += [arr2tensor(glyphs, img_count, model.load_device, model.dtype)]
|
|
info['gly_line'] += [arr2tensor(gly_line, img_count, model.load_device, model.dtype)]
|
|
info['positions'] += [arr2tensor(pos, img_count, model.load_device, model.dtype)]
|
|
info['colors'][i] = arr2tensor(info['colors'][i], img_count, model.load_device, model.dtype)/255.
|
|
font_hint += [font_hint_line]
|
|
font_hint_mimic_imgs = [font_hint_mimic_imgs] * img_count
|
|
model.model.embedding_manager.font_hint_mimic_imgs = font_hint_mimic_imgs
|
|
# get masked_x
|
|
masked_img = ((edit_image.astype(np.float32) / 127.5) - 1.0 - np_hint*10).clip(-1, 1)
|
|
masked_img = np.transpose(masked_img, (2, 0, 1))
|
|
masked_img = torch.from_numpy(masked_img.copy()).float().cuda(0)
|
|
|
|
# if self.use_fp16:
|
|
# masked_img = masked_img.half()
|
|
masked_img = masked_img.to(model.dtype)
|
|
|
|
encoder_posterior = model.model.encode_first_stage(masked_img[None, ...])
|
|
masked_x = model.model.get_first_stage_encoding(encoder_posterior).detach()
|
|
|
|
# if self.use_fp16:
|
|
# masked_x = masked_x.half()
|
|
masked_x = masked_x.to(model.dtype)
|
|
|
|
info['masked_x'] = torch.cat([masked_x for _ in range(img_count)], dim=0)
|
|
|
|
hint = arr2tensor(np_hint, img_count, model.load_device, model.dtype)
|
|
|
|
font_hint_fg = np.sum(font_hint, axis=0).clip(0, 1)[..., 0:1]*255
|
|
if fonts.font_hollow and font_hint_fg.mean() > 0:
|
|
img = cv2.imread(os.path.join(os.path.dirname(os.path.abspath(__file__)), "bg_noise.png"))
|
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
|
img = cv2.resize(img, (font_hint_fg.shape[1], font_hint_fg.shape[0]))
|
|
img[img < 230] = 0
|
|
font_hint_bg = cv2.adaptiveThreshold(img.astype(np.uint8), 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)
|
|
kernel1 = np.ones((2, 2), dtype=np.uint8)
|
|
kernel2 = np.ones((3, 3), dtype=np.uint8)
|
|
dilate_img1 = cv2.dilate(font_hint_fg[..., 0].astype(np.uint8), kernel1, iterations=1)
|
|
dilate_img2 = cv2.dilate(font_hint_fg[..., 0].astype(np.uint8), kernel2, iterations=1)
|
|
dilate_text = dilate_img2 - dilate_img1
|
|
result = (font_hint_fg[..., 0]-font_hint_bg + dilate_text).clip(0, 255)
|
|
font_hint_bg[font_hint_fg[..., 0] > 0] = 0
|
|
result = (result + font_hint_bg).clip(0, 255)
|
|
font_hint_bg = result[..., None]
|
|
else:
|
|
font_hint_bg = font_hint_fg
|
|
|
|
text_prompt = "#"
|
|
for i in texts:
|
|
text_prompt+=" *"
|
|
|
|
info['font_hint'] = arr2tensor((font_hint_bg/255), img_count, model.load_device, model.dtype)
|
|
cond = model.model.get_learned_conditioning(dict(c_concat=[hint], c_crossattn=[[[prompt + ', ' + a_prompt] * img_count, [text_prompt.replace("# ", "")] * img_count]], text_info=info))
|
|
un_cond = model.model.get_learned_conditioning(dict(c_concat=[hint], c_crossattn=[[[n_prompt] * img_count, [""] * img_count]], text_info=info))
|
|
|
|
|
|
# font_img = []
|
|
# for f_i in gly_pos_imgs:
|
|
# f_i = torch.from_numpy(f_i).permute(2,0,1)
|
|
# font_img.append(f_i)
|
|
# font_img = torch.cat(font_img, dim=0)
|
|
|
|
# gly_line_img = []
|
|
# glyphs_img = []
|
|
# for g_la, g_lb in zip(gly_line_list, glyphs_list):
|
|
# g_la, g_lb = torch.from_numpy(g_la).permute(2,0,1), torch.from_numpy(g_lb).permute(2,0,1)
|
|
# gly_line_img.append(g_la)
|
|
# glyphs_img.append(g_lb)
|
|
# gly_line_img = torch.cat(gly_line_img, dim=0)
|
|
# glyphs_img = torch.cat(glyphs_img, dim=0)
|
|
|
|
# glyph_img = torch.from_numpy(np.sum(np.stack(gly_pos_imgs), axis=0).clip(0, 255).astype(np.uint8)).permute(2,0,1)
|
|
|
|
# font_hint_img = torch.from_numpy(np.repeat(font_hint_bg.astype(np.uint8), 3, axis=2)).permute(2,0,1)
|
|
|
|
glyph_img = font_hint_img = []
|
|
|
|
# return cond, un_cond, masked_x, font_img, gly_line_img, glyphs_img, masked_img.clone().cpu().float(), h, w, font_hint_img, glyph_img
|
|
return cond, un_cond, h, w, font_hint_img, glyph_img |