sksurgerytf.models.rgb_unet module

Module to implement a semantic (pixelwise) segmentation using UNet on 512x512.

class sksurgerytf.models.rgb_unet.RGBUNet(logs='logs/fit', data=None, working=None, omit=None, model=None, learning_rate=0.0001, epochs=50, batch_size=2, input_size=(512, 512, 3), patience=20)

Bases: object

Class to encapsulate RGB UNet semantic (pixelwise) segmentation network.

Thanks to Zhixuhao, and ShawDa for getting me started, and `Harshall Lamba <https://towardsdatascience.com/understanding-semantic-segmentation-with-unet-6be4f42d4b47>_, for further inspiration.

predict(rgb_image)

Method to test a single image. Image resized to match network, segmented and then resized back to match the input size.

Parameters:rgb_image – 3 channel RGB, [0-255], uchar. Returns:single channel, [0=bg|255=fg].

save_model(filename)

Method to save the whole trained network to disk.

Parameters:filename – file to save to.

train()

Method to train the neural network. Writes each epoch to tensorboard log files.

Returns:output of self.model.evaluate on validation set, or None.

sksurgerytf.models.rgb_unet.run_rgb_unet_model(logs, data, working, omit, model, save, test, prediction, epochs, batch_size, learning_rate, patience)

Helper function to run the RGBUnet model from the command line entry point.

Parameters:

• logs – directory for log files for tensorboard.
• data – root directory of training data.
• working – working directory for organising data.
• omit – patient identifier to omit, when doing Leave-One-Out.
• model – file of previously saved model.
• save – file to save model to.
• test – input image to test.
• prediction – output image, the result of the prediction on test image.
• epochs – number of epochs.
• batch_size – batch size.
• learning_rate – learning rate for optimizer.
• patience – number of steps to tolerate non-improving accuracy