image.cpp 852 Bytes
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14

#include "image.h"
#include <ATen/ATen.h>
#include <Python.h>

// If we are in a Windows environment, we need to define
// initialization functions for the _custom_ops extension
#ifdef _WIN32
PyMODINIT_FUNC PyInit_image(void) {
  // No need to do anything.
  return NULL;
}
#endif

15
16
static auto registry = torch::RegisterOperators()
                           .op("image::decode_png", &decodePNG)
17
18
                           .op("image::encode_png", &encodePNG)
                           .op("image::write_png", &writePNG)
19
20
                           .op("image::decode_jpeg", &decodeJPEG)
                           .op("image::encode_jpeg", &encodeJPEG)
Francisco Massa's avatar
Francisco Massa committed
21
                           .op("image::write_jpeg", &writeJPEG)
Francisco Massa's avatar
Francisco Massa committed
22
                           .op("image::read_file", &read_file)
Francisco Massa's avatar
Francisco Massa committed
23
                           .op("image::decode_image", &decode_image);