View Source Rotating Image

setup

Setup

Mix.install(
  [
    {:req, "~> 0.2.0"},
    {:kino, "~> 0.5.1"},
    {:exla, "~> 0.2"},
    {:stb_image, "~> 0.5"}
  ],
  config: [
    nx: [default_backend: EXLA.Backend]
  ]
)

preprocessing

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

  # there is no pi const in Nx yet.
  @pi 3.14159

  defp calculate_new_side_length(side1, side2, angle, radians?) do
    angle = if radians?, do: angle, else: angle |> Nx.divide(180) |> Nx.multiply(@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

Rotating image

defmodule Rotate do
  import Nx.Defn
  @num_channels 3
  # 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(Nx.divide(k, x_image), {:u, 32})
    j = rem(k, x_image)

    y =
      orig_ctr_height
      |> Nx.add(1)
      |> Nx.add(i)
      |> then(&Nx.subtract(y_image, &1))

    x =
      orig_ctr_width
      |> Nx.add(1)
      |> Nx.add(j)
      |> then(&Nx.subtract(x_image, &1))

    tangent =
      angle
      |> Nx.divide(2)
      |> Nx.tan()

    new_x =
      y
      |> Nx.multiply(tangent)
      |> then(&Nx.subtract(x, &1))
      |> Nx.round()

    new_y =
      angle
      |> Nx.sin()
      |> Nx.multiply(new_x)
      |> Nx.add(y)
      |> Nx.round()

    new_x =
      new_y
      |> Nx.multiply(tangent)
      |> then(&Nx.subtract(new_x, &1))
      |> Nx.round()

    new_y = Nx.subtract(new_ctr_height, new_y)
    new_x = Nx.subtract(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
    coordinate
    |> Nx.add(1)
    |> Nx.divide(2)
    |> Nx.subtract(1)
    |> Nx.round()
  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

  # there is no pi const in Nx yet.
  @pi 3.14159

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

  defn rotate_image_by_angle(image, angle, new_image, radians? \\ true) do
    # if radians? is set to false than assuming that angle is given in degrees
    angle = if radians?, do: angle, else: angle |> Nx.divide(180) |> Nx.multiply(@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

Example of usage

num_channels = 3

dir = File.cwd!()

default_input =
  "https://www.researchgate.net/publication/259521525/figure/fig8/AS:268024322195479@1440913386352/Original-standard-test-image-of-Peppers.png"

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

# set options
radians? = false
angle = 70

# 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)
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)