from __future__ import absolute_import, print_function

import tensorflow as tf
from tensorflow.contrib.layers.python.layers import regularizers

from niftynet.network.unet_2d import UNet2D


class UNet3DTest(tf.test.TestCase):
    def test_2d_shape(self):
        input_shape = (2, 572, 572, 3)
        x = tf.ones(input_shape)

        unet_instance = UNet2D(num_classes=2)
        out = unet_instance(x, is_training=True)
        print(unet_instance.num_trainable_params())

        with self.test_session() as sess:
            sess.run(tf.global_variables_initializer())
            out = sess.run(out)
            self.assertAllClose((2, 388, 388, 2), out.shape)

    def test_2d_reg_shape(self):
        input_shape = (2, 572, 572, 5)
        x = tf.ones(input_shape)

        unet_instance = UNet2D(num_classes=2,
                               w_regularizer=regularizers.l2_regularizer(0.4))
        out = unet_instance(x, is_training=True)
        print(unet_instance.num_trainable_params())

        with self.test_session() as sess:
            sess.run(tf.global_variables_initializer())
            out = sess.run(out)
            self.assertAllClose((2, 388, 388, 2), out.shape)


if __name__ == "__main__":
    tf.test.main()

@@ -250,32 +250,33 @@
 (self):%0A        
+#
 input_shape = (2
@@ -290,16 +290,55 @@
 572, 3)%0A
+        input_shape = (2, 180, 180, 3)%0A
         
@@ -645,32 +645,33 @@
 ut)%0A            
+#
 self.assertAllCl
@@ -698,24 +698,81 @@
 , out.shape)
+%0A            self.assertAllClose((2, 4, 4, 2), out.shape)
 %0A%0A    def te
@@ -794,32 +794,33 @@
 (self):%0A        
+#
 input_shape = (2
@@ -834,16 +834,55 @@
 572, 5)%0A
+        input_shape = (2, 180, 180, 5)%0A
         
@@ -1268,32 +1268,33 @@
 ut)%0A            
+#
 self.assertAllCl
@@ -1326,16 +1326,73 @@
 .shape)%0A
+            self.assertAllClose((2, 4, 4, 2), out.shape)%0A
 %0A%0Aif __n

fix slow unet tests