Add 3d rotations (#93)

* Add 3D rotations

This adds 3D support to the Rotate transformation, and introduces
RotateX, RotateY, and RotateZ for 3D rotations around the x, y, and z
axes.

* Add check to `getprojection` for `ComposedProjectiveTransform`

Ensure the number of transforms equals the number of random states.
Previously, the projection would silently be calculated only for the
number of transforms equal to the number of random states.

* Add RotateX, RotateY, and RotateZ to documentation

* Clean up documentation, add 3D rotation to gallery

* Use `@docs; canonical=false` for documentation outside of ref.md.
  This allows documentation to exist on a second page. Previously,
  documentation was split between multiple pages. Some were in ref.md,
  and some were in other pages (transformations.md, buffering.md). Now,
  all documentation is at least in ref.md, but can also be duplicated
  on other pages with `canonical=false`.
* Add missing documentation to ref.md, and transformations.md.
* Enable the visualizations in gallery.md, and clean up the code a bit.
* Add visualizations in gallery.md for 3D rotations.

Author: paulnovo

docs/src/buffering.md

@@ -22,7 +22,7 @@ buffer = apply(tfm, item)  # uses apply! internally
 
 Since `Buffered` only stores one buffer, you may run into problems when using it in a multi-threading context where different threads invalidate the buffer before it can be used. In that case, you can use [`DataAugmentation.BufferedThreadsafe`](@ref), a version of `Buffered` that keeps a separate buffer for every thread. 
 
-```@docs
+```@docs; canonical=false
 DataAugmentation.Buffered
 DataAugmentation.BufferedThreadsafe
 ```

docs/src/projective/gallery.md

@@ -7,15 +7,7 @@ the keypoint-based items [`Keypoints`](@ref), [`Polygon`](@ref), and [`BoundingB
 
 Let's take this picture of a light house:
 
-```@setup deps
-using DataAugmentation
-using MosaicViews
-using Images
-using TestImages
-using StaticArrays
-```
-
-```@example
+```@example deps
 using DataAugmentation
 using MosaicViews
 using Images
@@ -26,11 +18,6 @@ imagedata = testimage("lighthouse")
 imagedata = imresize(imagedata, ratio = 196 / size(imagedata, 1))
 ```
 
-```@example deps
-imagedata = testimage("lighthouse")
-imagedata
-```
-
 To apply a transformation `tfm` to it, wrap it in
 `Image`, apply the transformation and unwrap it using [`itemdata`](@ref):
 
@@ -74,59 +61,29 @@ apply(tfm, (image, bbox)) |> showitems
 
     Of course, you have to create a 3-dimensional transformation, i.e. `CenterCrop((128, 128, 128))` instead of `CenterCrop((128, 128))`.
 
-## Gallery
-```julia
-function showtransform(tfm, item, n = 8; ncol = 4)
-    return mosaicview(
-        [showitems(apply(tfm, item)) for _ in 1:n],
-        fillvalue = RGBA(1, 1, 1, 0),
-        npad = 8,
-        rowmajor = true,
-        ncol = ncol)
-end
-
-
-function showtransforms(tfms, item; ncol = length(tfms))
-    return mosaicview(
-        [parent(showitems(apply(tfm, item))) for tfm in tfms],
-        fillvalue = RGBA(1, 1, 1, 0),
-        npad = 8,
-        rowmajor = true,
-        ncol = ncol)
-end
-
-nothing # hide
-```
-
-### [`RandomResizeCrop`](@ref)`(sz)`
+## [`RandomResizeCrop`](@ref)`(sz)`
 
 Resizes the sides so that one of them is no longer than `sz` and crops a region of size `sz` *from a random location*.
 
-```julia
+```@example deps
 tfm = RandomResizeCrop((128, 128))
+showgrid([apply(tfm, (image, bbox)) for _ in 1:6]; ncol=6, npad=8)
 ```
 
-```julia
-o = showtransform(tfm, (image, bbox), 6, ncol=6)
-```
-
-### [`CenterResizeCrop`](@ref)
+## [`CenterResizeCrop`](@ref)
 
 Resizes the sides so that one of them is no longer than `sz` and crops a region of size `sz` *from the center*.
 
-```julia
+```@example deps
 tfm = CenterResizeCrop((128, 128))
+showgrid([apply(tfm, (image, bbox))]; ncol=6, npad=8)
 ```
 
-```julia
-o = showtransform(tfm, (image, bbox), 1)
-```
-
-### [`Crop`](@ref)`(sz[, from])`
+## [`Crop`](@ref)`(sz[, from])`
 
 Crops a region of size `sz` from the image, *without resizing* the image first.
 
-```julia
+```@example deps
 using DataAugmentation: FromOrigin, FromCenter, FromRandom
 tfms = [
     Crop((128, 128), FromOrigin()),
@@ -136,36 +93,43 @@ tfms = [
     Crop((128, 128), FromRandom()),
     Crop((128, 128), FromRandom()),
 ]
+showgrid([apply(tfm, (image, bbox)) for tfm in tfms]; ncol=6, npad=8)
 ```
 
-```julia
-o = showtransforms(tfms, (image, bbox))
-```
-
-### [`FlipX`](@ref), [`FlipY`](@ref), [`Reflect`](@ref)
+## [`FlipX`](@ref), [`FlipY`](@ref), [`Reflect`](@ref)
 
 Flip the data on the horizontally and vertically, respectively. More generally, reflect around an angle from the x-axis.
 
-```julia
+```@example deps
 tfms = [
     FlipX(),
     FlipY(),
     Reflect(30),
 ]
+showgrid([apply(tfm, (image, bbox)) for tfm in tfms]; ncol=6, npad=8)
 ```
 
-```julia
-o = showtransforms(tfms, (image, bbox))
-```
-
-### [`Rotate`](@ref)
+## [`Rotate`](@ref), [`RotateX`](@ref), [`RotateY`](@ref), [`RotateZ`](@ref)
 
-Rotate counter-clockwise by an angle.
+Rotate a 2D image counter-clockwise by an angle.
 
-```julia
+```@example deps
 tfm = Rotate(20) |> CenterCrop((256, 256))
+showgrid([apply(tfm, (image, bbox)) for _ in 1:6]; ncol=6, npad=8)
 ```
 
-```julia
-o = showtransform(tfm, (image, bbox), 1)
+Rotate also works with 3D images in addition to 3D specific transforms RotateX, RotateY, and RotateZ.
+
+```@example deps
+image3D = Image([RGB(i, j, k) for i=0:0.01:1, j=0:0.01:1, k=0:0.01:1])
+tfms = [
+    Rotate(20, 30, 40),
+    Rotate{3}(45),
+    RotateX(45),
+    RotateY(45),
+    RotateZ(45),
+]
+transformed = [apply(tfm, image3D) |> itemdata for tfm in tfms]
+slices = [Image(parent(t[:, :, 50])) for t in transformed]
+showgrid(slices; ncol=6, npad=8)
 ```

docs/src/ref.md

@@ -1,25 +1,44 @@
 # Reference
 ```@docs
+AdjustBrightness
+AdjustContrast
 BoundingBox
+CenterCrop
+CenterResizeCrop
+Crop
+FlipX
+FlipY
 Image
 Keypoints
 MaskBinary
 MaskMulti
+Maybe
+OneOf
+PermuteDims
 Polygon
+RandomCrop
 RandomResizeCrop
-CenterResizeCrop
-Crop
+Reflect
 Rotate
+RotateX
+RotateY
+RotateZ
+ScaleKeepAspect
+ScaleRatio
+WarpAffine
 itemdata
 showitems
 
 DataAugmentation.AbstractArrayItem
 DataAugmentation.AbstractItem
 DataAugmentation.ArrayItem
+DataAugmentation.Buffered
+DataAugmentation.BufferedThreadsafe
 DataAugmentation.Categorify
 DataAugmentation.ComposedProjectiveTransform
 DataAugmentation.FillMissing
 DataAugmentation.Identity
+DataAugmentation.ImageToTensor
 DataAugmentation.Item
 DataAugmentation.ItemWrapper
 DataAugmentation.MapElem

docs/src/transformations.md

@@ -25,20 +25,26 @@ Projective transformations include:
 
 Affine transformations are a subgroup of projective transformations that can be composed very efficiently: composing two affine transformations results in another affine transformation. Affine transformations can represent translation, scaling, reflection and rotation. Available `Transform`s are:
 
-```@docs
-ScaleRatio
-ScaleKeepAspect
+```@docs; canonical=false
 FlipX
 FlipY
 Reflect
+Rotate
+RotateX
+RotateY
+RotateZ
+ScaleKeepAspect
+ScaleFixed
+ScaleRatio
 WarpAffine
+Zoom
 ```
 
 ## Crops
 
 To get a cropped result, simply `compose` any `ProjectiveTransform` with
 
-```@docs
+```@docs; canonical=false
 CenterCrop
 RandomCrop
 ```
@@ -47,17 +53,17 @@ RandomCrop
 
 DataAugmentation.jl currently supports the following color transformations for augmentation:
 
-```@docs
-AdjustContrast
+```@docs; canonical=false
 AdjustBrightness
+AdjustContrast
 ```
 
 # Stochastic transformations
 When augmenting data, it is often useful to apply a transformation only with some probability or choose from a set of transformations. Unlike in other data augmentation libraries like *albumentations*, in DataAugmentation.jl you can use wrapper transformations for this functionality.
 
 - [`Maybe`](@ref)`(tfm, p = 0.5)` applies a transformation with probability `p`; and
 - [`OneOf`](@ref)`([tfm1, tfm2])` randomly selects a transformation to apply.
-```@docs
+```@docs; canonical=false
 Maybe
 OneOf
 ```
@@ -72,6 +78,6 @@ titems = [apply(tfm, item) for _ in 1:8]
 showgrid(titems; ncol = 4, npad = 16)
 ```
 
-```@docs
+```@docs; canonical=false
 DataAugmentation.ImageToTensor
 ```

src/DataAugmentation.jl

@@ -70,6 +70,9 @@ export Item,
     RandomCrop,
     ScaleFixed,
     Rotate,
+    RotateX,
+    RotateY,
+    RotateZ,
     RandomResizeCrop,
     CenterResizeCrop,
     Buffered,

src/projective/affine.jl

@@ -120,36 +120,92 @@ end
 
 """
     Rotate(γ)
-    Rotate(γs)
+    Rotate(distribution)
+    Rotate(α, β, γ)
+    Rotate(α_distribution, β_distribution, γ_distribution)
 
-Rotate 2D spatial data around the center by an angle chosen at
-uniformly from [-γ, γ], an angle given in degrees.
+Rotate spatial data around its center. Rotate(γ) is a 2D rotation
+by an angle chosen uniformly from [-γ, γ], an angle given in degrees.
+Rotate(α, β, γ) is a 3D rotation by angles chosen uniformly from
+[-α, α], [-β, β], and [-γ, γ], for X, Y, and Z rotations.
 
 You can also pass any `Distributions.Sampleable` from which the
 angle is selected.
 
 ## Examples
 
 ```julia
-tfm = Rotate(10)
-```
+tfm2d = Rotate(10)
+apply(tfm2d, Image(rand(Float32, 16, 16)))
 
+tfm3d = Rotate(10, 20, 30)
+apply(tfm3d, Image(rand(Float32, 16, 16, 16)))
+```
 """
-struct Rotate{S<:Sampleable} <: ProjectiveTransform
+struct Rotate{N, R, S<:Sampleable} <: ProjectiveTransform
     dist::S
 end
-Rotate(γ) = Rotate(Uniform(-abs(γ), abs(γ)))
+
+Rotate{N, R}(s::S) where {N, R, S} = Rotate{N, R, S}(s)
+
+function Rotate{N, R}(s::S) where {N, R, S<:Number}
+    if s == 0
+        return Identity()
+    end
+    Rotate{N, R, Uniform}(Uniform(-abs(s), abs(s)))
+end
+
+Rotate(s) = Rotate{2, Type{RotMatrix{2, Float64}}}(s)
+Rotate(sx, sy, sz) = RotateX(sx) |> RotateY(sy) |> RotateZ(sz)
+Rotate{3}(s) = Rotate(s, s, s)
+Rotate{2}(s) = Rotate(s)
+
+"""
+    RotateX(γ)
+    RotateX(distribution)
+
+X-Axis rotation of 3D spatial data around the center by an angle chosen
+uniformly from [-γ, γ], an angle given in degrees.
+
+You can also pass any `Distributions.Sampleable` from which the
+angle is selected.
+"""
+RotateX(s) = Rotate{3, Type{RotX{Float64}}}(s)
+
+"""
+    RotateY(γ)
+    RotateY(distribution)
+
+Y-Axis rotation of 3D spatial data around the center by an angle chosen
+uniformly from [-γ, γ], an angle given in degrees.
+
+You can also pass any `Distributions.Sampleable` from which the
+angle is selected.
+"""
+RotateY(s) = Rotate{3, Type{RotY{Float64}}}(s)
+
+"""
+    RotateZ(γ)
+    RotateZ(distribution)
+
+Z-Axis rotation of 3D spatial data around the center by an angle chosen
+uniformly from [-γ, γ], an angle given in degrees.
+
+You can also pass any `Distributions.Sampleable` from which the
+angle is selected.
+"""
+RotateZ(s) = Rotate{3, Type{RotZ{Float64}}}(s)
 
 getrandstate(tfm::Rotate) = rand(tfm.dist)
 
 function getprojection(
-        tfm::Rotate,
-        bounds::Bounds{2};
-        randstate = getrandstate(tfm))
+        tfm::Rotate{N, Type{R}, S},
+        bounds::Bounds{N};
+        randstate = getrandstate(tfm)) where {N, R, S}
     γ = randstate
-    middlepoint = SVector{2, Float32}(mean.(bounds.rs))
+    middlepoint = SVector{N, Float32}(mean.(bounds.rs))
     r = γ / 360 * 2pi
-    return recenter(RotMatrix(convert(Float32, r)), middlepoint)
+    return recenter(RotMatrix{N, Float32}(R(convert(Float64, r))), middlepoint)
 end
 
 

src/projective/compose.jl

@@ -40,6 +40,7 @@ function getprojection(
         composed::ComposedProjectiveTransform,
         bounds;
         randstate = getrandstate(composed))
+    @assert length(composed.tfms) == length(randstate)
     P = CoordinateTransformations.IdentityTransformation()
     for (tfm, r) in zip(composed.tfms, randstate)
         P_tfm = getprojection(tfm, bounds; randstate = r)