-
Notifications
You must be signed in to change notification settings - Fork 0
/
soft_skeleton.py
37 lines (29 loc) · 1.04 KB
/
soft_skeleton.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
import torch
import torch.nn as nn
import torch.nn.functional as F
def soft_erode(img):
if len(img.shape)==4:
p1 = -F.max_pool2d(-img, (3,1), (1,1), (1,0))
p2 = -F.max_pool2d(-img, (1,3), (1,1), (0,1))
return torch.min(p1,p2)
elif len(img.shape)==5:
p1 = -F.max_pool3d(-img,(3,1,1),(1,1,1),(1,0,0))
p2 = -F.max_pool3d(-img,(1,3,1),(1,1,1),(0,1,0))
p3 = -F.max_pool3d(-img,(1,1,3),(1,1,1),(0,0,1))
return torch.min(torch.min(p1, p2), p3)
def soft_dilate(img):
if len(img.shape)==4:
return F.max_pool2d(img, (3,3), (1,1), (1,1))
elif len(img.shape)==5:
return F.max_pool3d(img,(3,3,3),(1,1,1),(1,1,1))
def soft_open(img):
return soft_dilate(soft_erode(img))
def soft_skel(img, iter_):
img1 = soft_open(img)
skel = F.relu(img-img1)
for j in range(iter_):
img = soft_erode(img)
img1 = soft_open(img)
delta = F.relu(img-img1)
skel = skel + F.relu(delta-skel*delta)
return skel