-
Notifications
You must be signed in to change notification settings - Fork 0
/
pyLayer.py
executable file
·158 lines (112 loc) · 6.21 KB
/
pyLayer.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
import caffe
import numpy as np
class DiceLoss(caffe.Layer):
"""
Compute energy based on dice coefficient.
"""
union = None
intersection = None
result = None
gt = None
#Dice = [0,0]
def setup(self, bottom, top):
# check input pair
if len(bottom) != 2:
raise Exception("Need two inputs to compute the dice. the result of the softmax and the ground truth.")
def reshape(self, bottom, top):
# check input dimensions match
# if bottom[0].count != 2*bottom[1].count:
# print bottom[0].data.shape
# print bottom[1].data.shape
# raise Exception("the dimension of inputs should match")
# loss output is two scalars (mean and std)
top[0].reshape(1)
def forward(self, bottom, top):
dice = np.zeros(bottom[0].data.shape[0],dtype=np.float32)
self.union = np.zeros(bottom[0].data.shape[0],dtype=np.float32)
self.intersection = np.zeros(bottom[0].data.shape[0],dtype=np.float32)
#self.result = np.reshape(np.squeeze(np.argmax(bottom[0].data[...],axis=1)),[bottom[0].data.shape[0],bottom[0].data.shape[2]])
# V
self.result = np.reshape(np.squeeze(bottom[0].data[:,0,:]),[bottom[0].data.shape[0],bottom[0].data.shape[2]])
self.gt = np.reshape(np.squeeze(bottom[1].data[...]),[bottom[1].data.shape[0],bottom[1].data.shape[2]])
self.gt = (self.gt > 0.5).astype(dtype=np.float32)
self.result = self.result.astype(dtype=np.float32)
self.result[self.result[:]>=0.5]=1
self.result[self.result[:]<0.5]=0
for i in range(0,bottom[0].data.shape[0]):
# compute dice
CurrResult = (self.result[i,:]).astype(dtype=np.float32)
CurrGT = (self.gt[i,:]).astype(dtype=np.float32)
self.union[i]=(np.sum(CurrResult) + np.sum(CurrGT))
self.intersection[i]=(np.sum(CurrResult * CurrGT))
dice[i] = 1- (2 * self.intersection[i] / (self.union[i]+0.00001))
print "Dice Loss: " + str(dice[i])
#self.Dice[i]=dice[i]
top[0].data[0]=np.sum(dice)
def backward(self, top, propagate_down, bottom):
for btm in [0]:
prob = bottom[0].data[...]
bottom[btm].diff[...] = np.zeros(bottom[btm].diff.shape, dtype=np.float32)
for i in range(0, bottom[btm].diff.shape[0]):
#self.union[i] = np.sum(prob[i,1,:]) + np.sum(self.gt[i,:])
#self.intersection[i] = np.sum(prob[i,1,:] * self.gt[i,:])
#if self.Dice[0] == 1 & self.Dice[1] == 1:
# bottom[btm].diff[i, 0, :] += 0.5
# bottom[btm].diff[i, 1, :] -= 0.5
#else: # V # |
bottom[btm].diff[i, 0, :] -= 2.0 * ( # V
(self.gt[i, :] * self.union[i]) / ((self.union[i]) ** 2) - 2.0*prob[i,0,:]*(self.intersection[i]) / (self.union[i]) ** 2)
if bottom[btm].diff.shape[1]==2:
bottom[btm].diff[i, 1, :] += 2.0 * (
(self.gt[i, :] * self.union[i]) / ((self.union[i]) ** 2) - 2.0*prob[i,0,:]*(self.intersection[i]) / (self.union[i]) ** 2)
class DistLoss(caffe.Layer):
"""
Compute loss based on Distance-Maps.
"""
alpha = 0
beta = 0
dmap_o = None
dmap_i = None
result_o = None
result_i = None
union = None
def setup(self, bottom, top):
# check input pair
if len(bottom) != 2:
raise Exception("Need two inputs to compute the loss. the result of the softmax and the distance map.")
def reshape(self, bottom, top):
# check input dimensions match
if bottom[0].count != 2*bottom[1].count:
print bottom[0].data.shape
print bottom[1].data.shape
raise Exception("the dimension of inputs should match")
# loss output is two scalars (mean and std)
top[0].reshape(1)
def forward(self, bottom, top):
dist = np.zeros(bottom[0].data.shape[0],dtype=np.float32)
self.union = np.zeros(bottom[0].data.shape[0],dtype=np.float32)
self.result_i = np.reshape(np.squeeze(bottom[0].data[:,0,:]), [bottom[0].data.shape[0], bottom[0].data.shape[2]]) #inside labelmap
self.result_o = np.reshape(np.squeeze(bottom[0].data[:,1,:]), [bottom[0].data.shape[0], bottom[0].data.shape[2]]) #outside labelmap
self.result_i[self.result_i < 0.5] = 0; self.result_i[self.result_i >= 0.5] = 1
self.result_o[self.result_o < 0.5] = 0; self.result_o[self.result_o >= 0.5] = 1
gt = np.reshape(np.squeeze(bottom[1].data[...]), [bottom[1].data.shape[0], bottom[1].data.shape[2]])
self.dmap_i = gt.astype(dtype=np.float32)
self.dmap_o = gt.astype(dtype=np.float32)
self.result_i = self.result_i.astype(dtype=np.float32)
self.result_o = self.result_o.astype(dtype=np.float32)
self.dmap_o[self.dmap_o < 0] = -self.alpha
self.dmap_i[self.dmap_i > 0] = self.beta
for i in range(0,bottom[0].data.shape[0]):
# compute distance loss
self.union[i] = (np.sum(self.result_i[i]) + np.sum(self.result_o[i])) #bottom[0].data.shape[0]*bottom[0].data.shape[1]*bottom[0].data.shape[2]
dist[i] = (self.dmap_o[i].dot(self.result_i[i]) - self.dmap_i[i].dot(self.result_o[i])) / (self.union[i])
print "Distance Loss: " + str(dist[i])
top[0].data[0]=np.sum(dist)
def backward(self, top, propagate_down, bottom):
btm = 0
prob_i = np.squeeze(bottom[0].data[:, 0, :])
prob_o = np.squeeze(bottom[0].data[:, 1, :])
bottom[btm].diff[...] = np.zeros(bottom[btm].diff.shape, dtype=np.float32)
for i in range(0, bottom[btm].diff.shape[0]):
bottom[btm].diff[i, 0, :] = (self.dmap_o[i] * self.union[i] - 2 * prob_i[i] *(self.dmap_o[i].dot(self.result_i[i]) - self.dmap_i[i].dot(self.result_o[i]))) / (self.union[i] ** 2)
bottom[btm].diff[i, 1, :] = (-self.dmap_i[i] * self.union[i] - 2 * prob_o[i] * (self.dmap_o[i].dot(self.result_i[i]) - self.dmap_i[i].dot(self.result_o[i]))) / (self.union[i] ** 2)