text: from tensorflow.compat.v1 import ConfigProto
  from tensorflow.compat.v1 import InteractiveSession
  
  config = ConfigProto()
  config.gpu_options.allow_growth = True
  session = InteractiveSession(config=config)
  
  from keras_segmentation.models.unet import vgg_unet
  
  model = vgg_unet(n_classes=51 ,  input_height=416, input_width=608  )
  
  model.train(
      train_images =  "CelebA-HQ-img",
      train_annotations = "mass",
      checkpoints_path = "/tmp/vgg_unet_1" , epochs=20
  )
  
  # out = model.predict_segmentation(
  #     inp="dataset1/images_prepped_test/0016E5_07965.png",
  #     out_fname="/tmp/out.png"
  # )
  
  # import matplotlib.pyplot as plt
  # plt.imshow(out)