io.write_png saves RGB image when given greyscale tensor

[rharish101]

🐛 Bug

When giving a 3D 1xHxW uint8 tensor to torchvision.io.write_png, it saves it as a coloured PNG instead of a greyscale PNG. torchvision.io.write_jpeg works correctly.

To Reproduce

Steps to reproduce the behavior:

1. Create a 3D 1x64x64 uint8 tensor.
2. Give it to torchvision.io.write_png.
3. Check the saved PNG's attributes, e.g. using identify from ImageMagick. It should say sRGB instead of greyscale.

import torch
from torchvision.io import write_png

img = (torch.rand(1, 64, 64) * 255).byte()
write_png(img, "test.png")

Expected behavior

The saved PNG should be greyscale.

Environment

Please copy and paste the output from our
environment collection script
(or fill out the checklist below manually).

You can get the script and run it with:

wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py

PyTorch version: 1.8.0
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Arch Linux (x86_64)
GCC version: (GCC) 10.2.0
Clang version: 11.1.0
CMake version: Could not collect

Python version: 3.9 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 11.2.142
GPU models and configuration: GPU 0: GeForce GTX 960M
Nvidia driver version: 460.56
cuDNN version: Probably one of the following:
/usr/lib/libcudnn.so.8.1.0
/usr/lib/libcudnn_adv_infer.so.8.1.0
/usr/lib/libcudnn_adv_train.so.8.1.0
/usr/lib/libcudnn_cnn_infer.so.8.1.0
/usr/lib/libcudnn_cnn_train.so.8.1.0
/usr/lib/libcudnn_ops_infer.so.8.1.0
/usr/lib/libcudnn_ops_train.so.8.1.0
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.5
[pip3] torch==1.8.0
[pip3] torchvision==0.9.0
[conda] Could not collect

[fmassa]

Looks like a bug to me, thanks for spotting this @rharish101 !

I would have expected this to work normally given that we specify grayscale in

vision/torchvision/csrc/io/image/cpu/encode_png.cpp

Line 131 in 551193b

auto color_type = PNG_COLOR_TYPE_GRAY ? channels == 1 : PNG_COLOR_TYPE_RGB;

, so this might require further investigation

[fmassa]

@andfoy could you have a look?

[andfoy]

Sure, I'll try to reproduce it and come up with a solution!

[rharish101]

vision/torchvision/csrc/io/image/cpu/encode_png.cpp

Line 131 in 551193b

auto color_type = PNG_COLOR_TYPE_GRAY ? channels == 1 : PNG_COLOR_TYPE_RGB;

@fmassa Sorry, I'm not that experienced in C++, but shouldn't it be:

    auto color_type = channels == 1 ? PNG_COLOR_TYPE_GRAY : PNG_COLOR_TYPE_RGB;

This article says that the condition should be before the ?.

[fmassa]

@rharish101 oh you are right, the ternary operator was not implemented correctly there and changing this will probably fix the issue