Fix UNet implementation with arbitrary channel sizes (#243)

src/convnets/unet.jl

@@ -71,18 +71,34 @@ 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
+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)
@@ -114,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)
-    layers = unet(encoder_backbone, (imsize..., inchannels, 1), outplanes)
+    # 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