try:
    from tensorflow.python.keras.models import Sequential
    from tensorflow.python.keras.layers import Conv2D, MaxPooling2D
    from tensorflow.python.keras.layers import BatchNormalization
except:
    from tensorflow.contrib.keras.python.keras.layers import Conv2D, MaxPooling2D
    from tensorflow.contrib.keras.python.keras.models import Sequential
    from tensorflow.contrib.keras.python.keras.layers.normalization import BatchNormalization
    from tensorflow.contrib.keras.python.keras.engine.topology import Layer
    from tensorflow.contrib.keras.python.keras import backend


class Squeeze(Layer):
    def __init__(self, output_dim=None, **kwargs):
        self.output_dim = output_dim
        super(Squeeze, self).__init__(**kwargs)

    def call(self, x):
        x = backend.squeeze(x, axis=2)
        return backend.squeeze(x, axis=1)

    def compute_output_shape(self, input_shape):
        return (input_shape[0], input_shape[3])



model = Sequential([
    Conv2D(64, kernel_size=(3, 3), activation='relu', input_shape=(None, None, 1)),
    BatchNormalization(),
    Conv2D(64, kernel_size=(4, 4), activation='relu'),
    MaxPooling2D(),
    Conv2D(64, kernel_size=(3, 3), activation='relu'),
    Conv2D(64, kernel_size=(3, 3), activation='relu'),
    MaxPooling2D(),
    Conv2D(64, kernel_size=(3, 3), activation='relu'),
    Conv2D(64, kernel_size=(3, 3), activation='relu'),
    MaxPooling2D(),
    Conv2D(200, kernel_size=(4, 4), activation='relu'),
    Conv2D(200, kernel_size=(1, 1), activation='relu'),
    Conv2D(3, kernel_size=(1, 1), activation='softmax'),  # filters num == # of labels
    Squeeze(3),
])