Skip to content

builder

data.dataloader.builder

SaigeDataLoader module-attribute 

SaigeDataLoader = Union[DataLoader, Full2PatchDataLoader]

Full2PatchDataLoader 

Full2PatchDataLoader(dataset: Dataset, collate_fn: Optional[Callable[[List[dict]], dict]] = None, batch_sampler: Optional[Iterable[int]] = None, batch_size: int = 1, shuffle: bool = False, drop_last: bool = True, **cfg_dataloader: dict)

DataLoader class for Segmentation and SceneTextRecognition tasks. Whole image is divided into multiple image patches, based on polygon labels. Full2PatchDataLoader has dataloader as class attribute, which returns croppped image patches. Image patches are stocked in buffer attribute, then passed to deep network in batch_size.

Attributes:

  • dataset (Dataset) –

    base dataset

  • batch_size (int) –

    batch size

  • drop_last (bool) –

    whether drop remaining data items less than batch_size at last iter.

  • buffer_multiple (int) –

    maximum size of buffer is decided by {buffer_multiple} x {batch_size}.

  • loader (DataLoader) –

    DataLoader that provides cropped image patches.

  • loader_iter (Iterable[list]) –

    iterator of loader attribute.

  • buffer (List[dict]) –

    buffer where image patches are stacking.

  • collate_function (Callable[[List[dict]], dict]) –

    function collating image patches into a batch.

initializing Full2PatchLoader

Parameters:

  • dataset (Dataset) –

    base dataset to be loaded.

  • collate_fn (Optional[Callable[[List[dict]], dict]], default: None ) –

    function collating image patches into a batch. Defaults to None.

  • sampler (Optional[Iterable[int]]) –

    specific sampler such as balanced sampler. Defaults to None.

  • batch_size (int, default: 1 ) –

    size of a batch for one iteration. Defaults to 1.

  • drop_last (bool, default: True ) –

    drop_last boolean. Defaults to True.

Source code in SaigeToolkit/data/dataloader/full2patch_loader.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
def __init__(
    self,
    dataset: data.Dataset,
    collate_fn: Optional[Callable[[List[dict]], dict]] = None,
    batch_sampler: Optional[Iterable[int]] = None,
    batch_size: int = 1,
    shuffle: bool = False,
    drop_last: bool = True,
    **cfg_dataloader: dict,
):
    """initializing Full2PatchLoader

    Args:
        dataset (data.Dataset): base dataset to be loaded.
        collate_fn (Optional[Callable[[List[dict]], dict]], optional):
            function collating image patches into a batch. Defaults to None.
        sampler (Optional[Iterable[int]], optional):
            specific sampler such as balanced sampler. Defaults to None.
        batch_size (int, optional): size of a batch for one iteration. Defaults to 1.
        drop_last (bool, optional): drop_last boolean. Defaults to True.
    """

    self.dataset = dataset
    self.batch_size = batch_size
    self.shuffle = shuffle
    self.drop_last = drop_last

    self.buffer_multiple = 10
    self.loader = data.DataLoader(
        dataset,
        collate_fn=self.collate_patch,
        batch_sampler=batch_sampler,
        **cfg_dataloader,
    )
    self.loader_iter = None
    self.buffer = []

    self.collate_function = collate_fn if collate_fn is not None else self.collate_batch

collate_patch 

collate_patch(batch: List[List[dict]]) -> List[dict]

dataset with crop_fn returns List[List[dict]] type. Since batch_size of self.loader is hard-defined as 1, simply returning first item of input is enough.

Parameters:

  • batch (List[List[dict]]) –

    batch (1) of cropped patches (n)

Returns:

  • List[dict] –

    List[dict]: cropped patches (n)

Source code in SaigeToolkit/data/dataloader/full2patch_loader.py 
68
69
70
71
72
73
74
75
76
77
78
79
def collate_patch(self, batch: List[List[dict]]) -> List[dict]:
    """dataset with crop_fn returns List[List[dict]] type.
    Since batch_size of self.loader is hard-defined as 1,
    simply returning first item of input is enough.

    Args:
        batch (List[List[dict]]): batch (1) of cropped patches (n)

    Returns:
        List[dict]: cropped patches (n)
    """
    return batch[0]

collate_batch 

collate_batch(batch: List[dict]) -> dict

collating image patches into a batch

