From 538cc3e960971416eb8f2ff4e8412f6c73687454 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Tue, 27 Jul 2021 08:08:01 -0700 Subject: [PATCH] Adding visibility of StyleGAN architecture for another project. PiperOrigin-RevId: 387110450 --- .../projects/gan/architectures_style_gan_v2.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tensorflow_graphics/projects/gan/architectures_style_gan_v2.py b/tensorflow_graphics/projects/gan/architectures_style_gan_v2.py index 4817d40a7..4eed2ff8e 100644 --- a/tensorflow_graphics/projects/gan/architectures_style_gan_v2.py +++ b/tensorflow_graphics/projects/gan/architectures_style_gan_v2.py @@ -242,6 +242,7 @@ def create_discriminator( kernel_initializer: Optional[_KerasInitializer] = None, use_fan_in_scaled_kernels: bool = True, use_antialiased_bilinear_downsampling: bool = False, + num_channels: int = 3, name: str = 'style_gan_v2_discriminator'): """Creates a Keras model for the discriminator architecture. @@ -262,6 +263,7 @@ def create_discriminator( ani-aliased bilinear downsampling with a [1, 3, 3, 1] tent kernel. If false standard bilinear downsampling, i.e. average pooling is used ([1, 1] tent kernel). + num_channels: The number of channels of the input tensor. name: The name of the Keras model. Returns: @@ -271,7 +273,7 @@ def create_discriminator( kernel_initializer = tf.keras.initializers.TruncatedNormal( mean=0.0, stddev=1.0) - input_tensor = tf.keras.Input(shape=(None, None, 3)) + input_tensor = tf.keras.Input(shape=(None, None, num_channels)) tensor = architectures_progressive_gan.from_rgb( input_tensor=input_tensor, use_fan_in_scaled_kernel=use_fan_in_scaled_kernels,