From 4bf74f5d3b497ed33722e9a8c4e33c9088c70a44 Mon Sep 17 00:00:00 2001 From: Vinayakjeet Singh Karki <139736674+vinayakjeet@users.noreply.github.com> Date: Fri, 22 Mar 2024 16:04:56 +0530 Subject: [PATCH 1/9] Fixes #243 Fix UNet implementation to support input with channel sizes other than 3 --- src/convnets/unet.jl | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/src/convnets/unet.jl b/src/convnets/unet.jl index 9a06a040..49f44bf5 100644 --- a/src/convnets/unet.jl +++ b/src/convnets/unet.jl @@ -71,18 +71,21 @@ Backbone of any Metalhead ResNet-like model can be used as encoder - `final`: final block as described in original paper - `fdownscale`: downscale factor """ -function unet(encoder_backbone, imgdims, outplanes::Integer, - final::Any = unet_final_block, fdownscale::Integer = 0) - backbonelayers = collect(flatten_chains(encoder_backbone)) - layers = unetlayers(backbonelayers, imgdims; m_middle = unet_middle_block, - skip_upscale = fdownscale) +function unet(encoder_backbone, imgdims, inchannels::Integer, outplanes::Integer, + final::Any = unet_final_block, fdownscale::Integer = 0) +backbonelayers = collect(flatten_chains(encoder_backbone)) - outsz = Flux.outputsize(layers, imgdims) - layers = Chain(layers, final(outsz[end - 1], outplanes)) +# Adjusting input size to include channels +adjusted_imgdims = (imgdims..., inchannels, 1) - return layers -end +layers = unetlayers(backbonelayers, adjusted_imgdims; m_middle = unet_middle_block, + skip_upscale = fdownscale) + +outsz = Flux.outputsize(layers, adjusted_imgdims) +layers = Chain(layers, final(outsz[end - 1], outplanes)) +return layers +end """ UNet(imsize::Dims{2} = (256, 256), inchannels::Integer = 3, outplanes::Integer = 3, encoder_backbone = Metalhead.backbone(DenseNet(121)); pretrain::Bool = false) @@ -114,7 +117,7 @@ end function UNet(imsize::Dims{2} = (256, 256), inchannels::Integer = 3, outplanes::Integer = 3, encoder_backbone = Metalhead.backbone(DenseNet(121)); pretrain::Bool = false) - layers = unet(encoder_backbone, (imsize..., inchannels, 1), outplanes) + layers = unet(encoder_backbone, imsize, inchannels, outplanes) model = UNet(layers) if pretrain artifact_name = "UNet" From 3c92c88c465206ca199cf141a288d51072b2318e Mon Sep 17 00:00:00 2001 From: Vinayakjeet Singh Karki <139736674+vinayakjeet@users.noreply.github.com> Date: Sat, 23 Mar 2024 17:26:59 +0530 Subject: [PATCH 2/9] Update unet.jl Modified the first convolutional layer of the encoder backbone to ensure compatibility with the input's channel size and dimension mismatch error is thus prevented #1 --- src/convnets/unet.jl | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/src/convnets/unet.jl b/src/convnets/unet.jl index 49f44bf5..5e10d3b0 100644 --- a/src/convnets/unet.jl +++ b/src/convnets/unet.jl @@ -86,6 +86,19 @@ layers = Chain(layers, final(outsz[end - 1], outplanes)) return layers end +function modify_first_conv_layer(encoder_backbone, inchannels) + for (index, layer) in enumerate(encoder_backbone.layers) + if isa(layer, Flux.Conv) # Checking for a convolutional layer + # Extracting the parameters + outchannels, kernel_size, stride, pad, activation = layer.out_channels, layer.kernel_size, layer.stride, layer.pad, layer.activation + # new convolutional layer created for desired input + new_conv_layer = Flux.Conv(kernel_size, inchannels => outchannels, stride=stride, pad=pad, activation=activation) + encoder_backbone.layers[index] = new_conv_layer + break + end + end + return encoder_backbone +end """ UNet(imsize::Dims{2} = (256, 256), inchannels::Integer = 3, outplanes::Integer = 3, encoder_backbone = Metalhead.backbone(DenseNet(121)); pretrain::Bool = false) @@ -117,12 +130,18 @@ end function UNet(imsize::Dims{2} = (256, 256), inchannels::Integer = 3, outplanes::Integer = 3, encoder_backbone = Metalhead.backbone(DenseNet(121)); pretrain::Bool = false) + # Modify the encoder backbone to adjust the first convolutional layer's input channels + encoder_backbone = modify_first_conv_layer(encoder_backbone, inchannels) + layers = unet(encoder_backbone, imsize, inchannels, outplanes) model = UNet(layers) + if pretrain + artifact_name = "UNet" loadpretrain!(model, artifact_name) end + return model end From 4f0114574078024dd7782ce5908cfabd5e8d95a4 Mon Sep 17 00:00:00 2001 From: Vinayakjeet Singh Karki <139736674+vinayakjeet@users.noreply.github.com> Date: Sun, 24 Mar 2024 05:58:04 +0530 Subject: [PATCH 3/9] Update unet.jl Indentation issue resolved --- src/convnets/unet.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/convnets/unet.jl b/src/convnets/unet.jl index 5e10d3b0..faf9dbd6 100644 --- a/src/convnets/unet.jl +++ b/src/convnets/unet.jl @@ -71,7 +71,7 @@ Backbone of any Metalhead ResNet-like model can be used as encoder - `final`: final block as described in original paper - `fdownscale`: downscale factor """ -function unet(encoder_backbone, imgdims, inchannels::Integer, outplanes::Integer, +function unet(encoder_backbone, imgdims, inchannels::Integer,outplanes::Integer, final::Any = unet_final_block, fdownscale::Integer = 0) backbonelayers = collect(flatten_chains(encoder_backbone)) From 6c7cfaae97b05511cd82ba25efb0d3b3986f9e9e Mon Sep 17 00:00:00 2001 From: Vinayakjeet Singh Karki <139736674+vinayakjeet@users.noreply.github.com> Date: Sun, 24 Mar 2024 06:27:43 +0530 Subject: [PATCH 4/9] Update unet.jl 2nd try --- src/convnets/unet.jl | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/src/convnets/unet.jl b/src/convnets/unet.jl index faf9dbd6..5c3da9ba 100644 --- a/src/convnets/unet.jl +++ b/src/convnets/unet.jl @@ -88,10 +88,16 @@ return layers end function modify_first_conv_layer(encoder_backbone, inchannels) for (index, layer) in enumerate(encoder_backbone.layers) - if isa(layer, Flux.Conv) # Checking for a convolutional layer - # Extracting the parameters - outchannels, kernel_size, stride, pad, activation = layer.out_channels, layer.kernel_size, layer.stride, layer.pad, layer.activation - # new convolutional layer created for desired input + if isa(layer, Flux.Conv) + # Correctly infer the number of output channels from the layer's weight dimensions + outchannels = size(layer.weight, 1) # The first dimension for Flux.Conv weight is the number of output channels + + kernel_size = (size(layer.weight, 3), size(layer.weight, 4)) # height and width of the kernel + stride = layer.stride + pad = layer.pad + activation = layer.activation + + # Create a new convolutional layer with the adjusted number of input channels new_conv_layer = Flux.Conv(kernel_size, inchannels => outchannels, stride=stride, pad=pad, activation=activation) encoder_backbone.layers[index] = new_conv_layer break @@ -99,6 +105,8 @@ function modify_first_conv_layer(encoder_backbone, inchannels) end return encoder_backbone end + + """ UNet(imsize::Dims{2} = (256, 256), inchannels::Integer = 3, outplanes::Integer = 3, encoder_backbone = Metalhead.backbone(DenseNet(121)); pretrain::Bool = false) From 689c3b140c4fdc79072b1745939f7e0327c3a3e4 Mon Sep 17 00:00:00 2001 From: Vinayakjeet Singh Karki <139736674+vinayakjeet@users.noreply.github.com> Date: Sun, 24 Mar 2024 06:46:09 +0530 Subject: [PATCH 5/9] Update unet.jl 3rd try --- src/convnets/unet.jl | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/src/convnets/unet.jl b/src/convnets/unet.jl index 5c3da9ba..5586fdd1 100644 --- a/src/convnets/unet.jl +++ b/src/convnets/unet.jl @@ -88,25 +88,26 @@ return layers end function modify_first_conv_layer(encoder_backbone, inchannels) for (index, layer) in enumerate(encoder_backbone.layers) - if isa(layer, Flux.Conv) + if isa(layer, Flux.Conv) # Correctly infer the number of output channels from the layer's weight dimensions - outchannels = size(layer.weight, 1) # The first dimension for Flux.Conv weight is the number of output channels + outchannels = size(layer.weight, 1) - kernel_size = (size(layer.weight, 3), size(layer.weight, 4)) # height and width of the kernel + # Recreate the convolutional layer with the new number of input channels + kernel_size = (size(layer.weight, 3), size(layer.weight, 4)) stride = layer.stride pad = layer.pad - activation = layer.activation + new_conv_layer = Flux.Conv(kernel_size, inchannels => outchannels, stride=stride, pad=pad) - # Create a new convolutional layer with the adjusted number of input channels - new_conv_layer = Flux.Conv(kernel_size, inchannels => outchannels, stride=stride, pad=pad, activation=activation) + # Update the layer in the backbone encoder_backbone.layers[index] = new_conv_layer - break + break # Assume only the first Conv layer needs adjustment end end return encoder_backbone end + """ UNet(imsize::Dims{2} = (256, 256), inchannels::Integer = 3, outplanes::Integer = 3, encoder_backbone = Metalhead.backbone(DenseNet(121)); pretrain::Bool = false) From 238bfa9750c0fc0173239ff3fe92c9b9d1db5925 Mon Sep 17 00:00:00 2001 From: Vinayakjeet Singh Karki <139736674+vinayakjeet@users.noreply.github.com> Date: Sun, 24 Mar 2024 07:36:04 +0530 Subject: [PATCH 6/9] Update unet.jl 4th try --- src/convnets/unet.jl | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/src/convnets/unet.jl b/src/convnets/unet.jl index 5586fdd1..030bbdb4 100644 --- a/src/convnets/unet.jl +++ b/src/convnets/unet.jl @@ -86,24 +86,25 @@ layers = Chain(layers, final(outsz[end - 1], outplanes)) return layers end -function modify_first_conv_layer(encoder_backbone, inchannels) - for (index, layer) in enumerate(encoder_backbone.layers) - if isa(layer, Flux.Conv) - # Correctly infer the number of output channels from the layer's weight dimensions - outchannels = size(layer.weight, 1) - - # Recreate the convolutional layer with the new number of input channels - kernel_size = (size(layer.weight, 3), size(layer.weight, 4)) +function modify_first_conv_layer_advanced(encoder_backbone, inchannels) + layers = [layer for layer in encoder_backbone.layers] # Create a mutable array from the layers + modified = false + for index in 1:length(layers) + if isa(layers[index], Flux.Conv) && !modified + layer = layers[index] + outchannels = size(layer.weight, 1) # The number of output channels + kernel_size = (size(layer.weight, 3), size(layer.weight, 4)) # Kernel size stride = layer.stride pad = layer.pad + + # Create a new convolutional layer with the updated input channels new_conv_layer = Flux.Conv(kernel_size, inchannels => outchannels, stride=stride, pad=pad) + layers[index] = new_conv_layer # Replace the old layer with the new one - # Update the layer in the backbone - encoder_backbone.layers[index] = new_conv_layer - break # Assume only the first Conv layer needs adjustment + modified = true # Mark as modified to avoid changing any other Conv layer end end - return encoder_backbone + return Flux.Chain(layers...) # Reconstruct the model with the modified layers end From c98640b13693e2feb9133528f5ca851b93c988ad Mon Sep 17 00:00:00 2001 From: Vinayakjeet Singh Karki <139736674+vinayakjeet@users.noreply.github.com> Date: Sun, 24 Mar 2024 07:56:38 +0530 Subject: [PATCH 7/9] Update unet.jl 5th try --- src/convnets/unet.jl | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/src/convnets/unet.jl b/src/convnets/unet.jl index 030bbdb4..8ad053d6 100644 --- a/src/convnets/unet.jl +++ b/src/convnets/unet.jl @@ -136,21 +136,18 @@ See also [`Metalhead.unet`](@ref). struct UNet layers::Any end + @functor UNet function UNet(imsize::Dims{2} = (256, 256), inchannels::Integer = 3, outplanes::Integer = 3, encoder_backbone = Metalhead.backbone(DenseNet(121)); pretrain::Bool = false) - # Modify the encoder backbone to adjust the first convolutional layer's input channels - encoder_backbone = modify_first_conv_layer(encoder_backbone, inchannels) + # Correct the function name to `modify_first_conv_layer_advanced` + encoder_backbone = modify_first_conv_layer_advanced(encoder_backbone, inchannels) layers = unet(encoder_backbone, imsize, inchannels, outplanes) model = UNet(layers) - if pretrain - - artifact_name = "UNet" - loadpretrain!(model, artifact_name) - end + # Pretraining code here (if applicable) return model end From 1e892a4968e6187d21d5ccd2078d328301d9f744 Mon Sep 17 00:00:00 2001 From: Vinayakjeet Singh Karki <139736674+vinayakjeet@users.noreply.github.com> Date: Sun, 24 Mar 2024 08:32:50 +0530 Subject: [PATCH 8/9] Update unet.jl 6th try --- src/convnets/unet.jl | 23 +++++++++++++++-------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/src/convnets/unet.jl b/src/convnets/unet.jl index 8ad053d6..9b3628b3 100644 --- a/src/convnets/unet.jl +++ b/src/convnets/unet.jl @@ -138,18 +138,25 @@ struct UNet end @functor UNet - function UNet(imsize::Dims{2} = (256, 256), inchannels::Integer = 3, outplanes::Integer = 3, encoder_backbone = Metalhead.backbone(DenseNet(121)); pretrain::Bool = false) - # Correct the function name to `modify_first_conv_layer_advanced` - encoder_backbone = modify_first_conv_layer_advanced(encoder_backbone, inchannels) - - layers = unet(encoder_backbone, imsize, inchannels, outplanes) + # Modify the first convolutional layer of the encoder backbone to have the correct `inchannels`. + # This is a conceptual step; the actual implementation will depend on the structure of your backbone. + if inchannels != 3 + encoder_backbone = modify_first_conv_layer_advanced(encoder_backbone, inchannels) + + end + + layers = unet(encoder_backbone, (imsize..., inchannels, 1), outplanes) model = UNet(layers) - - # Pretraining code here (if applicable) - + if pretrain + # Note: As per the original comment, pre-trained weights are not supported in this context. + # This block is left as-is from your original code for completeness. + artifact_name = "UNet" + loadpretrain!(model, artifact_name) + end return model end + (m::UNet)(x::AbstractArray) = m.layers(x) From 18478fe58b5c40c0a3e2d04c50b0a8be32ed8864 Mon Sep 17 00:00:00 2001 From: Vinayakjeet Singh Karki <139736674+vinayakjeet@users.noreply.github.com> Date: Sun, 24 Mar 2024 08:39:37 +0530 Subject: [PATCH 9/9] Update unet.jl --- src/convnets/unet.jl | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/convnets/unet.jl b/src/convnets/unet.jl index 9b3628b3..d78756f5 100644 --- a/src/convnets/unet.jl +++ b/src/convnets/unet.jl @@ -97,7 +97,7 @@ function modify_first_conv_layer_advanced(encoder_backbone, inchannels) stride = layer.stride pad = layer.pad - # Create a new convolutional layer with the updated input channels + new_conv_layer = Flux.Conv(kernel_size, inchannels => outchannels, stride=stride, pad=pad) layers[index] = new_conv_layer # Replace the old layer with the new one @@ -140,8 +140,7 @@ end @functor UNet function UNet(imsize::Dims{2} = (256, 256), inchannels::Integer = 3, outplanes::Integer = 3, encoder_backbone = Metalhead.backbone(DenseNet(121)); pretrain::Bool = false) - # Modify the first convolutional layer of the encoder backbone to have the correct `inchannels`. - # This is a conceptual step; the actual implementation will depend on the structure of your backbone. + if inchannels != 3 encoder_backbone = modify_first_conv_layer_advanced(encoder_backbone, inchannels) @@ -150,8 +149,7 @@ function UNet(imsize::Dims{2} = (256, 256), inchannels::Integer = 3, outplanes:: layers = unet(encoder_backbone, (imsize..., inchannels, 1), outplanes) model = UNet(layers) if pretrain - # Note: As per the original comment, pre-trained weights are not supported in this context. - # This block is left as-is from your original code for completeness. + artifact_name = "UNet" loadpretrain!(model, artifact_name) end