import tensorflow as tf

  model.fit(
      train_dataset,
      epochs=20,
      steps_per_epoch=steps_per_epoch,
      validation_steps=val_steps,
      validation_data=test_dataset,
      callbacks=[])

  model.save_weights('./test/segmentation')

  print(create_mask(model(tf.ones((1, 512, 512, 3)), False)))


if __name__ == '__main__':
  logging.set_verbosity(logging.WARNING)
  app.run(main)
