From 39e054e718f91c2dfa44bc626119480bf975917c Mon Sep 17 00:00:00 2001 From: kba Date: Thu, 9 Jul 2026 19:23:21 +0200 Subject: [PATCH] train: remove tf-specifics and weird __getitem__ def from transformer-ocr setup --- src/eynollah/training/train.py | 21 +++++---------------- 1 file changed, 5 insertions(+), 16 deletions(-) diff --git a/src/eynollah/training/train.py b/src/eynollah/training/train.py index 304dc87..a6b5331 100644 --- a/src/eynollah/training/train.py +++ b/src/eynollah/training/train.py @@ -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