Remove write_jpeg and write_png in favor of encode + write_file (#2766)

torchvision/csrc/cpu/image/image.cpp

@@ -15,10 +15,8 @@ PyMODINIT_FUNC PyInit_image(void) {
 static auto registry = torch::RegisterOperators()
                            .op("image::decode_png", &decodePNG)
                            .op("image::encode_png", &encodePNG)
-                           .op("image::write_png", &writePNG)
                            .op("image::decode_jpeg", &decodeJPEG)
                            .op("image::encode_jpeg", &encodeJPEG)
-                           .op("image::write_jpeg", &writeJPEG)
                            .op("image::read_file", &read_file)
                            .op("image::write_file", &write_file)
                            .op("image::decode_image", &decode_image);

torchvision/csrc/cpu/image/writejpeg_cpu.cpp

@@ -10,14 +10,6 @@ torch::Tensor encodeJPEG(const torch::Tensor& data, int64_t quality) {
       false, "encodeJPEG: torchvision not compiled with libjpeg support");
 }
 
-void writeJPEG(
-    const torch::Tensor& data,
-    std::string filename,
-    int64_t quality) {
-  TORCH_CHECK(
-      false, "writeJPEG: torchvision not compiled with libjpeg support");
-}
-
 #else
 
 #include <jpeglib.h>
@@ -110,20 +102,4 @@ torch::Tensor encodeJPEG(const torch::Tensor& data, int64_t quality) {
 
   return outTensor;
 }
-
-void writeJPEG(
-    const torch::Tensor& data,
-    std::string filename,
-    int64_t quality) {
-  auto jpegBuf = encodeJPEG(data, quality);
-  auto fileBytes = jpegBuf.data_ptr<uint8_t>();
-  auto fileCStr = filename.c_str();
-  FILE* outfile = fopen(fileCStr, "wb");
-
-  TORCH_CHECK(outfile != NULL, "Error opening output jpeg file");
-
-  fwrite(fileBytes, sizeof(uint8_t), jpegBuf.numel(), outfile);
-  fclose(outfile);
-}
-
 #endif

torchvision/csrc/cpu/image/writejpeg_cpu.h

@@ -3,7 +3,3 @@
 #include <torch/torch.h>
 
 C10_EXPORT torch::Tensor encodeJPEG(const torch::Tensor& data, int64_t quality);
-C10_EXPORT void writeJPEG(
-    const torch::Tensor& data,
-    std::string filename,
-    int64_t quality);

torchvision/csrc/cpu/image/writepng_cpu.cpp

@@ -9,12 +9,6 @@ torch::Tensor encodePNG(const torch::Tensor& data, int64_t compression_level) {
   TORCH_CHECK(false, "encodePNG: torchvision not compiled with libpng support");
 }
 
-void writePNG(
-    const torch::Tensor& data,
-    std::string filename,
-    int64_t compression_level) {
-  TORCH_CHECK(false, "writePNG: torchvision not compiled with libpng support");
-}
 #else
 
 #include <png.h>
@@ -179,19 +173,4 @@ torch::Tensor encodePNG(const torch::Tensor& data, int64_t compression_level) {
   return outTensor;
 }
 
-void writePNG(
-    const torch::Tensor& data,
-    std::string filename,
-    int64_t compression_level) {
-  auto pngBuf = encodePNG(data, compression_level);
-  auto fileBytes = pngBuf.data_ptr<uint8_t>();
-  auto fileCStr = filename.c_str();
-  FILE* outfile = fopen(fileCStr, "wb");
-
-  TORCH_CHECK(outfile != NULL, "Error opening output png file");
-
-  fwrite(fileBytes, sizeof(uint8_t), pngBuf.numel(), outfile);
-  fclose(outfile);
-}
-
 #endif

torchvision/csrc/cpu/image/writepng_cpu.h

@@ -5,7 +5,3 @@
 C10_EXPORT torch::Tensor encodePNG(
     const torch::Tensor& data,
     int64_t compression_level);
-C10_EXPORT void writePNG(
-    const torch::Tensor& data,
-    std::string filename,
-    int64_t compression_level);

torchvision/io/image.py

@@ -104,7 +104,8 @@ def write_png(input: torch.Tensor, filename: str, compression_level: int = 6):
         Compression factor for the resulting file, it must be a number
         between 0 and 9. Default: 6
     """
-    torch.ops.image.write_png(input, filename, compression_level)
+    output = encode_png(input, compression_level)
+    write_file(filename, output)
 
 
 def decode_jpeg(input: torch.Tensor) -> torch.Tensor:
@@ -162,11 +163,8 @@ def write_jpeg(input: torch.Tensor, filename: str, quality: int = 75):
         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)
+    output = encode_jpeg(input, quality)
+    write_file(filename, output)
 
 
 def decode_image(input: torch.Tensor) -> torch.Tensor: