File size: 1,192 Bytes
4acbfc7 | 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 | import torch
import pickle
import numpy as np
from PIL import Image
class UnifontModule(torch.nn.Module):
def __init__(self, out_dim, alphabet, device='cuda', input_type='unifont', linear=True):
super(UnifontModule, self).__init__()
self.device = device
self.alphabet = alphabet
self.symbols = self.get_symbols('unifont')
self.symbols_repr = self.get_symbols(input_type)
if linear:
self.linear = torch.nn.Linear(self.symbols_repr.shape[1], out_dim)
else:
self.linear = torch.nn.Identity()
def get_symbols(self, input_type):
with open(f"./File/{input_type}.pickle", "rb") as f:
symbols = pickle.load(f)
symbols = {sym['idx'][0]: sym['mat'].astype(np.float32).flatten() for sym in symbols}
# self.special_symbols = [self.symbols[ord(char)] for char in special_alphabet]
symbols = [symbols[ord(char)] for char in self.alphabet]
symbols.insert(0, np.zeros_like(symbols[0]))
symbols = np.stack(symbols)
return torch.from_numpy(symbols).float().to(self.device)
def forward(self, QR):
return self.linear(self.symbols_repr[QR]) |