Make Flips ND compatible and fix ND ScaleKeepAspect, ScaleFixed and PinOrigin

docs/literate/projective/gallery.md

@@ -147,15 +147,15 @@ tfms = [
 o = showtransforms(tfms, (image, bbox))
 ```
 
-### [`FlipX`](#), [`FlipY`](#), [`Reflect`](#)
+### [`FlipX`](#), [`FlipY`](#), [`FlipZ`](#), [`FlipDim`](#), [`Reflect`](#)
 
 Flip the data on the horizontally and vertically, respectively. More generally, reflect around an angle from the x-axis.
 
 {cell=main result=false}
 ```julia
 tfms = [
-    FlipX(),
-    FlipY(),
+    FlipX(2),
+    FlipY(2),
     Reflect(30),
 ]
 ```

docs/literate/projective/intro.md

@@ -19,7 +19,7 @@ We can break down most augmentation used in practive into a single (possibly sto
 As an example, consider an image augmentation pipeline: A random horizontal flip, followed by a random resized crop. The latter resizes and crops (irregularly sized) images to a common size without distorting the aspect ratio.
 
 ```julia
-Maybe(FlipX()) |> RandomResizeCrop((h, w))
+Maybe(FlipX(2)) |> RandomResizeCrop((h, w))
 ```
 
 Let's pull apart the steps involved. 

docs/literate/projective/usage.md

@@ -3,7 +3,7 @@
 Using projective transformations is as simple as any other transformations. Simply `compose` them:
 
 ```julia
-Rotate(-10:10) |> ScaleRatio(0.7:0.1:1.2) |> FlipX() |> Crop((128, 128))
+Rotate(-10:10) |> ScaleRatio(0.7:0.1:1.2) |> FlipX(2) |> Crop((128, 128))
 ```
 
 The composition will automatically create a single projective transformation and evaluate only the cropped area.
@@ -14,7 +14,7 @@ Affine transformations are a subgroup of projective transformations that can be
 
 - [`ScaleRatio`](#), [`ScaleKeepAspect`](#)
 - [`Rotate`](#)
-- [`FlipX`](#), [`FlipY`](#), [`Reflect`](#)
+- [`FlipX`](#), [`FlipY`](#), [`FlipZ`](#), [`FlipDim`](#), [`Reflect`](#)
 - [`WarpAffine`](#)
 
 ## Crops

docs/literate/stochastic.md

@@ -13,7 +13,7 @@ Let's say we have an image classification dataset. For most datasets, horizontal
 ```julia
 using DataAugmentation, TestImages
 item = Image(testimage("lighthouse"))
-tfm = Maybe(FlipX())
+tfm = Maybe(FlipX(2))
 titems = [apply(tfm, item) for _ in 1:8]
 showgrid(titems; ncol = 4, npad = 16)
 ```

src/DataAugmentation.jl

@@ -82,6 +82,7 @@ export Item,
     WarpAffine,
     FlipX,
     FlipY,
+    FlipZ,
     PinOrigin,
     AdjustBrightness,
     AdjustContrast,

src/projective/affine.jl

@@ -39,15 +39,15 @@ struct ScaleKeepAspect{N} <: ProjectiveTransform
 end
 
 
-function getprojection(scale::ScaleKeepAspect{N}, bounds; randstate = nothing) where N
+function getprojection(scale::ScaleKeepAspect{N}, bounds::Bounds{N}; randstate = nothing) where N
     # If no scaling needs to be done, return a noop transform
     scale.minlengths == length.(bounds.rs) && return IdentityTransformation()
 
     # Offset `minlengths` by 1 to avoid black border on one side
     ratio = maximum((scale.minlengths .+ 1) ./ length.(bounds.rs))
     upperleft = SVector{N, Float32}(minimum.(bounds.rs)) .- 0.5
     P = scaleprojection(Tuple(ratio for _ in 1:N))
-    if upperleft != SVector(0, 0)
+    if any(upperleft .!= 0)
         P = P ∘ Translation((Float32.(P(upperleft)) .+ 0.5f0))
     end
     return P
@@ -75,11 +75,11 @@ struct ScaleFixed{N} <: ProjectiveTransform
 end
 
 
-function getprojection(scale::ScaleFixed, bounds; randstate = nothing)
+function getprojection(scale::ScaleFixed, bounds::Bounds{N}; randstate = nothing) where N
     ratios = (scale.sizes .+ 1) ./ length.(bounds.rs)
-    upperleft = SVector{2, Float32}(minimum.(bounds.rs)) .- 1
+    upperleft = SVector{N, Float32}(minimum.(bounds.rs)) .- 1
     P = scaleprojection(ratios)
-    if upperleft != SVector(0, 0)
+    if any(upperleft .!= 0)
         P = P  ∘ Translation(-upperleft)
     end
     return P
@@ -88,7 +88,7 @@ end
 
 function projectionbounds(tfm::ScaleFixed{N}, P, bounds::Bounds{N}; randstate = nothing) where N
     bounds_ = transformbounds(bounds, P)
-    return offsetcropbounds(tfm.sizes, bounds_, (1., 1.))
+    return offsetcropbounds(tfm.sizes, bounds_, ntuple(_ -> 1., N))
 end
 
 """
@@ -167,7 +167,7 @@ struct Reflect <: ProjectiveTransform
 end
 
 
-function getprojection(tfm::Reflect, bounds; randstate = getrandstate(tfm))
+function getprojection(tfm::Reflect, bounds::Bounds{2}; randstate = getrandstate(tfm))
     r = tfm.γ / 360 * 2pi
     return centered(LinearMap(reflectionmatrix(r)), bounds)
 end
@@ -178,24 +178,45 @@ end
 Transform `P` so that is applied around the center of `bounds`
 instead of the origin
 """
-function centered(P, bounds::Bounds{2})
+function centered(P, bounds::Bounds{N}) where N
     upperleft = minimum.(bounds.rs)
     bottomright = maximum.(bounds.rs)
 
-    midpoint = SVector{2, Float32}((bottomright .- upperleft) ./ 2) .+ SVector{2, Float32}(.5, .5)
+    midpoint = SVector{N, Float32}((bottomright .- upperleft) ./ 2) .+ .5f0
     return recenter(P, midpoint)
 end
 
-
-FlipX() = Reflect(180)
-FlipY() = Reflect(90)
-
 function reflectionmatrix(r)
     A = SMatrix{2, 2, Float32}(cos(2r), sin(2r), sin(2r), -cos(2r))
     return round.(A; digits = 12)
 end
 
 
+"""
+    FlipDim{N}()
+Reflect `N` dimensional data along the axis of dimension `dim`. Must satisfy 1 <= `dim` <= `N`.
+## Examples
+```julia
+tfm = FlipDim{2}(1)
+```
+"""
+struct FlipDim{N} <: ProjectiveTransform
+    dim::Int
+    FlipDim{N}(dim) where {N} = 1 <= dim <= N ? new{N}(dim) : error("invalid dimension")
+end
+# 2D images use (r, c) = (y, x) convention
+FlipX(N) = FlipDim{N}(N==2 ? 2 : 1)
+FlipY(N) = FlipDim{N}(N==2 ? 1 : 2)
+FlipZ(N) = FlipDim{N}(3)
+
+function getprojection(tfm::FlipDim{N}, bounds::Bounds{N}; randstate = nothing) where N
+    arr = 1I(N)
+    arr[tfm.dim, tfm.dim] = -1
+    M = SMatrix{N, N, Float32}(arr)
+    return DataAugmentation.centered(LinearMap(M), bounds)
+end
+
+
 """
     PinOrigin()
 
@@ -213,8 +234,8 @@ at one.
 """
 struct PinOrigin <: ProjectiveTransform end
 
-function getprojection(::PinOrigin, bounds; randstate = nothing)
-    p = (-SVector{2, Float32}(minimum.(bounds.rs))) .+ 1
+function getprojection(::PinOrigin, bounds::Bounds{N}; randstate = nothing) where N
+    p = (-SVector{N, Float32}(minimum.(bounds.rs))) .+ 1
     P = Translation(p)
     return P
 end

test/projective/affine.jl

@@ -142,7 +142,7 @@ include("../imports.jl")
     @testset ExtendedTestSet "`RandomCrop` correct indices" begin
         # Flipping and cropping should be the same as reverse-indexing
         # the flipped dimension
-        tfm = FlipX() |> RandomCrop((64, 64)) |> PinOrigin()
+        tfm = FlipX(2) |> RandomCrop((64, 64)) |> PinOrigin()
         img = rand(RGB, 64, 64)
         item = Image(img)
         titem = apply(tfm, item)
@@ -157,8 +157,8 @@ end
     @testset ExtendedTestSet "2D" begin
         tfms = compose(
             Rotate(10),
-            FlipX(),
-            FlipY(),
+            FlipX(2),
+            FlipY(2),
             ScaleRatio((.8, .8)),
             WarpAffine(0.1),
             Zoom((1., 1.2)),
@@ -177,8 +177,13 @@ end
         )
 
         tfms = compose(
+            FlipX(3),
+            FlipY(3),
+            FlipZ(3),
+            ScaleFixed((30, 40, 50)),
             ScaleRatio((.8, .8, .8)),
             ScaleKeepAspect((12, 10, 10)),
+            Zoom((1., 1.2)),
             RandomCrop((10, 10, 10))
         )
         testprojective(tfms, items)