Rotating Image

Mix.install(
  [
    {:nx, github: "elixir-nx/nx", override: true, sparse: "nx"},
    {:req, "~> 0.3.5"},
    {:kino, "~> 0.8.1"},
    {:exla, "~> 0.5"},
    {:stb_image, "~> 0.6"}
  ],
  config: [
    nx: [default_backend: EXLA.Backend]
  ]
)

Preprocessing

> Note: if you want to use TPUs/GPUs, you must set the XLA_TARGET environment variable. > See XLA_TARGET.

Preprocessing a binary to the tensor representation and calculating shapes of the rotated image.

defmodule Preprocess do
  def image_to_tensor(file_or_url) do
    case URI.parse(file_or_url) do
      %URI{scheme: nil} ->
        img = StbImage.read_file!(file_or_url)
        Nx.from_binary(img.data, {:u, 8}) |> Nx.reshape(img.shape)

      _ ->
        %{body: binary_image} = Req.get!(file_or_url)
        img = StbImage.read_binary!(binary_image)
        Nx.from_binary(img.data, {:u, 8}) |> Nx.reshape(img.shape)
    end
  end

  defp calculate_new_side_length(side1, side2, angle, radians?) do
    angle =
      if radians?, do: angle, else: angle |> Nx.divide(180) |> Nx.multiply(Nx.Constants.pi())

    sine = Nx.sin(angle)
    cosine = Nx.cos(angle)

    side1
    |> Nx.multiply(cosine)
    |> Nx.abs()
    |> Nx.add(Nx.abs(Nx.multiply(side2, sine)))
    |> Nx.add(1)
    |> Nx.to_number()
    |> round()
  end

  def get_new_height_and_width(height, width, angle, radians? \\ true) do
    new_height = calculate_new_side_length(height, width, angle, radians?)
    new_width = calculate_new_side_length(width, height, angle, radians?)
    {new_height, new_width}
  end
end

Rotating image

defmodule Rotate do
  import Nx.Defn
  @num_channels 4
  # using three shears to avoid aliasing
  # https://datagenetics.com/blog/august32013/index.html
  defnp calculate_new_positions(
          {y_image, x_image, _},
          orig_ctr_height,
          orig_ctr_width,
          new_ctr_height,
          new_ctr_width,
          angle
        ) do
    k = Nx.iota({y_image * x_image})
    i = Nx.as_type(k / x_image, {:u, 32})
    j = Nx.remainder(k, x_image)

    y = y_image - (orig_ctr_height + i + 1)

    x = x_image - (orig_ctr_width + j + 1)

    tangent = Nx.tan(angle / 2)

    new_x = Nx.round(x - y * tangent)

    new_y = Nx.round(Nx.sin(angle) * new_x + y)

    new_x = Nx.round(new_x - new_y * tangent)

    new_y = new_ctr_height - new_y
    new_x = new_ctr_width - new_x

    new_y = Nx.as_type(Nx.round(new_y), {:u, 32})
    new_x = Nx.as_type(Nx.round(new_x), {:u, 32})

    Nx.stack([new_y, new_x], axis: 1)
  end

  defnp calculate_ctr(coordinate) do
    Nx.round((coordinate + 1) / 2 - 1)
  end

  defnp preprocess_position(pos, n) do
    Nx.concatenate(
      [
        Nx.new_axis(pos, 1) |> Nx.tile([1, @num_channels, 1]),
        Nx.iota({n, @num_channels, 1}, axis: 1)
      ],
      axis: 2
    )
    |> Nx.reshape({n * @num_channels, 3})
  end

  # new_image is a tensor with dims of rotated image filled with zeros

  defn rotate_image_by_angle(image, angle, new_image, radians?) do
    # if radians? is set to false than assuming that angle is given in degrees
    angle = if radians?, do: angle, else: angle / 180 * Nx.Constants.pi()

    {height, width, _} = image.shape
    {new_height, new_width, _} = new_image.shape

    orig_ctr_height = calculate_ctr(height)
    orig_ctr_width = calculate_ctr(width)

    new_ctr_height = calculate_ctr(new_height)
    new_ctr_width = calculate_ctr(new_width)

    pos =
      calculate_new_positions(
        image.shape,
        orig_ctr_height,
        orig_ctr_width,
        new_ctr_height,
        new_ctr_width,
        angle
      )

    {n, 2} = Nx.shape(pos)

    preprocessed_pos = preprocess_position(pos, n)
    image = image |> Nx.reshape({n * @num_channels})

    Nx.indexed_add(new_image, preprocessed_pos, image)
  end
end

Example of usage

num_channels = 4

dir = File.cwd!()

default_input =
  "https://upload.wikimedia.org/wikipedia/commons/7/7f/Earth_fluent_design_reflective_surface_icon.png"

input = Kino.Input.text("Image to rotate", default: default_input)
Kino.render(input)
path_to_file = Kino.Input.read(input)

# set options
radians? = 0
angle = 90

# rotate image by 70 degrees to the right
image = Preprocess.image_to_tensor(path_to_file)
{height, width, _} = image.shape
{new_height, new_width} = Preprocess.get_new_height_and_width(height, width, angle, radians?)
new_image = Nx.broadcast(Nx.tensor([0], type: {:u, 8}), {new_height, new_width, num_channels})
rotate_image = EXLA.jit(&Rotate.rotate_image_by_angle/4)
rotated_image = rotate_image.(image, angle, new_image, radians?)

# convert the tensor back to stb_image and render it as png
img = StbImage.from_nx(rotated_image)
content = StbImage.to_binary(img, :png)

# display the rotated image
Kino.Image.new(content, :png)