wehub-resource-sync
224 lines
8.6 KiB
Python
224 lines
8.6 KiB
Python
# ------------------------------------------
|
|
# TextDiffuser: Diffusion Models as Text Painters
|
|
# Paper Link: https://arxiv.org/abs/2305.10855
|
|
# Code Link: https://github.com/microsoft/unilm/tree/master/textdiffuser
|
|
# Copyright (c) Microsoft Corporation.
|
|
# This file aims to predict the layout of keywords in user prompts.
|
|
# ------------------------------------------
|
|
|
|
import warnings
|
|
warnings.filterwarnings("ignore", category=DeprecationWarning)
|
|
|
|
import re
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
from transformers import CLIPTokenizer
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
from util import get_width, get_key_words, adjust_overlap_box, shrink_box, adjust_font_size, alphabet_dic
|
|
from model.layout_transformer import LayoutTransformer, TextConditioner
|
|
from termcolor import colored
|
|
|
|
# import layout transformer
|
|
model = LayoutTransformer().cuda().eval()
|
|
model.load_state_dict(torch.load('textdiffuser-ckpt/layout_transformer.pth'))
|
|
|
|
# import text encoder and tokenizer
|
|
text_encoder = TextConditioner().cuda().eval()
|
|
tokenizer = CLIPTokenizer.from_pretrained('openai/clip-vit-large-patch14')
|
|
|
|
|
|
def process_caption(font_path, caption, keywords):
|
|
# remove punctuations. please remove this statement if you want to paint punctuations
|
|
caption = re.sub(u"([^\u0041-\u005a\u0061-\u007a\u0030-\u0039])", " ", caption)
|
|
|
|
# tokenize it into ids and get length
|
|
caption_words = tokenizer([caption], truncation=True, max_length=77, return_length=True, return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
|
|
caption_words_ids = caption_words['input_ids'] # (1, 77)
|
|
length = caption_words['length'] # (1, )
|
|
|
|
# convert id to words
|
|
words = tokenizer.convert_ids_to_tokens(caption_words_ids.view(-1).tolist())
|
|
words = [i.replace('</w>', '') for i in words]
|
|
words_valid = words[:int(length)]
|
|
|
|
# store the box coordinates and state of each token
|
|
info_array = np.zeros((77,5)) # (77, 5)
|
|
|
|
# split the caption into words and convert them into lower case
|
|
caption_split = caption.split()
|
|
caption_split = [i.lower() for i in caption_split]
|
|
|
|
start_dic = {} # get the start index of each word
|
|
state_list = [] # 0: start, 1: middle, 2: special token
|
|
word_match_list = [] # the index of the word in the caption
|
|
current_caption_index = 0
|
|
current_match = ''
|
|
for i in range(length):
|
|
|
|
# the first and last token are special tokens
|
|
if i == 0 or i == length-1:
|
|
state_list.append(2)
|
|
word_match_list.append(127)
|
|
continue
|
|
|
|
if current_match == '':
|
|
state_list.append(0)
|
|
start_dic[current_caption_index] = i
|
|
else:
|
|
state_list.append(1)
|
|
|
|
current_match += words_valid[i]
|
|
word_match_list.append(current_caption_index)
|
|
if current_match == caption_split[current_caption_index]:
|
|
current_match = ''
|
|
current_caption_index += 1
|
|
|
|
while len(state_list) < 77:
|
|
state_list.append(127)
|
|
while len(word_match_list) < 77:
|
|
word_match_list.append(127)
|
|
|
|
length_list = []
|
|
width_list =[]
|
|
for i in range(len(word_match_list)):
|
|
if word_match_list[i] == 127:
|
|
length_list.append(0)
|
|
width_list.append(0)
|
|
else:
|
|
length_list.append(len(caption.split()[word_match_list[i]]))
|
|
width_list.append(get_width(font_path, caption.split()[word_match_list[i]]))
|
|
|
|
while len(length_list) < 77:
|
|
length_list.append(127)
|
|
width_list.append(0)
|
|
|
|
length_list = torch.Tensor(length_list).long() # (77, )
|
|
width_list = torch.Tensor(width_list).long() # (77, )
|
|
|
|
boxes = []
|
|
duplicate_dict = {} # some words may appear more than once
|
|
for keyword in keywords:
|
|
keyword = keyword.lower()
|
|
if keyword in caption_split:
|
|
if keyword not in duplicate_dict:
|
|
duplicate_dict[keyword] = caption_split.index(keyword)
|
|
index = caption_split.index(keyword)
|
|
else:
|
|
if duplicate_dict[keyword]+1 < len(caption_split) and keyword in caption_split[duplicate_dict[keyword]+1:]:
|
|
index = duplicate_dict[keyword] + caption_split[duplicate_dict[keyword]+1:].index(keyword)
|
|
duplicate_dict[keyword] = index
|
|
else:
|
|
continue
|
|
|
|
index = caption_split.index(keyword)
|
|
index = start_dic[index]
|
|
info_array[index][0] = 1
|
|
|
|
box = [0,0,0,0]
|
|
boxes.append(list(box))
|
|
info_array[index][1:] = box
|
|
|
|
boxes_length = len(boxes)
|
|
if boxes_length > 8:
|
|
boxes = boxes[:8]
|
|
while len(boxes) < 8:
|
|
boxes.append([0,0,0,0])
|
|
|
|
return caption, length_list, width_list, torch.from_numpy(info_array), words, torch.Tensor(state_list).long(), torch.Tensor(word_match_list).long(), torch.Tensor(boxes), boxes_length
|
|
|
|
|
|
def get_layout_from_prompt(args):
|
|
|
|
# prompt = args.prompt
|
|
font_path = args.font_path
|
|
keywords = get_key_words(args.prompt)
|
|
|
|
print(f'{colored("[!]", "red")} Detected keywords: {keywords} from prompt {args.prompt}')
|
|
|
|
text_embedding, mask = text_encoder(args.prompt) # (1, 77 768) / (1, 77)
|
|
|
|
# process all relevant info
|
|
caption, length_list, width_list, target, words, state_list, word_match_list, boxes, boxes_length = process_caption(font_path, args.prompt, keywords)
|
|
target = target.cuda().unsqueeze(0) # (77, 5)
|
|
width_list = width_list.cuda().unsqueeze(0) # (77, )
|
|
length_list = length_list.cuda().unsqueeze(0) # (77, )
|
|
state_list = state_list.cuda().unsqueeze(0) # (77, )
|
|
word_match_list = word_match_list.cuda().unsqueeze(0) # (77, )
|
|
|
|
padding = torch.zeros(1, 1, 4).cuda()
|
|
boxes = boxes.unsqueeze(0).cuda()
|
|
right_shifted_boxes = torch.cat([padding, boxes[:,0:-1,:]],1) # (1, 8, 4)
|
|
|
|
# inference
|
|
return_boxes= []
|
|
with torch.no_grad():
|
|
for box_index in range(boxes_length):
|
|
|
|
if box_index == 0:
|
|
encoder_embedding = None
|
|
|
|
output, encoder_embedding = model(text_embedding, length_list, width_list, mask, state_list, word_match_list, target, right_shifted_boxes, train=False, encoder_embedding=encoder_embedding)
|
|
output = torch.clamp(output, min=0, max=1) # (1, 8, 4)
|
|
|
|
# add overlap detection
|
|
output = adjust_overlap_box(output, box_index) # (1, 8, 4)
|
|
|
|
right_shifted_boxes[:,box_index+1,:] = output[:,box_index,:]
|
|
xmin, ymin, xmax, ymax = output[0, box_index, :].tolist()
|
|
return_boxes.append([xmin, ymin, xmax, ymax])
|
|
|
|
|
|
# print the location of keywords
|
|
print(f'index\tkeyword\tx_min\ty_min\tx_max\ty_max')
|
|
for index, keyword in enumerate(keywords):
|
|
x_min = int(return_boxes[index][0] * 512)
|
|
y_min = int(return_boxes[index][1] * 512)
|
|
x_max = int(return_boxes[index][2] * 512)
|
|
y_max = int(return_boxes[index][3] * 512)
|
|
print(f'{index}\t{keyword}\t{x_min}\t{y_min}\t{x_max}\t{y_max}')
|
|
|
|
|
|
# paint the layout
|
|
render_image = Image.new('RGB', (512, 512), (255, 255, 255))
|
|
draw = ImageDraw.Draw(render_image)
|
|
segmentation_mask = Image.new("L", (512,512), 0)
|
|
segmentation_mask_draw = ImageDraw.Draw(segmentation_mask)
|
|
|
|
for index, box in enumerate(return_boxes):
|
|
box = [int(i*512) for i in box]
|
|
xmin, ymin, xmax, ymax = box
|
|
|
|
width = xmax - xmin
|
|
height = ymax - ymin
|
|
text = keywords[index]
|
|
|
|
font_size = adjust_font_size(args, width, height, draw, text)
|
|
font = ImageFont.truetype(args.font_path, font_size)
|
|
|
|
# draw.rectangle([xmin, ymin, xmax,ymax], outline=(255,0,0))
|
|
draw.text((xmin, ymin), text, font=font, fill=(0, 0, 0))
|
|
|
|
boxes = []
|
|
for i, char in enumerate(text):
|
|
|
|
# paint character-level segmentation masks
|
|
# https://github.com/python-pillow/Pillow/issues/3921
|
|
bottom_1 = font.getsize(text[i])[1]
|
|
right, bottom_2 = font.getsize(text[:i+1])
|
|
bottom = bottom_1 if bottom_1 < bottom_2 else bottom_2
|
|
width, height = font.getmask(char).size
|
|
right += xmin
|
|
bottom += ymin
|
|
top = bottom - height
|
|
left = right - width
|
|
|
|
char_box = (left, top, right, bottom)
|
|
boxes.append(char_box)
|
|
|
|
char_index = alphabet_dic[char]
|
|
segmentation_mask_draw.rectangle(shrink_box(char_box, scale_factor = 0.9), fill=char_index)
|
|
|
|
print(f'{colored("[√]", "green")} Layout is successfully generated')
|
|
return render_image, segmentation_mask
|