Repository: pytorch/vision

Issue:
Reading PNG/JPG images into a torch::tensor and saving a torch::tensor to PNG/JPG in C++ without OpenCV 
## 🚀 Feature

After integrating Siv3D with Libtorch (https://github.com/QuantScientist/Siv3DTorch) I am now trying to read and write images from and to Siv3D **in C++, not Python**. The way it works is:

## Motivation
In C++ I need to do the following:
1. An image is read from disk (usually using OpenCV which is easy but I am trying to avoid)
2. The image is converted to torch::tensor
3. A DL model is applied on the tensor
4. A tensor is returned from the model
5. The tensor is converted to an image for display/saving purposes.

This is one example where they used stb_image to this, avoiding the use of OpenCV.
https://github.com/prabhuomkar/pytorch-cpp/blob/master/utils/image_io/src/image_io.cpp

## Pitch

## Alternatives

For reference this is the OpenCV to Libtorch conversion utils which I use, I would like something very similiar:
```
at::Tensor matToTensor(cv::Mat frame, int h, int w, int c) {
    cv::cvtColor(frame, frame, CV_BGR2RGB);
    frame.convertTo(frame, CV_32FC3, 1.0f / 255.0f);
    auto input_tensor = torch::from_blob(frame.data, {1, h, w, c});
    input_tensor = input_tensor.permute({0, 3, 1, 2});

    torch::DeviceType device_type = torch::kCPU;
//    if (torch::cuda::is_available()) {
    device_type = torch::kCUDA;
//    }
    input_tensor = input_tensor.to(device_type);
    return input_tensor;
}

cv::Mat tensorToOpenCv(at::Tensor out_tensor, int h, int w, int c) {
    out_tensor = out_tensor.squeeze().detach().permute({1, 2, 0});
    out_tensor = out_tensor.mul(255).clamp(0, 255).to(torch::kU8);
    out_tensor = out_tensor.to(torch::kCPU);
    cv::Mat resultImg(h, w, CV_8UC3);
    // cv::Mat resultImg(h, w, CV_8UC1);
    std::memcpy((void *) resultImg.data, out_tensor.data_ptr(), sizeof(torch::kU8) * out_tensor.numel());
    return resultImg;
}

```
## Additional context

I found this:"https://github.com/pytorch/vision/blob/5e4a9f6d1a2bf85137f4826dbf76e4f25986f878/torchvision/csrc/cpu/image/readpng_cpu.cpp
however, could not get any useful method out of it. 

Thanks, 






Discussion:
Hi,

We have recently introduced a `decode_png` and `decode_jpeg` functions in C++, which are exposed to Python for image decoding.

Those methods take the raw read bytes from the image, pass them as a Tensor and get the returned result back.
So the recommended approach for reading images in C++ would be the C++ equivalent of
https://github.com/pytorch/vision/blob/2801a76dd91a8b5c0f63e74e47b3697084f3af28/torchvision/io/image.py#L58-L65

We don't yet have C++-only tests for this interface, but we could add them soon now that we have CMake-only tests.

About implementing image saving operations, this would be a good thing to have, but we haven't currently implemented it.
We would be glad to accept PRs adding it (including Python frontend)

cc @andfoy for thoughts.
See this:
https://github.com/koba-jon/pytorch_cpp/issues/6

The author provided two methods based on png++ that I am going to test. 

```
#include <iostream>
#include <torch/torch.h>
#include <png++/png.hpp>

torch::Tensor ConvertRGBintoTensor(png::image<png::rgb_pixel> &image);
png::image<png::rgb_pixel> ConvertTensorintoRGB(torch::Tensor &tensor_);

int main(void){

    // Input PNG-image
    png::image<png::rgb_pixel> imageI("input.png");

    // Convert png::image into torch::Tensor
    torch::Tensor tensor = ConvertRGBintoTensor(imageI);
    std::cout << "C:" << tensor.size(0) << " H:" << tensor.size(1) << " W:" << tensor.size(2) << std::endl;

    // Convert torch::Tensor into png::image
    png::image<png::rgb_pixel> imageO = ConvertTensorintoRGB(tensor);

    // Output PNG-image
    imageO.write("output.png");

    return 0;
}

torch::Tensor ConvertRGBintoTensor(png::image<png::rgb_pixel> &image){
    size_t width = image.get_width();
    size_t height = image.get_height();
    unsigned char *pointer = new unsigned char[width * height * 3];
    for (size_t j = 0; j < height; j++){
        for (size_t i = 0; i < width; i++){
            pointer[j * width * 3 + i * 3 + 0] = image[j][i].red;
            pointer[j * width * 3 + i * 3 + 1] = image[j][i].green;
            pointer[j * width * 3 + i * 3 + 2] = image[j][i].blue;
        }
    }
    torch::Tensor tensor = torch::from_blob(pointer, {image.get_height(), image.get_width(), 3}, torch::kUInt8).clone();  // copy
    tensor = tensor.permute({2, 0, 1});  // {H,W,C} ===> {C,H,W}
    delete[] pointer;
    return tensor;
}

png::image<png::rgb_pixel> ConvertTensorintoRGB(torch::Tensor &tensor_){
    torch::Tensor tensor = tensor_.permute({1, 2, 0});  // {C,H,W} ===> {H,W,C}
    size_t width = tensor.size(1);
    size_t height = tensor.size(0);
    unsigned char *pointer = tensor.data_ptr<unsigned char>();
    png::image<png::rgb_pixel> image(width, height);
    for (size_t j = 0; j < height; j++){
        for (size_t i = 0; i < width; i++){
            image[j][i].red = pointer[j * width * 3 + i * 3 + 0];
            image[j][i].green = pointer[j * width * 3 + i * 3 + 1];
            image[j][i].blue = pointer[j * width * 3 + i * 3 + 2];
        }
    }
    return image;
}
```
We will probably not rely on png++, but only on png, given that it is a very standard library which is present in many distributions.
As I said, we already have loading functionality for jpeg and png, relying on `libjpeg` and `libpng`, and using it from C++ should be straightforward.
I wrote a small lib for it:
https://github.com/QuantScientist/PngTorch
> I wrote a small lib for it:

still has OpenCV as dependency and  
```
target_link_libraries(${PROJECT_NAME} ${TORCH_LIBRARIES}  ${OpenCV_LIBRARIES} "${CMAKE_CURRENT_LIST_DIR}/lib/libpng/libpng16.lib" "${CMAKE_CURRENT_LIST_DIR}/lib/zlib/zlib.lib")
```
https://github.com/QuantScientist/PngTorch/blob/9475c8adc1f2eb1c3eb6534257af65e6b1de5cb0/CMakeLists.txt#L82

> > I wrote a small lib for it:
> 
> still has OpenCV as dependency and
> 
> ```
> target_link_libraries(${PROJECT_NAME} ${TORCH_LIBRARIES}  ${OpenCV_LIBRARIES} "${CMAKE_CURRENT_LIST_DIR}/lib/libpng/libpng16.lib" "${CMAKE_CURRENT_LIST_DIR}/lib/zlib/zlib.lib")
> ```
> 
> https://github.com/QuantScientist/PngTorch/blob/9475c8adc1f2eb1c3eb6534257af65e6b1de5cb0/CMakeLists.txt#L82

Yes, it evaluated to an empty string, but you are right and I removed it. 
I think we could add encoding support for both jpeg an png directly on our C++ frontend, similar to what we have currently for decoding, that works also for the Python frontend

Resolving patch:
diff --git a/torchvision/io/image.py b/torchvision/io/image.py
--- a/torchvision/io/image.py
+++ b/torchvision/io/image.py
@@ -102,3 +102,42 @@ def read_jpeg(path: str) -> torch.Tensor:
         raise ValueError("Expected a non empty file.")
     data = torch.from_file(path, dtype=torch.uint8, size=size)
     return decode_jpeg(data)
+
+
+def encode_jpeg(input: torch.Tensor, quality: int = 75) -> torch.Tensor:
+    """
+    Takes an input tensor in CHW layout (or HW in the case of grayscale images)
+    and returns a buffer with the contents of its corresponding JPEG file.
+    Arguments:
+        input (Tensor[channels, image_height, image_width]): int8 image tensor
+    of `c` channels, where `c` must be 1 or 3.
+        quality (int): Quality of the resulting JPEG file, it must be a number
+    between 1 and 100. Default: 75
+    Returns
+        output (Tensor[1]): A one dimensional int8 tensor that contains the raw
+    bytes of the JPEG file.
+    """
+    if quality < 1 or quality > 100:
+        raise ValueError('Image quality should be a positive number '
+                         'between 1 and 100')
+
+    output = torch.ops.image.encode_jpeg(input, quality)
+    return output
+
+
+def write_jpeg(input: torch.Tensor, filename: str, quality: int = 75):
+    """
+    Takes an input tensor in CHW layout (or HW in the case of grayscale images)
+    and saves it in a JPEG file.
+    Arguments:
+        input (Tensor[channels, image_height, image_width]): int8 image tensor
+    of `c` channels, where `c` must be 1 or 3.
+        filename (str): Path to save the image.
+        quality (int): Quality of the resulting JPEG file, it must be a number
+    between 1 and 100. Default: 75
+    """
+    if quality < 1 or quality > 100:
+        raise ValueError('Image quality should be a positive number '
+                         'between 1 and 100')
+
+    torch.ops.image.write_jpeg(input, filename, quality)


Test patch:
diff --git a/test/assets/jpeg_write/grace_hopper_517x606_pil.jpg b/test/assets/jpeg_write/grace_hopper_517x606_pil.jpg
new file mode 100644
Binary files /dev/null and b/test/assets/jpeg_write/grace_hopper_517x606_pil.jpg differ
diff --git a/test/test_image.py b/test/test_image.py
--- a/test/test_image.py
+++ b/test/test_image.py
@@ -1,4 +1,5 @@
 import os
+import io
 import glob
 import unittest
 import sys
@@ -6,7 +7,8 @@
 import torch
 import torchvision
 from PIL import Image
-from torchvision.io.image import read_png, decode_png, read_jpeg, decode_jpeg
+from torchvision.io.image import (
+    read_png, decode_png, read_jpeg, decode_jpeg, encode_jpeg, write_jpeg)
 import numpy as np
 
 IMAGE_ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets")
@@ -17,7 +19,7 @@
 def get_images(directory, img_ext):
     assert os.path.isdir(directory)
     for root, _, files in os.walk(directory):
-        if os.path.basename(root) == 'damaged_jpeg':
+        if os.path.basename(root) in {'damaged_jpeg', 'jpeg_write'}:
             continue
 
         for fl in files:
@@ -66,6 +68,71 @@ def test_damaged_images(self):
             with self.assertRaises(RuntimeError):
                 read_jpeg(image_path)
 
+    def test_encode_jpeg(self):
+        for img_path in get_images(IMAGE_ROOT, ".jpg"):
+            dirname = os.path.dirname(img_path)
+            filename, _ = os.path.splitext(os.path.basename(img_path))
+            write_folder = os.path.join(dirname, 'jpeg_write')
+            expected_file = os.path.join(
+                write_folder, '{0}_pil.jpg'.format(filename))
+            img = read_jpeg(img_path)
+
+            with open(expected_file, 'rb') as f:
+                pil_bytes = f.read()
+                pil_bytes = torch.as_tensor(list(pil_bytes), dtype=torch.uint8)
+            for src_img in [img, img.contiguous()]:
+                # PIL sets jpeg quality to 75 by default
+                jpeg_bytes = encode_jpeg(src_img, quality=75)
+                self.assertTrue(jpeg_bytes.equal(pil_bytes))
+
+        with self.assertRaisesRegex(
+                RuntimeError, "Input tensor dtype should be uint8"):
+            encode_jpeg(torch.empty((3, 100, 100), dtype=torch.float32))
+
+        with self.assertRaisesRegex(
+                ValueError, "Image quality should be a positive number "
+                "between 1 and 100"):
+            encode_jpeg(torch.empty((3, 100, 100), dtype=torch.uint8), quality=-1)
+
+        with self.assertRaisesRegex(
+                ValueError, "Image quality should be a positive number "
+                "between 1 and 100"):
+            encode_jpeg(torch.empty((3, 100, 100), dtype=torch.uint8), quality=101)
+
+        with self.assertRaisesRegex(
+                RuntimeError, "The number of channels should be 1 or 3, got: 5"):
+            encode_jpeg(torch.empty((5, 100, 100), dtype=torch.uint8))
+
+        with self.assertRaisesRegex(
+                RuntimeError, "Input data should be a 3-dimensional tensor"):
+            encode_jpeg(torch.empty((1, 3, 100, 100), dtype=torch.uint8))
+
+        with self.assertRaisesRegex(
+                RuntimeError, "Input data should be a 3-dimensional tensor"):
+            encode_jpeg(torch.empty((100, 100), dtype=torch.uint8))
+
+    def test_write_jpeg(self):
+        for img_path in get_images(IMAGE_ROOT, ".jpg"):
+            img = read_jpeg(img_path)
+
+            basedir = os.path.dirname(img_path)
+            filename, _ = os.path.splitext(os.path.basename(img_path))
+            torch_jpeg = os.path.join(
+                basedir, '{0}_torch.jpg'.format(filename))
+            pil_jpeg = os.path.join(
+                basedir, 'jpeg_write', '{0}_pil.jpg'.format(filename))
+
+            write_jpeg(img, torch_jpeg, quality=75)
+
+            with open(torch_jpeg, 'rb') as f:
+                torch_bytes = f.read()
+
+            with open(pil_jpeg, 'rb') as f:
+                pil_bytes = f.read()
+
+            os.remove(torch_jpeg)
+            self.assertEqual(torch_bytes, pil_bytes)
+
     def test_read_png(self):
         # Check across .png
         for img_path in get_images(IMAGE_DIR, ".png"):