qurator-spk/eynollah
1
0
Fork 0
mirror of https://github.com/qurator-spk/eynollah.git synced 2026-07-12 06:39:31 +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 Wraps preprocess_imgs in a format consumable by torch
""" """
def __init__(self, config, dir_img, dir_lab, char_to_num): def __init__(self, config, dir_img, dir_lab):
self.samples = list( self.samples = preprocess_imgs(
preprocess_imgs(
config, config,
dir_img, dir_img,
dir_lab, dir_lab,
char_to_num=char_to_num,
processor=processor, processor=processor,
) )
)
def __len__(self): def __len__(self):
return len(self.samples) return len(self.samples)
def __getitem__(self, idx): def __iter__(self):
image, label = self.samples[idx] yield from self.samples
return { dataset = TransformerOCRTorchDataset(_config, dir_img, dir_lab)
"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)
data_loader = torch.utils.data.DataLoader(dataset, batch_size=1) data_loader = torch.utils.data.DataLoader(dataset, batch_size=1)
train_dataset = data_loader.dataset train_dataset = data_loader.dataset