Fail to `im_write` after divided by Nx

[zacky1972]

Hi,

My student found the issue by the following script:

Mix.install([:evision])

src_file = "sample.jpg"
div_size = 256
dst_file_ext = Path.extname(src_file)
dst_file_basename = Path.basename(src_file, dst_file_ext)

dst_files =
  Stream.unfold(0, fn counter -> {counter, counter + 1} end)
  |> Stream.map(& "#{dst_file_basename}_#{&1}#{dst_file_ext}")

div_img =
  Evision.imread(src_file)
  |> Evision.Mat.to_nx()
  |> Nx.to_batched(div_size)
  |> Enum.map(& Evision.Mat.from_nx(&1))

Enum.zip(div_img, dst_files)
|> Enum.map(fn {img, dst_file} -> Evision.imwrite(dst_file, img) end)

The writing by imwrite is failed on the following environment on macOS Sonoma 14.2:

% elixir -v
Erlang/OTP 26 [erts-14.1] [source] [64-bit] [smp:12:12] [ds:12:12:10] [async-threads:1] [jit]

Elixir 1.15.7 (compiled with Erlang/OTP 26)

I found Evision.imread(src_file) has three channels, but img in the divided image has only one channel, so I guess it may cause the issue.

[cocoa-xu]

Hi @zacky1972, you can use Evision.Mat.from_nx_2d/1 instead of Evision.Mat.from_nx/1 to convert images back :)

iex> Evision.imread(src_file)
%Evision.Mat{
  channels: 3,
  dims: 2,
  type: {:u, 8},
  raw_type: 16,
  shape: {2160, 2880, 3},
  ref: #Reference<0.708639968.471203865.168169>
}
iex> |> Evision.Mat.to_nx()
#Nx.Tensor<
  u8[2160][2880][3]
  Evision.Backend
  [
    [
      [194, 196, 197],
      [201, 203, 204],
      [196, 198, 199],
      [198, 200, 201],
      [198, 200, 201],
      [195, 197, 198],
      [202, 204, 205],
      [199, 201, 202],
      [204, 206, 207],
      [202, 204, 205],
      [194, 196, 197],
      [192, 194, 195],
      [205, 207, 208],
      [196, 198, 199],
      [201, 203, 204],
      [195, 197, 198],
      [200, 201, ...],
      ...
    ],
    ...
  ]
>
iex> |> Nx.to_batched(div_size)
[
  #Nx.Tensor<
    u8[256][2880][3]
    Evision.Backend
    [
      [
        [194, 196, 197],
        [201, 203, 204],
        [196, 198, 199],
        [198, 200, 201],
        [198, 200, 201],
        [195, 197, 198],
        [202, 204, 205],
        [199, 201, 202],
        [204, 206, 207],
        [202, 204, 205],
        [194, 196, 197],
        [192, 194, 195],
        [205, 207, 208],
        [196, 198, 199],
        [201, 203, 204],
        [195, 197, 198],
        [200, ...],
        ...
      ],
      ...
    ]
  >,
  #Nx.Tensor<
    u8[256][2880][3]
    Evision.Backend
    [
      [
        [214, 218, 219],
        [218, 222, 223],
        [223, 227, 228],
        [225, 229, 230],
        [223, 225, 226],
        [224, 226, 227],
        [227, 229, 230],
        [220, 222, 223],
        [222, 224, 225],
        [215, 217, 218],
        [221, 220, 222],
        [225, 224, 226],
        [225, 224, 226],
        [216, 215, ...],
        ...
      ],
      ...
    ]
  >
]
iex>  |> Enum.map(& Evision.Mat.from_nx_2d(&1))
[
  %Evision.Mat{
    channels: 3,
    dims: 2,
    type: {:u, 8},
    raw_type: 16,
    shape: {256, 2880, 3},
    ref: #Reference<0.708639968.471203863.164793>
  },
  %Evision.Mat{
    channels: 3,
    dims: 2,
    type: {:u, 8},
    raw_type: 16,
    shape: {256, 2880, 3},
    ref: #Reference<0.708639968.471203863.164794>
  },
  ...
]

[zacky1972]

Thank you so much! The student has just succeeded in running the program code.