Torchvision v2.Grayscale not returning 1 channel

It looks like torchvision v2.Grayscale is not returning 1 channel, even if passing num_output_channels=1:

import cv2
from torchvision.transforms import v2

img = cv2.imread('/tmp/test.png')
gray = v2.Grayscale(num_output_channels=1)
out_img = gray(img)
# (339, 510, 3)

Does anybody knows a solution? Using functional.rgb_to_grayscale works as intended (outputs only 1 channel), however that function doesn’t work with numpy arrays, only PIL or tensors. Which means I need to do a whole lot of conversions for it to work:

import numpy as np
from PIL import Image
from torchvision.transforms.functional import rgb_to_grayscale

# my inputs are float32, but PIL can't read it
img = Image.fromarray(img_gt.astype(np.uint8))
img = rgb_to_grayscale(img)
img = np.array(img, dtype=np.float32)
# my other functions expect it to output (h, w, 1), but rgb_to_grayscale outputs (h, w) only
img = np.expand_dims(img, axis=-1)

Any help is appreciated :+1:

Indeed it works with PIL.Images but seems to fail with numpy.arrays (loaded via OpenCV).
CC @pmeier