Hi, thank you for your U-net implementation. I am using your code to segment the aortic lumen from CTA images. I've tried both cross entropy and dice losses, with bad results. I think that the reason is that the classes are highly unbalanced. I have implemented my own loss function to separately handle the class errors in this way:

elif cost_name == "my_cost": eps = 1e-5

Define error in aorta identification

            prediction = pixel_wise_softmax(logits)
            prediction = prediction[..., 1]
            gt_map = self.y[...,1]                
            AortaError = tf.reduce_sum(gt_map - tf.math.multiply(prediction,gt_map)) / (tf.reduce_sum(gt_map) + eps)
            # Define background error
            BGError = tf.reduce_sum(tf.math.multiply(prediction, (1-gt_map))) / (tf.reduce_sum(1-gt_map) + eps)                
            # Use different weights                
            alpha = 0.5                 
            loss = (alpha*AortaError + (1 - alpha)*BGError)

I have used layers=4, features_root=64 and 400 iterations with batch_size = 6. After the first epoch I get this result, where the last colum show the output map with values > 0.7: epoch_0 The average loss after the first epoch is low (0.12). I think the network is segmenting the aorta without considering the contextual information but only the graylevels. Do you think the problem is in my loss function?

Thanks for your help,

Alice

jakeret / tf_unet

Modified Loss Function #233

Define error in aorta identification