forked from Project-MONAI/MONAI
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathpadders.py
178 lines (165 loc) · 8.69 KB
/
padders.py
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
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
# 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 unittest
import numpy as np
import torch
from monai.data.meta_tensor import MetaTensor
from monai.transforms import Compose
from monai.transforms.lazy.functional import apply_pending
from monai.transforms.transform import MapTransform
from monai.utils.enums import NumpyPadMode, PytorchPadMode
from tests.test_utils import TEST_NDARRAYS_ALL, assert_allclose
MODES = []
# Test modes
NP_MODES: list = [
"constant",
"edge",
# `reflect` mode is not supported in some PyTorch versions, skip the test
# "reflect",
"wrap",
"median",
"mean",
]
MODES += NP_MODES
MODES += [NumpyPadMode(i) for i in NP_MODES]
PT_MODES: list = [
"constant",
"replicate",
"circular",
# `reflect` mode is not supported in some PyTorch versions, skip the test
# "reflect",
]
MODES += PT_MODES
MODES += [PytorchPadMode(i) for i in PT_MODES]
TESTS_PENDING_MODE = [["constant", "zeros"], ["edge", "border"]]
class PadTest(unittest.TestCase):
@staticmethod
def get_arr(shape):
return np.random.randint(100, size=shape).astype(float)
def pad_test(self, input_param, input_shape, expected_shape, modes=None):
# loop over each mode
for mode in modes or MODES:
with self.subTest(mode=mode):
base_comparison = None
im = self.get_arr(input_shape)
padder = self.Padder(mode=mode, **input_param)
is_map = isinstance(padder, MapTransform)
# check result is the same regardless of input type
for im_type in TEST_NDARRAYS_ALL:
with self.subTest(im_type=im_type):
input_image = im_type(im)
input_data = {"img": im_type(im)} if is_map else im_type(im)
# our array transforms can also take `mode` as an argument to `__call__`
# Check this gives equivalent results
for call_extra_args in [{}] if is_map else [{}, {"mode": mode}]:
with self.subTest(call_extra_args=call_extra_args):
r_out = padder(input_data, **call_extra_args)
r_im = r_out["img"] if is_map else r_out
# check shape, type, etc.
np.testing.assert_allclose(r_im.shape, expected_shape)
self.assertIsInstance(r_im, MetaTensor)
self.assertEqual(len(r_im.applied_operations), 1)
# check results are same regardless of input type
if base_comparison is None:
base_comparison = r_im
else:
assert_allclose(r_im, base_comparison)
# test inverse
if isinstance(r_im, MetaTensor):
r_out = padder.inverse(r_out)
r_im = r_out["img"] if is_map else r_out
self.assertIsInstance(r_im, MetaTensor)
assert_allclose(r_im, input_image, type_test=False)
self.assertEqual(r_im.applied_operations, [])
def pad_test_kwargs(self, unchanged_slices, **input_param):
for im_type in TEST_NDARRAYS_ALL:
with self.subTest(im_type=im_type):
for kwargs in ({"value": 2}, {"constant_values": ((0, 0), (1, 1), (2, 2))}):
with self.subTest(kwargs=kwargs):
im = im_type(np.random.randint(-100, -10, size=(3, 8, 4)))
padder = self.Padder(**input_param, **kwargs)
result = padder(im)
if isinstance(result, torch.Tensor):
result = result.cpu()
assert_allclose(result[unchanged_slices], im, type_test=False)
# we should have the same as the input plus some 2s (if value) or 1s and 2s (if constant_values)
if isinstance(im, torch.Tensor):
im = im.detach().cpu().numpy()
expected_vals = np.unique(im).tolist()
expected_vals += [2] if "value" in kwargs else [1, 2]
assert_allclose(np.unique(result), expected_vals, type_test=False)
# check inverse
if isinstance(result, MetaTensor):
inv = padder.inverse(result)
assert_allclose(im, inv, type_test=False)
self.assertEqual(inv.applied_operations, [])
def pad_test_pending_ops(self, input_param, input_shape):
for mode in TESTS_PENDING_MODE:
# TODO: One of the dim in the input data contains 1 report error.
pad_fn = self.Padder(mode=mode[0], **input_param)
data = self.get_arr(input_shape)
is_map = isinstance(pad_fn, MapTransform)
im = MetaTensor(data, meta={"a": "b", "affine": np.eye(len(input_shape))})
input_data = {"img": im} if is_map else im
# non-lazy
result_non_lazy = pad_fn(input_data)
expected = result_non_lazy["img"] if is_map else result_non_lazy
self.assertIsInstance(expected, MetaTensor)
# lazy
pad_fn.lazy = True
pending_result = pad_fn(input_data)
pending_result = pending_result["img"] if is_map else pending_result
self.assertIsInstance(pending_result, MetaTensor)
assert_allclose(pending_result.peek_pending_affine(), expected.affine)
assert_allclose(pending_result.peek_pending_shape(), expected.shape[1:])
# TODO: mode="bilinear" may report error
overrides = {"mode": "nearest", "padding_mode": mode[1], "align_corners": False}
result = apply_pending(pending_result, overrides=overrides)[0]
# lazy in constructor
pad_fn_lazy = self.Padder(mode=mode[0], lazy=True, **input_param)
self.assertTrue(pad_fn_lazy.lazy)
# compare
assert_allclose(result, expected, rtol=1e-5)
if isinstance(result, MetaTensor) and not isinstance(pad_fn, MapTransform):
pad_fn.lazy = False
inverted = pad_fn.inverse(result)
self.assertTrue((not inverted.pending_operations) and (not inverted.applied_operations))
self.assertEqual(inverted.shape, im.shape)
def pad_test_combine_ops(self, funcs, input_shape, expected_shape):
for mode in TESTS_PENDING_MODE:
# non-lazy
_funcs = []
for func in funcs:
for _func, _params in func.items():
_funcs.append(_func(mode=mode[0], **_params))
trans = Compose(_funcs)
data = self.get_arr(input_shape)
is_map = isinstance(_funcs[0], MapTransform)
im = MetaTensor(data, meta={"a": "b", "affine": np.eye(len(input_shape))})
input_data = {"img": im} if is_map else im
result_non_lazy = trans(input_data)
expected = result_non_lazy["img"] if is_map else result_non_lazy
self.assertIsInstance(expected, MetaTensor)
# lazy
pending_result = input_data
for _func in _funcs:
_func.lazy = True
pending_result = _func(pending_result)
pending_result = pending_result["img"] if is_map else pending_result
self.assertIsInstance(pending_result, MetaTensor)
assert_allclose(pending_result.peek_pending_affine(), expected.affine)
assert_allclose(pending_result.peek_pending_shape(), expected.shape[1:])
# TODO: mode="bilinear" may report error
overrides = {"mode": "nearest", "padding_mode": mode[1], "align_corners": False}
result = apply_pending(pending_result, overrides=overrides)[0]
# compare
assert_allclose(result, expected, rtol=1e-5)