Parameters:

  • batch (List[dict]) –

    list cropped patch data dicts (n)

Returns:

  • dict ( dict ) –

    a batch dict

Source code in SaigeToolkit/data/dataloader/full2patch_loader.py 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def collate_batch(self, batch: List[dict]) -> dict:
    """collating image patches into a batch

    Args:
        batch (List[dict]): list cropped patch data dicts (n)

    Returns:
        dict: a batch dict
    """

    data_out = {k: [] for k in batch[0].keys()}
    for data in batch:
        for key, value in data.items():
            if torch.is_tensor(value):
                value = value[None]
            data_out[key].append(value)

    for key, value in data_out.items():
        if torch.is_tensor(value[0]):
            try:
                data_out[key] = torch.cat(value)
            except:
                print(f'[{"WARNING".center(9)}] {key} is not concatenated')

    return data_out

__next__ 

__next__() -> dict

generating a batch data, extracting from buffer stack.

Raises:

  • StopIteration –

    self.get_buffer() called due to (len(self.buffer) < self.batch_size) and (len(self.buffer) == 0) even after self.get_buffer() : raise StopIteration.

  • StopIteration –

    self.get_buffer() called due to (len(self.buffer) < self.batch_size) and (len(self.buffer) < self.batch_size) even after self.get_buffer() and (drop_last == True) : raise StopIteration.

Returns:

  • dict ( dict ) –

    a batch dict

Source code in SaigeToolkit/data/dataloader/full2patch_loader.py 
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
def __next__(self) -> dict:
    """generating a batch data, extracting from buffer stack.

    Raises:
        StopIteration:
            self.get_buffer() called due to (len(self.buffer) < self.batch_size)
            and (len(self.buffer) == 0) even after self.get_buffer()
            : raise StopIteration.
        StopIteration:
            self.get_buffer() called due to (len(self.buffer) < self.batch_size)
            and (len(self.buffer) < self.batch_size) even after self.get_buffer()
            and (drop_last == True)
            : raise StopIteration.

    Returns:
        dict: a batch dict
    """

    batch_size = self.batch_size
    if len(self.buffer) < self.batch_size:
        self.get_buffer()

        if len(self.buffer) == 0:
            raise StopIteration

        if len(self.buffer) < self.batch_size:
            if self.drop_last:
                raise StopIteration
            batch_size = len(self.buffer)

    if self.shuffle:
        idx = np.sort(np.random.choice(range(len(self.buffer)), size=batch_size, replace=False))
    else:
        idx = np.array([0 for _ in range(batch_size)])

    batched_buffer = [self.buffer.pop(idx[-i - 1]) for i in range(len(idx))]

    return self.collate_function(batched_buffer)

get_dataloader_class 

get_dataloader_class(cfg_dataloader_name: Optional[str] = None, crop_fn: Optional[Callable[[dict], List[dict]]] = None) -> Type[SaigeDataLoader]

getting dataloader class

Parameters:

  • cfg_dataloader_name (Optional[str], default: None ) –

    name of dataloader class. If None, choose appropriate class considering crop_fn. Defaults to None.

  • crop_fn (Optional[Callable[[dict], List[dict]]], default: None ) –

    image crop function. Used for segmentation/OCR patch-wise training. Defaults to None.

Returns:

Source code in SaigeToolkit/data/dataloader/builder.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def get_dataloader_class(
    cfg_dataloader_name: Optional[str] = None, crop_fn: Optional[Callable[[dict], List[dict]]] = None
) -> Type[SaigeDataLoader]:
    """getting dataloader class

    Args:
        cfg_dataloader_name (Optional[str], optional): name of dataloader class.
            If None, choose appropriate class considering `crop_fn`. Defaults to None.
        crop_fn (Optional[Callable[[dict], List[dict]]], optional): image crop function.
            Used for segmentation/OCR patch-wise training. Defaults to None.

    Returns:
        Type[SaigeDataLoader]: dataloader class
    """
    if cfg_dataloader_name:
        try:
            return {
                "torch": data.DataLoader,
                "full2patch": Full2PatchDataLoader,
            }[cfg_dataloader_name]
        except:
            raise (f"Dataset {cfg_dataloader_name} not available")

    if crop_fn:
        return Full2PatchDataLoader

    # monkey patch
    def _stop(self):
        pass

    data.DataLoader.stop = _stop

    return data.DataLoader