-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathmathematical_morphology.py
147 lines (116 loc) · 6.26 KB
/
mathematical_morphology.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
from scipy.ndimage import grey_erosion, grey_dilation, grey_closing, grey_opening
from parameters import *
from torch.nn import functional as f
from typing import Union
def erode(image: np.ndarray, strel: np.ndarray, origin: tuple = (0, 0)):
assert len(image.shape) == 2
assert len(strel.shape) == 2
assert len(origin)
assert origin[0] == origin[1]
assert 0 <= origin[0] < strel.shape[0]
origin = - strel.shape[0] // 2 + 1 + origin[0]
return grey_erosion(image, origin=origin, structure=strel, mode='constant')
def dilate(image, strel, origin):
assert len(image.shape) == 2
assert len(strel.shape) == 2
assert len(origin)
assert origin[0] == origin[1]
assert 0 <= origin[0] < strel.shape[0]
origin = - strel.shape[0] // 2 + 1 + origin[0]
return grey_dilation(image, origin=origin, structure=strel, mode='constant')
def closing(image, strel, origin=None):
# TODO: change origin from relative to absolute
if not origin:
return grey_closing(image, structure=strel, mode='constant')
else:
return grey_closing(image, origin=origin, structure=strel, mode='constant')
def opening(image, strel, origin=None):
# TODO: change origin from relative to absolute
if not origin:
return grey_opening(image, structure=strel, mode='constant')
else:
return grey_opening(image, origin=origin, structure=strel, mode='constant')
def erosion(image: torch.Tensor, strel: torch.Tensor, origin: tuple = (0, 0), border_value: Union[int, float] = 0):
image_pad = f.pad(image, [origin[0], strel.shape[0] - origin[0] - 1, origin[1], strel.shape[1] - origin[1] - 1],
mode='constant', value=border_value)
image_extended = f.unfold(image_pad.unsqueeze(0).unsqueeze(0), kernel_size=strel.shape)
strel_flatten = torch.flatten(strel).unsqueeze(0).unsqueeze(-1)
differences = image_extended - strel_flatten
result, _ = differences.min(dim=1)
return torch.reshape(result, image.shape)
def erosion_fast(image: torch.Tensor, strel: torch.Tensor, origin: tuple = (0, 0), border_value: Union[int, float] = 0):
result, _ = (f.unfold(f.pad(image, [origin[0], strel.shape[0] - origin[0] - 1, origin[1], strel.shape[1] - origin[1] - 1], mode='constant',
value=border_value).unsqueeze(0).unsqueeze(0), kernel_size=strel.shape)
- torch.flatten(strel).unsqueeze(0).unsqueeze(-1)).min(dim=1)
return torch.reshape(result, image.shape)
def dilation(image: torch.Tensor, strel: torch.Tensor, origin: tuple = (0, 0), border_value: Union[int, float] = 0):
image_pad = f.pad(image, [origin[0], strel.shape[0] - origin[0] - 1, origin[1], strel.shape[1] - origin[1] - 1],
mode='constant', value=border_value)
image_extended = f.unfold(image_pad.unsqueeze(0).unsqueeze(0), kernel_size=strel.shape)
strel_flatten = torch.flatten(strel).unsqueeze(0).unsqueeze(-1)
sums = image_extended + strel_flatten
result, _ = sums.max(dim=1)
return torch.reshape(result, image.shape)
if __name__ == '__main__':
library = 'cv2'
operation = 'dilation'
_image = np.zeros((5, 5))
_image[1, :] = 1
_image[:, 1] = 2
_image[3, :] = 3
_image[:, 3] = 3
plt.figure()
plt.imshow(_image, cmap='Greys', vmin=_image.min(), vmax=_image.max(), origin='lower')
_strel = np.zeros((3, 3)) - 1000
_strel[1, :] = 0
_strel[:, 1] = 0
_origin = (1, 1)
plt.figure()
plt.imshow(_strel, cmap='Greys', vmin=-100, vmax=0, origin='lower')
plt.scatter(_origin[0], _origin[1], marker='x', c='r')
if library == 'scipy':
if operation == 'erosion':
_erosion = erode(_image, _strel, origin=_origin)
plt.figure()
plt.imshow(_erosion, cmap='Greys', vmin=_image.min(), vmax=_image.max(), origin='lower')
elif operation == 'dilation':
_dilation = dilate(_image, _strel, origin=_origin)
plt.figure()
plt.imshow(_dilation, cmap='Greys', vmin=_image.min(), vmax=_image.max(), origin='lower')
if library == 'cv2':
from cv2 import dilate as dilate_cv2
from cv2 import erode as erode_cv2
from cv2 import getStructuringElement
if operation == 'erosion':
_erosion = erode_cv2(_image.astype(np.int8), np.exp(_strel).astype(np.int8))
plt.figure()
plt.imshow(_erosion, cmap='Greys', vmin=_image.min(), vmax=_image.max(), origin='lower')
elif operation == 'dilation':
k_size = _strel.shape[0]
_dilation = dilate_cv2(_image.astype(np.int8), getStructuringElement(0, k_size))
plt.figure()
plt.imshow(_dilation, cmap='Greys', vmin=_image.min(), vmax=_image.max(), origin='lower')
elif library == 'torch':
_image_tensor = torch.tensor(_image, device=DEVICE)
_strel_tensor = torch.tensor(_strel, device=DEVICE)
if operation == 'erosion':
_erosion_tensor = erosion(_image_tensor, _strel_tensor, origin=_origin)
plt.figure()
plt.imshow(_erosion_tensor.cpu().numpy(), cmap='Greys', vmin=_image.min(), vmax=_image.max(), origin='lower')
elif operation == 'dilation':
_dilation_tensor = dilation(_image_tensor, _strel_tensor, origin=_origin, border_value=-1000)
plt.figure()
plt.imshow(_dilation_tensor.cpu().numpy(), cmap='Greys', vmin=_image.min(), vmax=_image.max(), origin='lower')
elif library == 'kornia':
from kornia.morphology import erosion as erosion_kornia
from kornia.morphology import dilation as dilation_kornia
_image_tensor = torch.tensor(_image, device=DEVICE)
_strel_tensor = torch.tensor(_strel, device=DEVICE)
if operation == 'erosion':
_erosion_tensor = erosion_kornia(_image_tensor.unsqueeze(0).unsqueeze(0), _strel_tensor)
plt.figure()
plt.imshow(_erosion_tensor.cpu().numpy()[0, 0, :, :], cmap='Greys', vmin=_image.min(), vmax=_image.max(), origin='lower')
elif operation == 'dilation':
_dilation_tensor = dilation_kornia(_image_tensor.unsqueeze(0).unsqueeze(0), _strel_tensor)
plt.figure()
plt.imshow(_dilation_tensor.cpu().numpy()[0, 0, :, :], cmap='Greys', vmin=_image.min(), vmax=_image.max(), origin='lower')