chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,223 @@
|
||||
# ------------------------------------------
|
||||
# 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
|
||||
@@ -0,0 +1,107 @@
|
||||
# ------------------------------------------
|
||||
# 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 define the Layout Transformer for predicting the layout of keywords.
|
||||
# ------------------------------------------
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from transformers import CLIPTokenizer, CLIPTextModel
|
||||
|
||||
class TextConditioner(nn.Module):
|
||||
|
||||
def __init__(self):
|
||||
super(TextConditioner, self).__init__()
|
||||
self.transformer = CLIPTextModel.from_pretrained('openai/clip-vit-large-patch14')
|
||||
self.tokenizer = CLIPTokenizer.from_pretrained('openai/clip-vit-large-patch14')
|
||||
|
||||
# fix
|
||||
self.transformer.eval()
|
||||
for param in self.transformer.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
def forward(self, prompt_list):
|
||||
batch_encoding = self.tokenizer(prompt_list, truncation=True, max_length=77, return_length=True, return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
|
||||
text_embedding = self.transformer(batch_encoding["input_ids"].cuda())
|
||||
return text_embedding.last_hidden_state.cuda(), batch_encoding["attention_mask"].cuda() # 1, 77, 768 / 1, 768
|
||||
|
||||
|
||||
class LayoutTransformer(nn.Module):
|
||||
|
||||
def __init__(self, layer_number=2):
|
||||
super(LayoutTransformer, self).__init__()
|
||||
|
||||
self.encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
|
||||
self.transformer = torch.nn.TransformerEncoder(
|
||||
self.encoder_layer, num_layers=layer_number
|
||||
)
|
||||
|
||||
self.decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
|
||||
self.decoder_transformer = torch.nn.TransformerDecoder(
|
||||
self.decoder_layer, num_layers=layer_number
|
||||
)
|
||||
|
||||
self.mask_embedding = nn.Embedding(2,512)
|
||||
self.length_embedding = nn.Embedding(256,512)
|
||||
self.width_embedding = nn.Embedding(256,512)
|
||||
self.position_embedding = nn.Embedding(256,512)
|
||||
self.state_embedding = nn.Embedding(256,512)
|
||||
self.match_embedding = nn.Embedding(256,512)
|
||||
|
||||
self.x_embedding = nn.Embedding(512,512)
|
||||
self.y_embedding = nn.Embedding(512,512)
|
||||
self.w_embedding = nn.Embedding(512,512)
|
||||
self.h_embedding = nn.Embedding(512,512)
|
||||
|
||||
self.encoder_target_embedding = nn.Embedding(256,512)
|
||||
|
||||
self.input_layer = nn.Sequential(
|
||||
nn.Linear(768, 512),
|
||||
nn.ReLU(),
|
||||
nn.Linear(512, 512),
|
||||
)
|
||||
|
||||
self.output_layer = nn.Sequential(
|
||||
nn.Linear(512, 128),
|
||||
nn.ReLU(),
|
||||
nn.Linear(128, 4),
|
||||
)
|
||||
|
||||
def forward(self, x, length, width, mask, state, match, target, right_shifted_boxes, train=False, encoder_embedding=None):
|
||||
|
||||
# detect whether the encoder_embedding is cached
|
||||
if encoder_embedding is None:
|
||||
# augmentation
|
||||
if train:
|
||||
width = width + torch.randint(-3, 3, (width.shape[0], width.shape[1])).cuda()
|
||||
|
||||
x = self.input_layer(x) # (1, 77, 512)
|
||||
width_embedding = self.width_embedding(torch.clamp(width, 0, 255).long()) # (1, 77, 512)
|
||||
encoder_target_embedding = self.encoder_target_embedding(target[:,:,0].long()) # (1, 77, 512)
|
||||
pe_embedding = self.position_embedding(torch.arange(77).cuda()).unsqueeze(0) # (1, 77, 512)
|
||||
total_embedding = x + width_embedding + pe_embedding + encoder_target_embedding # combine all the embeddings (1, 77, 512)
|
||||
total_embedding = total_embedding.permute(1,0,2) # (77, 1, 512)
|
||||
encoder_embedding = self.transformer(total_embedding) # (77, 1, 512)
|
||||
|
||||
right_shifted_boxes_resize = (right_shifted_boxes * 512).long() # (1, 8, 4)
|
||||
right_shifted_boxes_resize = torch.clamp(right_shifted_boxes_resize, 0, 511) # (1, 8, 4)
|
||||
|
||||
# decoder pe
|
||||
pe_decoder = torch.arange(8).cuda() # (8, )
|
||||
pe_embedding_decoder = self.position_embedding(pe_decoder).unsqueeze(0) # (1, 8, 512)
|
||||
decoder_input = pe_embedding_decoder + self.x_embedding(right_shifted_boxes_resize[:,:,0]) + self.y_embedding(right_shifted_boxes_resize[:,:,1]) + self.w_embedding(right_shifted_boxes_resize[:,:,2]) + self.h_embedding(right_shifted_boxes_resize[:,:,3]) # (1, 8, 512)
|
||||
decoder_input = decoder_input.permute(1,0,2) # (8, 1, 512)
|
||||
|
||||
# generate triangular mask
|
||||
mask = nn.Transformer.generate_square_subsequent_mask(8) # (8, 8)
|
||||
mask = mask.cuda() # (8, 8)
|
||||
decoder_result = self.decoder_transformer(decoder_input, encoder_embedding, tgt_mask=mask) # (8, 1, 512)
|
||||
decoder_result = decoder_result.permute(1,0,2) # (1, 8, 512)
|
||||
|
||||
box_prediction = self.output_layer(decoder_result) # (1, 8, 4)
|
||||
return box_prediction, encoder_embedding
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
# ------------------------------------------
|
||||
# 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 define the architecture of unet.
|
||||
# ------------------------------------------
|
||||
|
||||
import torch.nn.functional as F
|
||||
from model.text_segmenter.unet_parts import *
|
||||
|
||||
|
||||
class UNet(nn.Module):
|
||||
def __init__(self, n_channels, n_classes, bilinear=True):
|
||||
super(UNet, self).__init__()
|
||||
self.n_channels = n_channels
|
||||
self.n_classes = n_classes
|
||||
self.bilinear = bilinear
|
||||
|
||||
self.inc = DoubleConv(n_channels, 64)
|
||||
self.down1 = Down(64, 128)
|
||||
self.down2 = Down(128, 256)
|
||||
self.down3 = Down(256, 512)
|
||||
factor = 2 if bilinear else 1
|
||||
self.down4 = Down(512, 1024 // factor)
|
||||
self.up1 = Up(1024, 512 // factor, bilinear)
|
||||
self.up2 = Up(512, 256 // factor, bilinear)
|
||||
self.up3 = Up(256, 128 // factor, bilinear)
|
||||
self.up4 = Up(128, 64, bilinear)
|
||||
self.outc = OutConv(64, n_classes)
|
||||
|
||||
def forward(self, x):
|
||||
x1 = self.inc(x)
|
||||
x2 = self.down1(x1)
|
||||
x3 = self.down2(x2)
|
||||
x4 = self.down3(x3)
|
||||
x5 = self.down4(x4)
|
||||
x = self.up1(x5, x4)
|
||||
x = self.up2(x, x3)
|
||||
x = self.up3(x, x2)
|
||||
x = self.up4(x, x1)
|
||||
logits = self.outc(x)
|
||||
# logits = torch.sigmoid(logits)
|
||||
return logits
|
||||
|
||||
if __name__ == '__main__':
|
||||
net = UNet(39,39,True)
|
||||
|
||||
net = net.cuda()
|
||||
|
||||
image = torch.Tensor(32,39,64,64).cuda()
|
||||
result = net(image)
|
||||
print(result.shape)
|
||||
@@ -0,0 +1,82 @@
|
||||
# ------------------------------------------
|
||||
# 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 define the architecture of unet.
|
||||
# ------------------------------------------
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class DoubleConv(nn.Module):
|
||||
"""(convolution => [BN] => ReLU) * 2"""
|
||||
|
||||
def __init__(self, in_channels, out_channels, mid_channels=None):
|
||||
super().__init__()
|
||||
if not mid_channels:
|
||||
mid_channels = out_channels
|
||||
self.double_conv = nn.Sequential(
|
||||
nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1),
|
||||
nn.BatchNorm2d(mid_channels),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1),
|
||||
nn.BatchNorm2d(out_channels),
|
||||
nn.ReLU(inplace=True)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.double_conv(x)
|
||||
|
||||
|
||||
class Down(nn.Module):
|
||||
"""Downscaling with maxpool then double conv"""
|
||||
|
||||
def __init__(self, in_channels, out_channels):
|
||||
super().__init__()
|
||||
self.maxpool_conv = nn.Sequential(
|
||||
nn.MaxPool2d(2),
|
||||
DoubleConv(in_channels, out_channels)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.maxpool_conv(x)
|
||||
|
||||
|
||||
class Up(nn.Module):
|
||||
"""Upscaling then double conv"""
|
||||
|
||||
def __init__(self, in_channels, out_channels, bilinear=True):
|
||||
super().__init__()
|
||||
|
||||
# if bilinear, use the normal convolutions to reduce the number of channels
|
||||
if bilinear:
|
||||
self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
|
||||
self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
|
||||
else:
|
||||
self.up = nn.ConvTranspose2d(in_channels , in_channels // 2, kernel_size=2, stride=2)
|
||||
self.conv = DoubleConv(in_channels, out_channels)
|
||||
|
||||
|
||||
def forward(self, x1, x2):
|
||||
x1 = self.up(x1)
|
||||
# input is CHW
|
||||
diffY = x2.size()[2] - x1.size()[2]
|
||||
diffX = x2.size()[3] - x1.size()[3]
|
||||
|
||||
x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
|
||||
diffY // 2, diffY - diffY // 2])
|
||||
|
||||
x = torch.cat([x2, x1], dim=1)
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class OutConv(nn.Module):
|
||||
def __init__(self, in_channels, out_channels):
|
||||
super(OutConv, self).__init__()
|
||||
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
|
||||
|
||||
def forward(self, x):
|
||||
return self.conv(x)
|
||||
Reference in New Issue
Block a user