def run(): """Builds model, loads data, trains and evaluates""" model = UNet(CFG) model.load_data() model.build() model.train() model.evaluate()
class UnetTest(tf.test.TestCase): def setUp(self): super(UnetTest, self).setUp() self.unet = UNet(CFG) def tearDown(self): pass def test_normalize(self): input_image = np.array([[1., 1.], [1., 1.]]) input_mask = 1 expected_image = np.array([[0.00392157, 0.00392157], [0.00392157, 0.00392157]]) result = self.unet._normalize(input_image, input_mask) self.assertAllClose(expected_image, result[0]) def test_ouput_size(self): shape = (1, self.unet.image_size, self.unet.image_size, 3) image = tf.ones(shape) self.unet.build() self.assertEqual(self.unet.model.predict(image).shape, shape) @patch('model.unet.DataLoader.load_data') def test_load_data(self, mock_data_loader): mock_data_loader.side_effect = dummy_load_data shape = tf.TensorShape([None, self.unet.image_size, self.unet.image_size, 3]) self.unet.load_data() mock_data_loader.assert_called() self.assertItemsEqual(self.unet.train_dataset.element_spec[0].shape, shape) self.assertItemsEqual(self.unet.test_dataset.element_spec[0].shape, shape)