import os
import numpy as np
import cv2
import random
import math
from PIL import Image, ImageDraw, ImageFont
from torch.utils.data import Dataset, DataLoader
from dataset_util import load, show_bbox_on_image
phrase_list = [
', content and position of the texts are ',
', textual material depicted in the image are ',
', texts that says ',
', captions shown in the snapshot are ',
', with the words of ',
', that reads ',
', the written materials on the picture: ',
', these texts are written on it: ',
', captions are ',
', content of the text in the graphic is '
]
def insert_spaces(string, nSpace):
if nSpace == 0:
return string
new_string = ""
for char in string:
new_string += char + " " * nSpace
return new_string[:-nSpace]
def draw_glyph(font, text):
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))
text_width, text_height = new_font.getsize(text)
offset_x, offset_y = new_font.getoffset(text)
x = (img.width - text_width) // 2
y = (img.height - text_height) // 2 - offset_y//2
draw.text((x, y), text, font=new_font, fill='white')
img = np.expand_dims(np.array(img), axis=2).astype(np.float64)
return img
def draw_glyph2(font, text, polygon, vertAng=10, scale=1, width=512, height=512, add_space=True):
enlarge_polygon = 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 = np.zeros((height*scale, width*scale, 3), np.uint8)
img = Image.fromarray(img)
# infer font size
image4ratio = Image.new("RGB", img.size, "white")
draw = ImageDraw.Draw(image4ratio)
_, _, _tw, _th = draw.textbbox(xy=(0, 0), text=text, font=font)
text_w = min(w, h) * (_tw / _th)
if text_w <= max(w, h):
# add space
if len(text) > 1 and not vert and add_space:
for i in range(1, 100):
text_space = insert_spaces(text, i)
_, _, _tw2, _th2 = draw.textbbox(xy=(0, 0), text=text_space, font=font)
if min(w, h) * (_tw2 / _th2) > max(w, h):
break
text = insert_spaces(text, i-1)
font_size = min(w, h)*0.80
else:
shrink = 0.75 if vert else 0.85
font_size = min(w, h) / (text_w/max(w, h)) * shrink
new_font = font.font_variant(size=int(font_size))
left, top, right, bottom = new_font.getbbox(text)
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:
draw.text((rect[0][0]-text_width//2, rect[0][1]-text_height//2-top), text, font=new_font, fill=(255, 255, 255, 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=(255, 255, 255, 255))
_, _t, _, _b = new_font.getbbox(c)
y_s += _b
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)
img = np.expand_dims(np.array(img.convert('1')), axis=2).astype(np.float64)
return img
def get_caption_pos(ori_caption, pos_idxs, prob=1.0, place_holder='*'):
idx2pos = {
0: " top left",
1: " top",
2: " top right",
3: " left",
4: random.choice([" middle", " center"]),
5: " right",
6: " bottom left",
7: " bottom",
8: " bottom right"
}
new_caption = ori_caption + random.choice(phrase_list)
pos = ''
for i in range(len(pos_idxs)):
if random.random() < prob and pos_idxs[i] > 0:
pos += place_holder + random.choice([' located', ' placed', ' positioned', '']) + random.choice([' at', ' in', ' on']) + idx2pos[pos_idxs[i]] + ', '
else:
pos += place_holder + ' , '
pos = pos[:-2] + '.'
new_caption += pos
return new_caption
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)
class T3DataSet(Dataset):
def __init__(
self,
json_path,
max_lines=5,
max_chars=20,
place_holder='*',
font_path='./font/Arial_Unicode.ttf',
caption_pos_prob=1.0,
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,
):
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.caption_pos_porb = caption_pos_prob
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
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):
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 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 = []
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'])
languages.append(annotation['language'])
if 'pos' in annotation:
pos.append(annotation['pos'])
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
d.append(info)
print(f'{json_path} loaded, imgs={len(d)}, wm_skip={wm_skip}')
if count > 0:
print(f"Found {count} image's caption contain placeholder: {self.place_holder}, change to ' '...")
return d
def __getitem__(self, item):
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] != 512 or target.shape[1] != 512:
target = cv2.resize(target, (512, 512))
target = (target.astype(np.float32) / 127.5) - 1.0
item_dict['img'] = target
# caption
item_dict['caption'] = cur_item['caption']
item_dict['glyphs'] = []
item_dict['gly_line'] = []
item_dict['positions'] = []
item_dict['texts'] = []
item_dict['language'] = []
item_dict['inv_mask'] = []
texts = cur_item.get('texts', [])
if len(texts) > 0:
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 = []
if len(cur_item['pos']) > 0:
pos_idxs = [cur_item['pos'][i] for i in sel_idxs]
else:
pos_idxs = [-1 for i in sel_idxs]
item_dict['caption'] = get_caption_pos(item_dict['caption'], pos_idxs, self.caption_pos_porb, 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']):
gly_line = draw_glyph(self.font, text)
glyphs = draw_glyph2(self.font, text, item_dict['polygons'][idx], scale=self.glyph_scale)
item_dict['glyphs'] += [glyphs]
item_dict['gly_line'] += [gly_line]
# mask_pos
for polygon in item_dict['polygons']:
item_dict['positions'] += [self.draw_pos(polygon, self.mask_pos_prob)]
# 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]
item_dict['inv_mask'] = self.draw_inv_mask(invalid_polygons)
item_dict['hint'] = self.get_hint(item_dict['positions'])
if random.random() < self.mask_img_prob:
# randomly generate 0~3 masks
box_num = random.randint(0, 3)
boxes = generate_random_rectangles(512, 512, 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)]
mask = self.get_hint(pos_list)
masked_img = target*(1-mask)
else:
masked_img = np.zeros_like(target)
item_dict['masked_img'] = masked_img
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((512*self.glyph_scale, 512*self.glyph_scale, 1))] * n_pad
item_dict['gly_line'] += [np.zeros((80, 512, 1))] * n_pad
item_dict['positions'] += [np.zeros((512, 512, 1))] * n_pad
item_dict['texts'] += [' '] * n_pad
item_dict['language'] += [' '] * n_pad
return item_dict
def __len__(self):
return len(self.data_list)
def draw_inv_mask(self, polygons):
img = np.zeros((512, 512))
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):
img = np.zeros((512, 512))
rect = cv2.minAreaRect(ploygon)
w, h = rect[1]
small = False
if w < 20 or h < 20:
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)
img = img[..., None]
return img/255.
def get_hint(self, positions):
if len(positions) == 0:
return np.zeros((512, 512, 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 = 50
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=2, mask_img_prob=1.0, caption_pos_prob=0.0)
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['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().astype(np.int32)*255)
all_glyphs += [glyphs[0].numpy().astype(np.int32)*255]
cv2.imwrite(os.path.join(show_imgs_dir, f'plots_{i}_allglyphs.jpg'), np.sum(all_glyphs, axis=0))
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)
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()