Fix WrapperBlock behavior and add Many (#158)

* Add dataset recipes and dataset registry

* Add fastai dataset registry with some recipes

* move `typify` helepr

* add learning method registry

* add test for `ImageSegmentationFolders` recipe

* import method registry

* move file

* fix default data registry definitin

* `WrapperBlock` now inherits from `WrapperBlock` to make wrapping work

* Fix `en|decodedblock` for `ImagePreprocessing` to reflect dimensionality of `ImageTensor`

* add missing `mockblock` for `OneHotMulti`

* update learning methods

* add convenience `plotpredictions` method

* Improve `WrapperBlock` behavior

* fix `WrapperBlock`s and implement `Many`

Author: lorenzoh

Project.toml

@@ -29,6 +29,7 @@ MosaicViews = "e94cdb99-869f-56ef-bcf0-1ae2bcbe0389"
 Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 ShowCases = "605ecd9f-84a6-4c9e-81e2-4798472b76a3"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
@@ -53,8 +54,8 @@ Flux = "0.12"
 FluxTraining = "0.2"
 Glob = "1"
 IndirectArrays = "0.5"
-LearnBase = "0.3, 0.4"
 JLD2 = "0.4"
+LearnBase = "0.3, 0.4"
 MLDataPattern = "0.5"
 Makie = "0.15"
 MosaicViews = "0.2, 0.3"

src/FastAI.jl

@@ -30,6 +30,7 @@ using MLDataPattern
 using Parameters
 using PrettyTables
 using StaticArrays
+using Setfield
 using ShowCases
 using Statistics: mean
 using Test: @testset, @test, @test_nowarn
@@ -57,7 +58,6 @@ include("datablock/loss.jl")
 include("datablock/plot.jl")
 
 
-
 # submodules
 include("datasets/Datasets.jl")
 @reexport using .Datasets
@@ -78,6 +78,11 @@ include("training/metrics.jl")
 include("serialization.jl")
 
 
+include("fasterai/methodregistry.jl")
+include("fasterai/learningmethods.jl")
+include("fasterai/defaults.jl")
+
+
 
 
 export
@@ -115,6 +120,7 @@ export
     Label,
     LabelMulti,
     Keypoints,
+    Many,
 
     # encodings
     encode,
@@ -131,6 +137,14 @@ export
     describemethod,
     checkblock,
 
+    # learning methods
+    findlearningmethods,
+    ImageClassificationSingle,
+    ImageClassificationMulti,
+    ImageSegmentation,
+    ImageKeypointRegression,
+
+
     # training
     methodlearner,
     Learner,

src/datablock/block.jl

@@ -60,6 +60,13 @@ Randomly generate an instance of `block`.
 mockblock(blocks::Tuple) = map(mockblock, blocks)
 
 
+# ## Utilities
+
+typify(T::Type) = T
+typify(t::Tuple) = Tuple{map(typify, t)...}
+typify(block::FastAI.AbstractBlock) = typeof(block)
+
+
 # ## Block implementations
 
 abstract type AbstractLabel{T} <: Block end

src/datablock/encoding.jl

@@ -39,8 +39,8 @@ Replaces all `nothing`s in outblocks with the corresponding block in `inblocks`.
 `outblocks` may be obtained by
 """
 fillblock(inblocks::Tuple, outblocks::Tuple) = map(fillblock, inblocks, outblocks)
-fillblock(inblock::Block, ::Nothing) = inblock
-fillblock(::Block, outblock::Block) = outblock
+fillblock(inblock::AbstractBlock, ::Nothing) = inblock
+fillblock(::AbstractBlock, outblock::AbstractBlock) = outblock
 
 function encodedblock(enc, block, fill::Bool)
     outblock = encodedblock(enc, block)

src/datablock/wrappers.jl

@@ -1,17 +1,32 @@
-abstract type WrapperBlock <: Block end
+# # Wrapper blocks
+
+abstract type WrapperBlock <: AbstractBlock end
 
 wrapped(w::WrapperBlock) = w.block
+function setwrapped(w::WrapperBlock, b)
+    return Setfield.@set w.block = b
+end
 mockblock(w::WrapperBlock) = mockblock(wrapped(w))
 checkblock(w::WrapperBlock, data) = checkblock(wrapped(w), data)
 
 # If not overwritten, encodings are applied to the wrapped block
+
+function encodedblock(enc::Encoding, wrapper::WrapperBlock)
+    inner = encodedblock(enc, wrapped(wrapper))
+    return isnothing(inner) ? nothing : setwrapped(wrapper, inner)
+end
+function decodedblock(enc::Encoding, wrapper::WrapperBlock)
+    inner = decodedblock(enc, wrapped(wrapper))
+    return isnothing(inner) ? nothing : setwrapped(wrapper, inner)
+end
 function encode(enc::Encoding, ctx, wrapper::WrapperBlock, data; kwargs...)
-    return encode(enc, ctx, wrapper.block, data; kwargs...)
+    return encode(enc, ctx, wrapped(wrapper), data; kwargs...)
 end
 function decode(enc::Encoding, ctx, wrapper::WrapperBlock, data; kwargs...)
-    return decode(enc, ctx, wrapper.block, data; kwargs...)
+    return decode(enc, ctx, wrapped(wrapper), data; kwargs...)
 end
 
+# ## Named
 
 """
     Named(name, block)
@@ -36,7 +51,39 @@ function decodedblock(enc::Encoding, named::Named{Name}) where Name
     return isnothing(outblock) ? nothing : Named(Name, outblock)
 end
 
-# Wrapper encodings
+# ## Many
+
+"""
+    Many(block) <: WrapperBlock
+
+`Many` indicates that you can variable number of data instances for
+`block`. Consider a bounding box detection task where there may be any
+number of targets in an image and this number varies for different
+samples. The blocks `(Image{2}(), BoundingBox{2}()` imply that there is exactly
+one bounding box for every image, which is not the case. Instead you
+would want to use `(Image{2}(), Many(BoundingBox{2}())`.
+"""
+struct Many{B<:AbstractBlock} <: WrapperBlock
+    block::B
+end
+
+FastAI.checkblock(many::Many, datas) = all(checkblock(wrapped(many), data) for data in datas)
+FastAI.mockblock(many::Many) = [mockblock(wrapped(many)), mockblock(wrapped(many))]
+
+function FastAI.encode(enc::Encoding, ctx, many::Many, datas)
+    return map(datas) do data
+        encode(enc, ctx, wrapped(many), data)
+    end
+end
+
+function FastAI.decode(enc::Encoding, ctx, many::Many, datas)
+    return map(datas) do data
+        decode(enc, ctx, wrapped(many), data)
+    end
+end
+
+
+# # Wrapper encodings
 
 """
     Only(name, encoding)

src/datasets/Datasets.jl

@@ -14,6 +14,7 @@ module Datasets
 
 
 using ..FastAI
+using ..FastAI: typify
 
 using DataDeps
 using Glob
@@ -41,6 +42,9 @@ include("containers.jl")
 include("transformations.jl")
 
 include("load.jl")
+include("recipes.jl")
+include("registry.jl")
+include("fastairegistry.jl")
 
 
 export
@@ -71,6 +75,12 @@ export
     # datasets
     DATASETS,
     loadfolderdata,
-    datasetpath
+    datasetpath,
+
+    # recipes
+    loadrecipe,
+    finddatasets,
+    listdatasources,
+    loaddataset
 
 end  # module

src/datasets/fastairegistry.jl

@@ -0,0 +1,24 @@
+
+const FASTAI_DATA_RECIPES = Dict{String, Vector{DatasetRecipe}}(
+    # Image classification datasets
+    [name => [ImageClassificationFolders()] for name in (
+        "imagenette", "imagenette-160", "imagenette-320",
+        "imagenette2", "imagenette2-160", "imagenette2-320",
+        "imagewoof", "imagewoof-160", "imagewoof-320",
+        "imagewoof2", "imagewoof2-160", "imagewoof2-320",
+    )]...,
+
+    "camvid_tiny" => [ImageSegmentationFolders()],
+)
+
+
+"""
+    const FASTAI_DATA_REGISTRY
+
+The default `DataRegistry` containing every dataset in
+the fastai dataset collection.
+"""
+const FASTAI_DATA_REGISTRY = DatasetRegistry(
+    Dict(d => () -> datasetpath(d) for d in DATASETS),
+    FASTAI_DATA_RECIPES,
+)

src/datasets/load.jl

@@ -51,29 +51,6 @@ function getclassessegmentation(dir::AbstractPath)
 end
 getclassessegmentation(name::String) = getclassessegmentation(datasetpath(name))
 
-#=
-"""
-    loadtaskdata(dir, ImageSegmentation; [split = false])
-
-Load a data container for `ImageSegmentation` with observations
-`(input = image, target = mask)`.
-
-If `split` is `true`, returns a tuple of the data containers split by
-the name of the grandparent folder.
-
-"""
-function loadtaskdata(
-        dir,
-        ::Type{FastAI.ImageSegmentation};
-        split=false,
-        kwargs...)
-    imagedata = mapobs(loadfile, filterobs(isimagefile, FileDataset(joinpath(dir, "images"))))
-    maskdata = mapobs(maskfromimage ∘ loadfile, filterobs(isimagefile, FileDataset(joinpath(dir, "labels"))))
-    return mapobs((input = obs -> obs[1], target = obs -> obs[2]), (imagedata, maskdata))
-end
-=#
-
-
 
 
 maskfromimage(a::AbstractArray{<:Gray{T}}, classes) where T = maskfromimage(reinterpret(T, a), classes)

src/datasets/recipes.jl

@@ -0,0 +1,126 @@
+"""
+abstract type DatasetRecipe
+
+A recipe that contains configuration for loading a data container. Calling it with a path returns a data container and the blocks that each sample is made of.
+
+#### Interface
+
+- `loadrecipe(::DatasetRecipe, args...) -> (data, blocks)`
+- `recipeblocks(::Type{DatasetRecipe}) -> TBlocks`
+
+#### Invariants
+
+- `data` must be a data container of samples that are valid `blocks`, i.e. `checkblock(blocks, getobs(data, 1)) == true`
+"""
+abstract type DatasetRecipe end
+
+
+"""
+    loadrecipe(recipe, path)
+
+Load a recipe from a path. Return a data container `data` and concrete
+`blocks`.
+"""
+function loadrecipe end
+
+
+"""
+    recipeblocks(TRecipe) -> TBlocks
+    recipeblocks(recipe) -> TBlocks
+
+Return the `Block` _types_ for the data container that recipe
+type `TRecipe` creates. Does not return `Block` instances as the exact
+configuration may not be known until the dataset is being
+loaded.
+
+#### Examples
+
+```julia
+recipeblocks(ImageLabelClf) == Tuple{Image{2}, Label}
+```
+"""
+recipeblocks(::R) where {R<:DatasetRecipe} = recipeblocks(R)
+
+
+# ## Implementations
+
+# ImageClfFolders
+
+"""
+    ImageClfFolders(; labelfn = parentname, split = false)
+
+Recipe for loading a single-label image classification dataset
+stored in a hierarchical folder format. If `split == true`, split
+the data container on the name of the grandparent folder. The label
+defaults to the name of the parent folder but a custom function can
+be passed as `labelfn`.
+
+```julia
+julia> recipeblocks(ImageClassificationFolders)
+Tuple{Image{2}, Label}
+```
+"""
+Base.@kwdef struct ImageClassificationFolders <: DatasetRecipe
+    labelfn = parentname
+    split::Bool = false
+end
+
+function loadrecipe(recipe::ImageClassificationFolders, path)
+    isdir(path) || error("$path is not a directory")
+    data = loadfolderdata(
+        path,
+        filterfn=isimagefile,
+        loadfn=(loadfile, recipe.labelfn),
+        splitfn=recipe.split ? grandparentname : nothing)
+
+    (recipe.split ? length(data) > 0 : nobs(data) > 0) || error("No image files found in $path")
+
+    labels = recipe.split ? first(values(data))[2] : data[2]
+    blocks = Image{2}(), Label(unique(eachobs(labels)))
+    length(blocks[2].classes) > 1 || error("Expected multiple different labels, got: $(blocks[2].classes))")
+    return data, blocks
+end
+
+recipeblocks(::Type{ImageClassificationFolders}) = Tuple{Image{2}, Label}
+
+
+# ImageSegmentationFolders
+
+
+"""
+    ImageSegmentationFolders(; imagefolder="images", maskfolder="labels", labelfile="codes.txt")
+
+Dataset recipe for loading 2D image segmentation datasets from a common format
+where images and masks are stored as images in two different subfolders
+"<root>/<imagefolder>" and "<root>/<maskfolder>"
+The class labels should be in a newline-delimited file "<root>/<labelfile>".
+"""
+Base.@kwdef struct ImageSegmentationFolders <: DatasetRecipe
+    imagefolder::String = "images"
+    maskfolder::String = "labels"
+    labelfile::String = "codes.txt"
+end
+
+function loadrecipe(recipe::ImageSegmentationFolders, path)
+    isdir(path) || error("$path is not a directory")
+    imagepath = joinpath(path, recipe.imagefolder)
+    maskpath = joinpath(path, recipe.maskfolder)
+    classespath = joinpath(path, recipe.labelfile)
+
+    isdir(imagepath) || error("Image folder $imagepath is not a directory")
+    isdir(maskpath) || error("Mask folder $maskpath is not a directory")
+
+    isfile(classespath) || error("Classes file $classespath does not exist")
+    classes = readlines(open(joinpath(path, recipe.labelfile)))
+    length(classes) > 1 || error("Expected multiple different labels, got: $(blocks[2].classes))")
+
+    images = loadfolderdata(imagepath, filterfn=isimagefile, loadfn=loadfile)
+    masks = loadfolderdata(maskpath, filterfn=isimagefile, loadfn=f -> loadmask(f, classes))
+    nobs(images) == nobs(masks) || error("Expected the same number of images and masks, but found $(nobs(images)) images and $(nobs(masks)) masks")
+    nobs(images) > 0 || error("No images or masks found in folders $imagepath and $maskpath")
+
+    blocks = Image{2}(), Mask{2}(classes)
+    return (images, masks), blocks
+end
+
+recipeblocks(::Type{ImageSegmentationFolders}) = Tuple{Image{2}, Mask{2}}