How to save an image in YCbCr mode

One4ever · August 1, 2021, 5:53pm

Hi,
I have a 3D tensor of an image in YCbCr color space instead of RGB.
I want to save the image using torchvision.utils.save_image, how can I set the mode to YCbCr? Are there any other ways to do this?
Thanks

tom · August 1, 2021, 7:03pm

So under the hood, save_image does is lines of code:

github.com

pytorch/vision/blob/a839796328cf4f789c9de5da0b3367b742b5a00c/torchvision/utils.py#L132-L136

    
      
          grid = make_grid(tensor, **kwargs)
          # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
          ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
          im = Image.fromarray(ndarr)
          im.save(fp, format=format)

First it uses make_grid which will arrange a batch of tensors in a grid, giving you a c x h x w tensor even if you passed in b x c x h x w. You may or may not need this.
Then it maps from 0…1 to 0…255, clamps, permutes to h x w x c and converts to a numpy uint8 array. You would want that line, too.
Third it makes a PIL image from the numpy array. Here you can add a mode argument to the call to PIL.Image.fromarray.
Then it saves the PIL image, you want that, too.

It’s a bit tedious, but probably not too much.
Or you could propose a patch to torchvision to add an optional mode argument to save_image.

Best regards

Thomas

One4ever · August 2, 2021, 1:27am

Thanks for your accurate and complete answer