FluxML · darsnack · Jan 27, 2023 · Dec 27, 2022 · Dec 27, 2022 · Dec 28, 2022
diff --git a/docs/make.jl b/docs/make.jl
@@ -1,26 +1,34 @@
-using Documenter, Metalhead, Artifacts, LazyArtifacts, Images, OneHotArrays, DataAugmentation, Flux
+using Documenter, Metalhead, Artifacts, LazyArtifacts, Images, OneHotArrays,
+      DataAugmentation, Flux
 
 DocMeta.setdocmeta!(Metalhead, :DocTestSetup, :(using Metalhead); recursive = true)
 
-makedocs(modules = [Metalhead, Artifacts, LazyArtifacts, Images, OneHotArrays, DataAugmentation, Flux],
+makedocs(;
+         modules = [
+             Metalhead,
+             Artifacts,
+             LazyArtifacts,
+             Images,
+             OneHotArrays,
+             DataAugmentation,
+             Flux,
+         ],
          sitename = "Metalhead.jl",
          doctest = false,
          pages = ["Home" => "index.md",
-                  "Tutorials" => [
-                      "tutorials/quickstart.md",
-                   ],
-                  "Developer guide" => "contributing.md",
-                  "API reference" => [
-                      "api/reference.md",
-                   ],
-                 ],
-         format = Documenter.HTML(
-              canonical = "https://fluxml.ai/Metalhead.jl/stable/",
-            #   analytics = "UA-36890222-9",
-              assets = ["assets/flux.css"],
-              prettyurls = get(ENV, "CI", nothing) == "true"),
-        )
+             "Tutorials" => [
+                 "tutorials/quickstart.md",
+             ],
+             "Developer guide" => "contributing.md",
+             "API reference" => [
+                 "api/reference.md",
+             ],
+         ],
+         format = Documenter.HTML(; canonical = "https://fluxml.ai/Metalhead.jl/stable/",
+                                  #   analytics = "UA-36890222-9",
+                                  assets = ["assets/flux.css"],
+                                  prettyurls = get(ENV, "CI", nothing) == "true"))
 
