# Image Rotation Algorithm for CUDA and Zygote
**URL:** https://discourse.julialang.org/t/image-rotation-algorithm-for-cuda-and-zygote/54085
**Category:** GPU
**Tags:** images, cuda, array, zygote
**Created:** [January 27, 2021, 9:50pm UTC](https://discourse.julialang.org/t/image-rotation-algorithm-for-cuda-and-zygote/54085 "2021-01-27T21:50:37Z")
**Posts on this page:** 11
**Page:** 1
### Author:  [@roflmaostc](https://discourse.julialang.org/u/roflmaostc)
#### Post date: [January 27, 2021, 9:50pm UTC](https://discourse.julialang.org/t/image-rotation-algorithm-for-cuda-and-zygote/54085/1 "2021-01-27T21:50:37Z")
Hey,
today I wanted to use a image rotation algorithm to rotate a 3D array around one dimension (so basically many 2D images).
For my applications it needs to be fast and also fully differentiable by Zygote.
I tried ImageTransformations.jl but the algorithms don’t play well with CUDA (some Interpolations.jl errors).
Consequently, I was looking for potential algorithms to implement:
- [Three Shear Algorithm](https://graphicsinterface.org/proceedings/gi1986/gi1986-15/)
- FFT based rotation which is sketched [here](https://computergraphics.stackexchange.com/a/4846). Probably promising because FFTs are fast.
I’m not totally sure how to achieve this because I’m not really familiar with CUDA (last week I bought the first GPU since my old AGP 8x 😃). I would be happy if some people could point me into relevant directions or algorithms.
Or does there even exist an Julia code offering a rotation?
Thanks,
Felix
---
### Author:  [@maxfreu](https://discourse.julialang.org/u/maxfreu)
#### Post date: [February 5, 2021, 4:00pm UTC](https://discourse.julialang.org/t/image-rotation-algorithm-for-cuda-and-zygote/54085/2 "2021-02-05T16:00:44Z")
As far as I know there is nothing out of the box with works with GPU and gradients. As rotations are just a special case of affine transformations, one could take the shot and port the grid sampler ([docs](https://pytorch.org/docs/stable/nn.functional.html#grid-sample), [GPU source](https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/GridSampler.cu)) and the affine grid generator from pytorch. This can handle arbitrary transformations and is used to build spatial transformer networks. However this is _a lot_ of complex code and there must be a julian way to handle this much more elegantly, but I can’t think of any right now.
---
### Author:  [@roflmaostc](https://discourse.julialang.org/u/roflmaostc)
#### Post date: [February 5, 2021, 4:04pm UTC](https://discourse.julialang.org/t/image-rotation-algorithm-for-cuda-and-zygote/54085/3 "2021-02-05T16:04:18Z")
Thanks for your answer!
In the meanwhile, I created a FFT based rotation algorithm for 3D arrays for a single special case.
It’s pretty fast (you need basically 6 `fft(arr, [1]))`) and also fast with Zygote (since the gradient of fft is known).
The code is not public at the moment, and there must be invested even more work to generalize it to for any rotation axis.
But if desired, I could post my special case code here.
---
### Author:  [@maxfreu](https://discourse.julialang.org/u/maxfreu)
#### Post date: [February 5, 2021, 8:37pm UTC](https://discourse.julialang.org/t/image-rotation-algorithm-for-cuda-and-zygote/54085/4 "2021-02-05T20:37:17Z")
Yes, please! 🙂
---
### Author:  [@roflmaostc](https://discourse.julialang.org/u/roflmaostc)
#### Post date: [February 5, 2021, 10:52pm UTC](https://discourse.julialang.org/t/image-rotation-algorithm-for-cuda-and-zygote/54085/5 "2021-02-05T22:52:21Z")
I’ll try to boil it down to a 2D example. Should be easier to understand.
---
### Author:  [@magister-ludi](https://discourse.julialang.org/u/magister-ludi)
#### Post date: [February 6, 2021, 5:58am UTC](https://discourse.julialang.org/t/image-rotation-algorithm-for-cuda-and-zygote/54085/6 "2021-02-06T05:58:23Z")
I am one of the co-authors of [Larkin et al.](https://doi.org/10.1016/S0030-4018(97)00097-7) that is mentioned in the [link](https://computergraphics.stackexchange.com/a/4846) that you provided. My (unregistered) package [Eigenbroetler.jl](https://github.com/magister-ludi/Eigenbroetler.jl) has a Julia implementation of the 2D algorithm. I’m sure it will need changing to suit your needs, but it’s there for you to look at and modify.
---
### Author:  [@roflmaostc](https://discourse.julialang.org/u/roflmaostc)
#### Post date: [February 6, 2021, 10:50am UTC](https://discourse.julialang.org/t/image-rotation-algorithm-for-cuda-and-zygote/54085/7 "2021-02-06T10:50:10Z")
Here is the code for a 2D FFT based rotation.
`fftpos` is extracted from [PhysicalOptics.jl](https://github.com/JuliaPhysics/PhysicalOptics.jl) and I copied it for simplicity.
I didn’t include it here, but it should work pretty much the same with CUDA (tested it last week).
```julia
using FFTW
function shear(arr, Δx)
ĩ = 2
c = eltype(arr)(2π * Δx)
ϕ_1D_shift = c .* rfftfreq(size(arr)[ĩ], one(eltype(arr)))'
ϕ_shift_strength = fftpos(one(eltype(arr)), size(arr)[2])
ϕ_2D = exp.(1im .* ϕ_1D_shift .* ϕ_shift_strength)
arr_ft = rfft(arr, [ĩ])
return irfft(arr_ft .* ϕ_2D, size(arr)[ĩ], [ĩ])
end
function rotate(arr, θ)
α = -tan(θ/2)
β = sin(θ)
arr = shear(arr, α * size(arr)[1])
arr = permutedims(arr, (2,1))
arr = shear(arr, β * size(arr)[1])
arr = permutedims(arr, (2,1))
arr = shear(arr, α * size(arr)[1])
return arr
end
```
Here is a full Pluto example:
> **Full Example**
>
> ```julia
> ### A Pluto.jl notebook ###
> # v0.12.19
>
> using Markdown
> using InteractiveUtils
>
> # This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
> macro bind(def, element)
> quote
> local el = $(esc(element))
> global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
> el
> end
> end
>
> # ╔═╡ 7560fc84-6862-11eb-0f99-bfa76ae8a694
> using Revise, FFTW, FFTResampling, TestImages, Colors, PlutoUI
>
> # ╔═╡ 16e4aa72-6863-11eb-35ad-21c3ba47ebad
> function fftpos(l, N)
> if N % 2 == 0
> dx = l / N
> return range(-l/2, l/2-dx, length=N)
> else
> return range(-l/2, l/2, length=N)
> end
> end
>
> # ╔═╡ 79d641c0-6862-11eb-15db-77d9548980c1
> function shear(arr, Δx)
> ĩ = 2
> c = eltype(arr)(2π * Δx)
>
> ϕ_1D_shift = c .* rfftfreq(size(arr)[ĩ], one(eltype(arr)))'
> ϕ_shift_strength = fftpos(one(eltype(arr)), size(arr)[2])
> ϕ_2D = exp.(1im .* ϕ_1D_shift .* ϕ_shift_strength)
>
> arr_ft = rfft(arr, [ĩ])
>
> return irfft(arr_ft .* ϕ_2D, size(arr)[ĩ], [ĩ])
> end
>
> # ╔═╡ 0323a81a-6864-11eb-1c5f-eb449060c54b
> function rotate(arr, θ)
> α = -tan(θ/2)
> β = sin(θ)
>
> arr = shear(arr, α * size(arr)[1])
> arr = permutedims(arr, (2,1))
> arr = shear(arr, β * size(arr)[1])
> arr = permutedims(arr, (2,1))
> arr = shear(arr, α * size(arr)[1])
>
> return arr
> end
>
> # ╔═╡ 79b6251e-6862-11eb-0260-a71fc759825b
> begin
> img = Float32.(testimage("fabip_gray_256"))
> img_pad = FFTResampling.center_set!(zeros(eltype(img), (400, 400)), img)
> end
>
> # ╔═╡ 4c4f8fde-6866-11eb-3e26-1b5be5dfd4fa
> md"""
> $(@bind θ Slider(-180:180))
> """
>
> # ╔═╡ b20c01c8-6866-11eb-33c1-89d7c45484ed
> img_s = rotate(img_pad, θ / 180 * π)
>
> # ╔═╡ b4b405c6-6866-11eb-1b55-173146524f32
> md"""
>
> $ \theta= $ $(θ)°
> """
>
> # ╔═╡ 95d4846c-6863-11eb-11c5-5568bca3e754
> Gray.(img_s)
>
> # ╔═╡ Cell order:
> # ╠═7560fc84-6862-11eb-0f99-bfa76ae8a694
> # ╠═16e4aa72-6863-11eb-35ad-21c3ba47ebad
> # ╠═79d641c0-6862-11eb-15db-77d9548980c1
> # ╠═0323a81a-6864-11eb-1c5f-eb449060c54b
> # ╠═79b6251e-6862-11eb-0260-a71fc759825b
> # ╠═b20c01c8-6866-11eb-33c1-89d7c45484ed
> # ╠═4c4f8fde-6866-11eb-3e26-1b5be5dfd4fa
> # ╠═b4b405c6-6866-11eb-1b55-173146524f32
> # ╠═95d4846c-6863-11eb-11c5-5568bca3e754
>
> ```
---
### Author:  [@roflmaostc](https://discourse.julialang.org/u/roflmaostc)
#### Post date: [December 12, 2023, 3:54pm UTC](https://discourse.julialang.org/t/image-rotation-algorithm-for-cuda-and-zygote/54085/8 "2023-12-12T15:54:01Z")
Based on KernelAbstractions I finally pulled things together, and published [DiffImageRotation.jl](https://github.com/roflmaostc/DiffImageRotation.jl).
It works with CUDA and has a registered adjoint.
It’s going to be officially released in ~3 days.
---
### Author:  [@maxfreu](https://discourse.julialang.org/u/maxfreu)
#### Post date: [December 14, 2023, 8:46am UTC](https://discourse.julialang.org/t/image-rotation-algorithm-for-cuda-and-zygote/54085/9 "2023-12-14T08:46:34Z")
Haven’t tested it out yet, but nice work, thanks! I have the feeling that this would perfectly fit into NNlib! I think your code would be highly valued there and used by a wider audience. Also a nearest neighbor version would be nice to rotate categorical masks.
---
### Author:  [@roflmaostc](https://discourse.julialang.org/u/roflmaostc)
#### Post date: [December 14, 2023, 9:06am UTC](https://discourse.julialang.org/t/image-rotation-algorithm-for-cuda-and-zygote/54085/10 "2023-12-14T09:06:08Z")
I see, ths also depends on KernelAbstractions.jl already.
I can try to ask what’s required. Not sure how I would handle the multidimensional case though (of course I can reduce always to 2D rotations) because KernelAbstractions does hard coded array indexing.
---
### Author:  [@roflmaostc](https://discourse.julialang.org/u/roflmaostc)
#### Post date: [January 31, 2024, 12:29pm UTC](https://discourse.julialang.org/t/image-rotation-algorithm-for-cuda-and-zygote/54085/11 "2024-01-31T12:29:43Z")
Now finally [merged into NNlib.jl](https://fluxml.ai/NNlib.jl/dev/reference/#Rotation)