| 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} |
| |
| 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]) |
