| |
| from __future__ import division |
| from __future__ import print_function |
|
|
| import os |
| import gc |
| import numpy as np |
| import cv2 |
| from collections import namedtuple |
| from datasets import load_dataset |
|
|
| from SamplePreprocessor_Hindi_v2 import preprocess |
|
|
| Batch = namedtuple("Batch", ["imgs", "gtTexts"]) |
|
|
|
|
| class DataLoader: |
|
|
| def __init__(self, |
| dataset_name: str, |
| images_dir: str, |
| batchSize: int, |
| imgSize: tuple, |
| maxTextLen: int, |
| charList: str, |
| dataAugmentation: bool = False, |
| val_split: float = 0.1): |
|
|
| self.batchSize = batchSize |
| self.imgSize = imgSize |
| self.maxTextLen = maxTextLen |
| self.charList = charList |
| self.dataAugmentation = dataAugmentation |
| self.images_dir = images_dir |
| self.charSet = set(charList) |
|
|
| print(f"[DataLoader] Loading dataset: {dataset_name}") |
|
|
| dataset = load_dataset(dataset_name, split='train') |
| dataset = dataset.shuffle(seed=42) |
|
|
| split_idx = int(len(dataset) * (1 - val_split)) |
|
|
| self.train_dataset = dataset.select(range(split_idx)) |
| self.val_dataset = dataset.select(range(split_idx, len(dataset))) |
|
|
| print(f"[DataLoader] Train: {len(self.train_dataset)}") |
| print(f"[DataLoader] Val: {len(self.val_dataset)}") |
|
|
| self._current_dataset = None |
| self._idx = 0 |
|
|
| def _load_and_preprocess_image(self, img_path, augment): |
| img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) |
| if img is None: |
| img = np.zeros([self.imgSize[1], self.imgSize[0]], dtype=np.uint8) |
| return preprocess(img, self.imgSize, dataAugmentation=augment) |
|
|
| def trainSet(self, subset_size=None): |
| dataset = self.train_dataset.shuffle(seed=None) |
|
|
| if subset_size is not None: |
| subset_size = min(subset_size, len(dataset)) |
| indices = np.random.choice(len(dataset), subset_size, replace=False) |
| dataset = dataset.select(indices) |
| print(f"[DataLoader] Training on {subset_size} samples") |
| else: |
| print(f"[DataLoader] Training on FULL dataset") |
|
|
| self._current_dataset = dataset |
| self._idx = 0 |
|
|
| def validationSet(self): |
| self._current_dataset = self.val_dataset |
| self._idx = 0 |
|
|
| def hasNext(self): |
| return self._idx < len(self._current_dataset) |
|
|
| def getIteratorInfo(self): |
| """Return current batch index and total batches.""" |
| total = max(1, (len(self._current_dataset) + self.batchSize - 1) // self.batchSize) |
| current = self._idx // self.batchSize + 1 |
| return (current, total) |
| |
| def getNext(self): |
| end = min(self._idx + self.batchSize, len(self._current_dataset)) |
| batch_indices = range(self._idx, end) |
| self._idx = end |
|
|
| imgs = [] |
| gtTexts = [] |
|
|
| for idx in batch_indices: |
| sample = self._current_dataset[idx] |
|
|
| augment = (self._current_dataset != self.val_dataset) |
|
|
| img = np.array(sample["image"].convert("L")) |
| img = preprocess(img, self.imgSize, dataAugmentation=augment) |
|
|
| imgs.append(img) |
| gtTexts.append(sample['text']) |
| return Batch(imgs=imgs, gtTexts=gtTexts) |
|
|
| def cleanup(self): |
| gc.collect() |
|
|
