qurator-spk/eynollah
1
0
Fork 0
mirror of https://github.com/qurator-spk/eynollah.git synced 2026-07-11 22:29:29 +02:00
Code Issues Projects Releases Wiki Activity

train: remove tf-specifics and weird __getitem__ def from transformer-ocr setup

Browse source
This commit is contained in:
kba 2026-07-09 19:23:21 +02:00
parent affddd6c85
commit 39e054e718
1 changed files with 5 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

21
src/eynollah/training/train.py
View file

@ -718,32 +718,21 @@ def run(_config,
"""
Wraps preprocess_imgs in a format consumable by torch
"""
def __init__(self, config, dir_img, dir_lab, char_to_num):
self.samples = list(
preprocess_imgs(
def __init__(self, config, dir_img, dir_lab):
self.samples = preprocess_imgs(
config,
dir_img,
dir_lab,
char_to_num=char_to_num,
processor=processor,
)
)
def __len__(self):
return len(self.samples)
def __getitem__(self, idx):
image, label = self.samples[idx]
def __iter__(self):
yield from self.samples
return {
"image": torch.as_tensor(image, dtype=torch.float32),
"label": torch.as_tensor(label, dtype=torch.long),
}
assert characters_txt_file
with open(characters_txt_file, 'r') as char_txt_f:
characters = json.load(char_txt_f)
char_to_num = StringLookup(vocabulary=list(characters), mask_token=None)
dataset = TransformerOCRTorchDataset(_config, dir_img, dir_lab, char_to_num)
dataset = TransformerOCRTorchDataset(_config, dir_img, dir_lab)
data_loader = torch.utils.data.DataLoader(dataset, batch_size=1)
train_dataset = data_loader.dataset