-
Notifications
You must be signed in to change notification settings - Fork 1.5k
Expand file tree
/
Copy pathinverse_batch_transform.py
More file actions
162 lines (137 loc) · 6.89 KB
/
inverse_batch_transform.py
File metadata and controls
162 lines (137 loc) · 6.89 KB
1
2
3
4
5
6
7
8
9
10
11
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
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
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
106
107
108
109
110
111
112
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
151
152
153
154
155
156
157
158
159
160
161
162
# Copyright (c) MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import warnings
from collections.abc import Callable, Sequence
from typing import Any
from torch.utils.data import Dataset
from torch.utils.data.dataloader import DataLoader as TorchDataLoader
from monai.config import KeysCollection
from monai.data.dataloader import DataLoader
from monai.data.utils import decollate_batch, no_collation, pad_list_data_collate
from monai.transforms.croppad.batch import PadListDataCollate
from monai.transforms.inverse import InvertibleTransform
from monai.transforms.transform import MapTransform, Transform
from monai.utils import first
__all__ = ["BatchInverseTransform", "Decollated", "DecollateD", "DecollateDict"]
class _BatchInverseDataset(Dataset):
def __init__(self, data: Sequence[Any], transform: InvertibleTransform, pad_collation_used: bool) -> None:
self.data = data
self.invertible_transform = transform
self.pad_collation_used = pad_collation_used
def __getitem__(self, index: int):
data = dict(self.data[index])
# If pad collation was used, then we need to undo this first
if self.pad_collation_used:
data = PadListDataCollate.inverse(data)
if not isinstance(self.invertible_transform, InvertibleTransform):
warnings.warn("transform is not invertible, can't invert transform for the input data.")
return data
return self.invertible_transform.inverse(data)
def __len__(self) -> int:
return len(self.data)
class BatchInverseTransform(Transform):
"""
Perform inverse on a batch of data. This is useful if you have inferred a batch of images and want to invert
them all.
"""
def __init__(
self,
transform: InvertibleTransform,
loader: TorchDataLoader,
collate_fn: Callable | None = no_collation,
num_workers: int | None = 0,
detach: bool = True,
pad_batch: bool = True,
fill_value=None,
) -> None:
"""
Args:
transform: a callable data transform on input data.
loader: data loader used to run `transforms` and generate the batch of data.
collate_fn: how to collate data after inverse transformations.
default won't do any collation, so the output will be a list of size batch size.
num_workers: number of workers when run data loader for inverse transforms,
default to 0 as only run 1 iteration and multi-processing may be even slower.
if the transforms are really slow, set num_workers for multi-processing.
if set to `None`, use the `num_workers` of the transform data loader.
detach: whether to detach the tensors. Scalars tensors will be detached into number types
instead of torch tensors.
pad_batch: when the items in a batch indicate different batch size,
whether to pad all the sequences to the longest.
If False, the batch size will be the length of the shortest sequence.
fill_value: the value to fill the padded sequences when `pad_batch=True`.
"""
self.transform = transform
self.batch_size = loader.batch_size
self.num_workers = loader.num_workers if num_workers is None else num_workers
self.collate_fn = collate_fn
self.detach = detach
self.pad_batch = pad_batch
self.fill_value = fill_value
self.pad_collation_used = loader.collate_fn.__doc__ == pad_list_data_collate.__doc__ or isinstance(
loader.collate_fn, PadListDataCollate
)
def __call__(self, data: dict[str, Any]) -> Any:
decollated_data = decollate_batch(data, detach=self.detach, pad=self.pad_batch, fill_value=self.fill_value)
inv_ds = _BatchInverseDataset(decollated_data, self.transform, self.pad_collation_used)
inv_loader = DataLoader(
inv_ds, batch_size=self.batch_size, num_workers=self.num_workers, collate_fn=self.collate_fn
)
try:
return first(inv_loader)
except RuntimeError as re:
re_str = str(re)
if "equal size" in re_str:
re_str += "\nMONAI hint: try creating `BatchInverseTransform` with `collate_fn=lambda x: x`."
raise RuntimeError(re_str) from re
class Decollated(MapTransform):
"""
Decollate a batch of data. If input is a dictionary, it also supports to only decollate specified keys.
Note that unlike most MapTransforms, it will delete the other keys that are not specified.
if `keys=None`, it will decollate all the data in the input.
It replicates the scalar values to every item of the decollated list.
Args:
keys: keys of the corresponding items to decollate, note that it will delete other keys not specified.
if None, will decollate all the keys. see also: :py:class:`monai.transforms.compose.MapTransform`.
detach: whether to detach the tensors. Scalars tensors will be detached into number types
instead of torch tensors.
pad_batch: when the items in a batch indicate different batch size,
whether to pad all the sequences to the longest.
If False, the batch size will be the length of the shortest sequence.
fill_value: the value to fill the padded sequences when `pad_batch=True`.
allow_missing_keys: don't raise exception if key is missing.
"""
def __init__(
self,
keys: KeysCollection | None = None,
detach: bool = True,
pad_batch: bool = True,
fill_value=None,
allow_missing_keys: bool = False,
) -> None:
super().__init__(keys, allow_missing_keys)
self.detach = detach
self.pad_batch = pad_batch
self.fill_value = fill_value
def __call__(self, data: dict | list):
d: dict | list
if len(self.keys) == 1 and self.keys[0] is None:
# it doesn't support `None` as the key
d = data
else:
if not isinstance(data, dict):
raise TypeError("input data is not a dictionary, but specified keys to decollate.")
d = {}
for key in self.key_iterator(data):
d[key] = data[key]
return decollate_batch(d, detach=self.detach, pad=self.pad_batch, fill_value=self.fill_value)
DecollateD = DecollateDict = Decollated