Spaces:
Runtime error
Runtime error
File size: 4,983 Bytes
ca4133a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 |
import os
import numpy as np
import PIL
from PIL import Image
from torch.utils.data import Dataset
import random
from PIL import Image
from PIL import ImageDraw
from PIL import ImageFont
from .font_list import font_list
from .font_list_single import font_list_single
import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)
import torchvision.transforms as T
color_multi_font =[
'#a3001b',
'#175c7a',
'#5cd6ce',
'#d1440c',
'#1f1775',
'#9e1e46',
'#9b100d',
'#d3760c',
'#e8bb06',
'#5135ad',
'#366993',
'#470e7c',
'#070707',
'#053d22',
'#7a354c',
'#7c0b03'
]
alphabets = [
'A','B','C','D','E','F','G','H','I','J','K','L',
'M','N','O','P','Q','R','S','T','U','V','W','X', 'Y','Z'
]
class Rasterizer(Dataset):
def __init__(self,
text = "R",
style_word = "DRAGON",
data_path = "data/DRAGON",
alternate_glyph = None,
img_size = 256,
num_samples = 1001,
make_black = False,
one_font = True,
full_word = False,
font_name = None,
just_use_style=False,
use_alt=False):
self.text = text
self.img_size = img_size
self.interpolation = PIL.Image.BILINEAR
self.num_samples = num_samples
self.style_word = style_word
self.data_path = data_path
self.alternate_glyph = alternate_glyph
self.dict = {}
self.data_img = []
self.alt_gly = []
self.make_black = make_black
self.one_font = one_font
self.classes = []
if font_name is not None:
self.fontname = font_name
else:
fontname = random.choice(font_list_single)
self.fontname = fontname
self.full_word = full_word
self.just_use_style = just_use_style
self.use_alt = use_alt
self.load_back()
def load_back(self):
style_only = self.style_word.split(" ")[0]
self.data_path = f"data_style/{style_only}"
self.data_img = []
for file in os.listdir(self.data_path):
self.data_img.append(os.path.join(self.data_path, file))
def __len__(self):
return self.num_samples
def getSize(self, txt, font):
testImg = Image.new('RGB', (1, 1))
testDraw = ImageDraw.Draw(testImg)
return testDraw.textsize(txt, font)
def __getitem__(self, i):
output = {}
fontname = random.choice(font_list)
if self.one_font:
fontname = self.fontname
if self.full_word:
font = ImageFont.truetype(fontname, 256)
width, height = self.getSize(self.text, font)
image = Image.new('RGB', (width+32, height+128), "white")
d = ImageDraw.Draw(image)
colorFont = random.choice(color_multi_font)
d.text((16,64), self.text, fill=colorFont, font=font)
img = np.array(image).astype(np.uint8)
image = Image.fromarray(img)
image = image.resize((self.img_size, self.img_size), resample=self.interpolation)
image = np.array(image).astype(np.uint8)
image_text = (image / 127.5 - 1.0).astype(np.float32)
else:
rand_img = str(random.randint(0,15))
fontname_t = fontname.split(".")[0]
dir_font = f"data_fonts/{fontname_t}/{self.text}/{rand_img}"+".png"
image = Image.open(dir_font)
if not image.mode == "RGB":
image = image.convert("RGB")
image = image.resize((self.img_size, self.img_size), resample=self.interpolation)
image = np.array(image).astype(np.uint8)
image_text = (image / 127.5 - 1.0).astype(np.float32)
output["image"] = image_text
output["caption"] = self.text
##################################################################################
output2 = {}
ind = i % (len(self.data_img)-1)
image = Image.open(self.data_img[ind])
if not image.mode == "RGB":
image = image.convert("RGB")
image = image.resize((self.img_size, self.img_size), resample=self.interpolation)
image = np.array(image).astype(np.uint8)
image = (image / 127.5 - 1.0).astype(np.float32)
output2["image"] = image
output2["caption"] = self.style_word
batch = {}
batch["base"] = output
batch["style"] = output2
batch["font"] = fontname
batch["number"] = 0
batch["epochs"] = 800*2 if self.one_font else 1000*2
if self.full_word or self.just_use_style:
batch["cond"] = self.style_word
else:
batch["cond"] = self.style_word + " " + self.text
return batch
|