Skip to content

Robust T-Loss for Medical Image Segmentation with TensorFlow backend

License

Notifications You must be signed in to change notification settings

andreped/t-loss-tf

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

24 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

t-loss-tf

Robust T-Loss for Medical Image Segmentation with TensorFlow backend

GitHub Downloads License CI Code style: black

This T-loss implementation is an adaption of the original PyTorch code Digital-Dermatology/t-loss-loss.

More information about T-loss and the original paper can be found here.

pip install git+https://github.com/andreped/t-loss-tf.git

As the t-loss contains a trainable parameter, in keras the loss needed to be implemented as a custom layer. Hence, instead of setting the loss as normally through model.compile(loss=[...]), just add it to the model at an appropriate place (e.g., at the end of the network).

An example can be seen below:

import tensorflow as tf
from t_loss import TLoss

# create dummy inputs and GTs
input_shape = (16, 16, 1)
x = tf.ones((32,) + input_shape, dtype="float32")
y = tf.ones((32,) + input_shape, dtype="float32")

# define network
input_x = tf.keras.Input(shape=input_shape)
input_y = tf.keras.Input(shape=input_shape)
z = tf.keras.layers.Conv2D(filters=4, kernel_size=(1, 1), activation="relu")(input_x)
z = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(z)
z = tf.keras.layers.UpSampling2D(size=(2, 2))(z)
z = tf.keras.layers.Conv2D(filters=1, kernel_size=(1, 1), activation="sigmoid")(z)
z = TLoss(tensor_shape=input_shape, image_size=input_shape[0])(z, input_y)
model = tf.keras.Model(inputs=[input_x, input_y], outputs=[z])

# train model
model.compile(optimizer="adam")
model.fit(x=[x, y], y=y, batch_size=2, epochs=1, verbose="auto")

The code in this repository is released under MIT License.

The implementation is based on the original torch implementation hosted here.

Hence, if this code is used, please cite the original research article:

@inproceedings{gonzalezjimenezRobustTLoss2023,
  title     = {Robust T-Loss for Medical Image Segmentation},
  author    = {Gonzalez-Jimenez, Alvaro and Lionetti, Simone and Gottfrois, Philippe and Gröger, Fabian and Pouly, Marc and Navarini, Alexander},
  journal   = {Medical {{Image Computing}} and {{Computer Assisted Intervention}} – {{MICCAI}} 2023},
  publisher = {{Springer International Publishing}},
  year      = {2023},
}