-deploydocs(repo = "github.com/FluxML/Metalhead.jl.git",
+deploydocs(; repo = "github.com/FluxML/Metalhead.jl.git",
            target = "build",
            push_preview = true)
diff --git a/docs/src/index.md b/docs/src/index.md
@@ -36,6 +36,7 @@ julia> ]add Metalhead
 | [ViT](https://arxiv.org/abs/2010.11929)                | [`ViT`](@ref)               | N            |
 | [ConvNeXt](https://arxiv.org/abs/2201.03545)           | [`ConvNeXt`](@ref)          | N            |
 | [ConvMixer](https://arxiv.org/abs/2201.09792)          | [`ConvMixer`](@ref)         | N            |
+| [UNet](https://arxiv.org/abs/1505.04597v1)             | [`UNet`](@ref)              | N            |
 
 To contribute new models, see our [contributing docs](@ref Contributing-to-Metalhead.jl).
 

diff --git a/src/Metalhead.jl b/src/Metalhead.jl
@@ -2,6 +2,7 @@ module Metalhead
 
 using Flux
 using Flux: Zygote, outputsize
+using Distributions: Normal
 using Functors
 using BSON
 using Artifacts, LazyArtifacts
@@ -10,7 +11,7 @@ using MLUtils
 using PartialFunctions
 using Random
 
-import Functors
+using Functors: Functors
 
 include("utilities.jl")
 
@@ -28,6 +29,8 @@ include("convnets/builders/stages.jl")
 ## AlexNet and VGG
 include("convnets/alexnet.jl")
 include("convnets/vgg.jl")
+## Unet
+include("convnets/unet.jl")
 ## ResNets
 include("convnets/resnets/core.jl")
 include("convnets/resnets/res2net.jl")
@@ -66,7 +69,7 @@ include("vit-based/vit.jl")
 # Load pretrained weights
 include("pretrain.jl")
 
-export AlexNet, VGG, VGG11, VGG13, VGG16, VGG19,
+export AlexNet, VGG, VGG11, VGG13, VGG16, VGG19, UNet,
        ResNet, ResNet18, ResNet34, ResNet50, ResNet101, ResNet152,
        WideResNet, ResNeXt, SEResNet, SEResNeXt, Res2Net, Res2NeXt,
        DenseNet, DenseNet121, DenseNet161, DenseNet169, DenseNet201,
@@ -76,7 +79,7 @@ export AlexNet, VGG, VGG11, VGG13, VGG16, VGG19,
        MLPMixer, ResMLP, gMLP, ViT
 
 # use Flux._big_show to pretty print large models
-for T in (:AlexNet, :VGG, :SqueezeNet, :ResNet, :WideResNet, :ResNeXt,
+for T in (:AlexNet, :VGG, :UNet, :SqueezeNet, :ResNet, :WideResNet, :ResNeXt,
           :SEResNet, :SEResNeXt, :Res2Net, :Res2NeXt, :GoogLeNet, :DenseNet,
           :Inceptionv3, :Inceptionv4, :InceptionResNetv2, :Xception,
           :MobileNetv1, :MobileNetv2, :MobileNetv3, :MNASNet,

diff --git a/src/convnets/alexnet.jl b/src/convnets/alexnet.jl
@@ -42,7 +42,7 @@ Create a `AlexNet`.
   - `nclasses`: the number of output classes
 
 !!! warning
-    
+
     `AlexNet` does not currently support pretrained weights.
 
 See also [`alexnet`](@ref).

diff --git a/src/convnets/builders/resnet.jl b/src/convnets/builders/resnet.jl
@@ -5,7 +5,7 @@
 Creates a generic ResNet-like model.
 
 !!! info
-    
+
     This is a very generic, flexible but low level function that can be used to create any of the ResNet
     variants. For a more user friendly function, see [`Metalhead.resnet`](@ref).
 

diff --git a/src/convnets/convnext.jl b/src/convnets/convnext.jl
@@ -97,7 +97,7 @@ Creates a ConvNeXt model.
   - `nclasses`: number of output classes
 
 !!! warning
-    
+
     `ConvNeXt` does not currently support pretrained weights.
 
 See also [`Metalhead.convnext`](@ref).

diff --git a/src/convnets/densenet.jl b/src/convnets/densenet.jl
@@ -122,7 +122,7 @@ Create a DenseNet model with specified configuration. Currently supported values
 Set `pretrain = true` to load the model with pre-trained weights for ImageNet.
 
 !!! warning
-    
+
     `DenseNet` does not currently support pretrained weights.
 
 See also [`Metalhead.densenet`](@ref).

diff --git a/src/convnets/efficientnets/efficientnet.jl b/src/convnets/efficientnets/efficientnet.jl
@@ -73,7 +73,7 @@ Create an EfficientNet model ([reference](https://arxiv.org/abs/1905.11946v5)).
   - `nclasses`: number of output classes.
 
 !!! warning
-    
+
     EfficientNet does not currently support pretrained weights.
 
 See also [`Metalhead.efficientnet`](@ref).

diff --git a/src/convnets/efficientnets/efficientnetv2.jl b/src/convnets/efficientnets/efficientnetv2.jl
@@ -76,7 +76,7 @@ Create an EfficientNetv2 model ([reference](https://arxiv.org/abs/2104.00298)).
   - `nclasses`: number of output classes
 
 !!! warning
-    
+
     `EfficientNetv2` does not currently support pretrained weights.
 
 See also [`efficientnet`](#).

diff --git a/src/convnets/inceptions/googlenet.jl b/src/convnets/inceptions/googlenet.jl
@@ -86,7 +86,7 @@ Create an Inception-v1 model (commonly referred to as `GoogLeNet`)
   - `bias`: set to `true` to use bias in the convolution layers
 
 !!! warning
-    
+
     `GoogLeNet` does not currently support pretrained weights.
 
 See also [`Metalhead.googlenet`](@ref).

diff --git a/src/convnets/inceptions/inceptionresnetv2.jl b/src/convnets/inceptions/inceptionresnetv2.jl
@@ -109,7 +109,7 @@ Creates an InceptionResNetv2 model.
   - `nclasses`: the number of output classes.
 
 !!! warning
-    
+
     `InceptionResNetv2` does not currently support pretrained weights.
 
 See also [`Metalhead.inceptionresnetv2`](@ref).

diff --git a/src/convnets/inceptions/inceptionv3.jl b/src/convnets/inceptions/inceptionv3.jl
@@ -170,7 +170,7 @@ Create an Inception-v3 model ([reference](https://arxiv.org/abs/1512.00567v3)).
   - `nclasses`: the number of output classes
 
 !!! warning
-    
+
     `Inceptionv3` does not currently support pretrained weights.
 
 See also [`Metalhead.inceptionv3`](@ref).

diff --git a/src/convnets/inceptions/inceptionv4.jl b/src/convnets/inceptions/inceptionv4.jl
@@ -124,7 +124,7 @@ Creates an Inceptionv4 model.
   - `nclasses`: the number of output classes.
 
 !!! warning
-    
+
     `Inceptionv4` does not currently support pretrained weights.
 
 See also [`Metalhead.inceptionv4`](@ref).

diff --git a/src/convnets/inceptions/xception.jl b/src/convnets/inceptions/xception.jl
@@ -80,7 +80,7 @@ Creates an Xception model.
   - `nclasses`: the number of output classes.
 
 !!! warning
-    
+
     `Xception` does not currently support pretrained weights.
 
 See also [`Metalhead.xception`](@ref).

diff --git a/src/convnets/mobilenets/mnasnet.jl b/src/convnets/mobilenets/mnasnet.jl
@@ -86,7 +86,7 @@ Creates a MNASNet model with the specified configuration.
   - `nclasses`: The number of output classes
 
 !!! warning
-    
+
     `MNASNet` does not currently support pretrained weights.
 
 See also [`Metalhead.mnasnet`](@ref).

diff --git a/src/convnets/mobilenets/mobilenetv1.jl b/src/convnets/mobilenets/mobilenetv1.jl
@@ -52,7 +52,7 @@ Create a MobileNetv1 model with the baseline configuration
   - `nclasses`: The number of output classes
 
 !!! warning
-    
+
     `MobileNetv1` does not currently support pretrained weights.
 
 See also [`Metalhead.mobilenetv1`](@ref).

diff --git a/src/convnets/mobilenets/mobilenetv2.jl b/src/convnets/mobilenets/mobilenetv2.jl
@@ -58,7 +58,7 @@ Create a MobileNetv2 model with the specified configuration.
   - `nclasses`: The number of output classes
 
 !!! warning
-    
+
     `MobileNetv2` does not currently support pretrained weights.
 
 See also [`Metalhead.mobilenetv2`](@ref).

diff --git a/src/convnets/mobilenets/mobilenetv3.jl b/src/convnets/mobilenets/mobilenetv3.jl
@@ -78,7 +78,7 @@ Set `pretrain = true` to load the model with pre-trained weights for ImageNet.
   - `nclasses`: the number of output classes
 
 !!! warning
-    
+
     `MobileNetv3` does not currently support pretrained weights.
 
 See also [`Metalhead.mobilenetv3`](@ref).

diff --git a/src/convnets/resnets/core.jl b/src/convnets/resnets/core.jl
@@ -191,7 +191,7 @@ If `outplanes > inplanes`, it maps the input to `outplanes` channels using a 1x1
 layer and zero padding.
 
 !!! warning
-    
+
     This does not currently support the scenario where `inplanes > outplanes`.
 
 # Arguments
@@ -237,7 +237,7 @@ on how to use this function.
 # Arguments
 
   - `stem_type`: The type of stem to be built. One of `[:default, :deep, :deep_tiered]`.
-    
+
       + `:default`: Builds a stem based on the default ResNet stem, which consists of a single
         7x7 convolution with stride 2 and a normalisation layer followed by a 3x3 max pooling
         layer with stride 2.

diff --git a/src/convnets/resnets/res2net.jl b/src/convnets/resnets/res2net.jl
@@ -17,7 +17,7 @@ Creates a Res2Net model with the specified depth, scale, and base width.
   - `nclasses`: the number of output classes
 
 !!! warning
-    
+
     `Res2Net` does not currently support pretrained weights.
 
 Advanced users who want more configuration options will be better served by using [`resnet`](@ref).
@@ -64,7 +64,7 @@ Creates a Res2NeXt model with the specified depth, scale, base width and cardina
   - `nclasses`: the number of output classes
 
 !!! warning
-    
+
     `Res2NeXt` does not currently support pretrained weights.
 
 Advanced users who want more configuration options will be better served by using [`resnet`](@ref).

diff --git a/src/convnets/resnets/resnext.jl b/src/convnets/resnets/resnext.jl
@@ -11,7 +11,7 @@ Creates a ResNeXt model with the specified depth, cardinality, and base width.
 
   - `pretrain`: set to `true` to load the model with pre-trained weights for ImageNet.
     Supported configurations are:
-    
+
       + depth 50, cardinality of 32 and base width of 4.
       + depth 101, cardinality of 32 and base width of 8.
       + depth 101, cardinality of 64 and base width of 4.

diff --git a/src/convnets/resnets/seresnet.jl b/src/convnets/resnets/seresnet.jl
@@ -12,7 +12,7 @@ Creates a SEResNet model with the specified depth.
   - `nclasses`: the number of output classes
 
 !!! warning
-    
+
     `SEResNet` does not currently support pretrained weights.
 
 Advanced users who want more configuration options will be better served by using [`resnet`](@ref).
@@ -55,7 +55,7 @@ Creates a SEResNeXt model with the specified depth, cardinality, and base width.
   - `nclasses`: the number of output classes
 
 !!! warning
-    
+
     `SEResNeXt` does not currently support pretrained weights.
 
 Advanced users who want more configuration options will be better served by using [`resnet`](@ref).

diff --git a/src/convnets/unet.jl b/src/convnets/unet.jl
@@ -0,0 +1,82 @@
+function unet_block(in_chs::Int, out_chs::Int, kernel = (3, 3))
+    return Chain(;
+                 conv1 = Conv(kernel, in_chs => out_chs; pad = (1, 1)),
+                 norm1 = BatchNorm(out_chs, relu),
+                 conv2 = Conv(kernel, out_chs => out_chs; pad = (1, 1)),
+                 norm2 = BatchNorm(out_chs, relu))
+end
+
+function upconv_block(in_chs::Int, out_chs::Int, kernel = (2, 2))
+    return ConvTranspose(kernel, in_chs => out_chs; stride = (2, 2))
+end
+
+function cat_fn(layers...)
+    return cat(layers...; dims = 3)
+end
+
+function unet(in_channels::Integer = 3, out_channels::Integer = in_channels,
+              features::Integer = 32)
+    encoder_conv_layers = []
+    append!(encoder_conv_layers,
+            [unet_block(in_channels, features)])
+
+    append!(encoder_conv_layers,
+            [unet_block(features * 2^i, features * 2^(i + 1)) for i in 0:2])
+
+    encoder_conv = Chain(encoder_conv_layers)
+    encoder_pool = [Chain(encoder_conv[i], MaxPool((2, 2); stride = (2, 2))) for i in 1:4]
+
+    bottleneck = unet_block(features * 8, features * 16)
+    layers = Chain(encoder_conv, bottleneck)
+
+    upconv = Chain([upconv_block(features * 2^(i + 1), features * 2^i)
+                    for i in 0:3])
+
+    concat_layer = Chain([Parallel(cat_fn,
+                                   encoder_pool[i],
+                                   upconv[i])
+                          for i in 1:4])
+
+    decoder_layer = Chain([unet_block(features * 2^(i + 1), features * 2^i) for i in 3:-1:0])
+
+    layers = Chain(layers, decoder_layer)
+
+    decoder = Chain([Chain([
+                               concat_layer[i],
+                               decoder_layer[5 - i]])
+                     for i in 4:-1:1])
+
+    final_conv = Conv((1, 1), features => out_channels, σ)
+
+    decoder = Chain(decoder, final_conv)
+
+    return layers
+end
+
+"""
+    UNet(inplanes::Integer = 3, outplanes::Integer = 1, init_features::Integer = 32)
+
+    Create a UNet model
+    ([reference](https://arxiv.org/abs/1505.04597v1))
+
+    # Arguments
+    - `in_channels`: The number of input channels
+    - `inplanes`: The number of input planes to the network
+    - `outplanes`: The number of output features
+
+!!! warning
+
+    `UNet` does not currently support pretrained weights.
+"""
+struct UNet
+    layers::Any
+end
+@functor UNet
+
+function UNet(in_channels::Integer = 3, inplanes::Integer = 32,
+              outplanes::Integer = inplanes)
+    layers = unet(in_channels, inplanes, outplanes)
+    return UNet(layers)
+end
+
+(m::UNet)(x::AbstractArray) = m.layers(x)