image.cpp 1.03 KB
Newer Older
1
#include "image.h"
2

3
#include <ATen/core/op_registration/op_registration.h>
4
#ifdef USE_PYTHON
5
#include <Python.h>
6
#endif
7
8
9

// If we are in a Windows environment, we need to define
// initialization functions for the _custom_ops extension
10
#ifdef USE_PYTHON
11
12
13
14
15
16
#ifdef _WIN32
PyMODINIT_FUNC PyInit_image(void) {
  // No need to do anything.
  return NULL;
}
#endif
17
#endif // USE_PYTHON
18

19
20
21
namespace vision {
namespace image {

22
static auto registry = torch::RegisterOperators()
23
24
25
26
                           .op("image::decode_png", &decode_png)
                           .op("image::encode_png", &encode_png)
                           .op("image::decode_jpeg", &decode_jpeg)
                           .op("image::encode_jpeg", &encode_jpeg)
Francisco Massa's avatar
Francisco Massa committed
27
                           .op("image::read_file", &read_file)
Francisco Massa's avatar
Francisco Massa committed
28
                           .op("image::write_file", &write_file)
29
30
                           .op("image::decode_image", &decode_image)
                           .op("image::decode_jpeg_cuda", &decode_jpeg_cuda);
31
32
33

} // namespace image
} // namespace vision