Feeding weight maps into a CNN (UNET network) in keras - tensorflow

Question

I have implemented UNET network described in here
The network is working fine, but in the paper, they have mentioned adding weighted maps into the network for better boundary separation. The weight maps are calculated this way as far as I understand
def unet_weight_map(y, wc=None, w0 = 10, sigma = 5):
"""
Parameters
----------
mask: Numpy array
2D array of shape (image_height, image_width) representing binary mask
of objects.
wc: dict
Dictionary of weight classes.
w0: int
Border weight parameter.
sigma: int
Border width parameter.
Returns
-------
Numpy array
Training weights. A 2D array of shape (image_height, image_width).
"""
labels = label(y)
no_labels = labels == 0
label_ids = sorted(np.unique(labels))[1:]
if len(label_ids) > 1:
distances = np.zeros((y.shape[0], y.shape[1], len(label_ids)))
for i, label_id in enumerate(label_ids):
distances[:,:,i] = distance_transform_edt(labels != label_id)
distances = np.sort(distances, axis=2)
d1 = distances[:,:,0]
d2 = distances[:,:,1]
w = w0 * np.exp(-1/2*((d1 + d2) / sigma)**2) * no_labels
else:
w = np.zeros_like(y)
if wc:
class_weights = np.zeros_like(y)
for k, v in wc.items():
class_weights[y == k] = v
w = w + class_weights
return w
Until here everything is fine. But, my question is that how I can get use of these weight maps in the network. I have a weighted binary cross entropy loss defined as below
def weighted_binary_crossentropy( y_true, y_pred, weight=[1.,2.]):
y_true = K.clip(y_true, K.epsilon(), 1-K.epsilon())
y_pred = K.clip(y_pred, K.epsilon(), 1-K.epsilon())
logloss = -(y_true * K.log(y_pred) * weight[0] + (1 - y_true) * K.log(1 - y_pred)*weight[1])
return K.mean( logloss, axis=-1)
But, here I give the weights as a [a, b] array into the loss for class weights and then feed this loss to the network when compiling. My question is that should I feed those maps into this customized loss function? if so, how? if not, what other way can I use in Keras? Please help. I have read many stack overflow questions related to this problem, but I could not get an answer. I can provide any information regarding my network if needed.

Answer 1

In order to pass your own parameters to a custom loss function, you have 2 ways. You should either subclass loss, or use a wrapper function.
For example you can set a wrapper function like this:
def wrapper_loss(weights=[1.,2.]):
def weighted_binary_crossentropy(y_true, y_pred):
y_true = K.clip(y_true, K.epsilon(), 1-K.epsilon())
y_pred = K.clip(y_pred, K.epsilon(), 1-K.epsilon())
logloss = -(y_true * K.log(y_pred) * weight[0] + (1 - y_true) * K.log(1 - y_pred)*weight[1])
return K.mean(logloss, axis=-1)
return weighted_binary_crossentropy
Then, pass it to the model.compile() like this:
model.compile(loss=wrapper_loss(weights=[1.,2.]), optimizer=...)
P.S: You may need to check these out:
tf.nn.weighted_cross_entropy_with_logits
class_weight argument for model.fit()

Answer 2

I realized how to use those maps. First I define an Input (with the same shape as ground truth labels) as the way we define Input when feeding the input images. Something like
weights = Input(shape=(shape_of_groundtruth_labels))
I define the customized loss with the same structure as wrapper_loss defined above; with weight maps this time, not class weights [1, 2]. Then, when defining the model which needs input and output. I give the input as both input images and input weights. something like:
model = Model(inputs=[images, weights], outputs=...)
where weights are the one I defined in the Input layer. In the model.compile(), I give the loss as the name of my customized loss (wrapper_loss) with the inputs weights. something like
model.compile(optimizer=..., loss=wrapper_loss(weight = weights), ...)
where the second 'weights' is the one defined in Input layer.
Now, the last thing to do is to do the same in model.fit; I give the weight maps with the images with the same structure as above.