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

3
#include <ATen/core/op_registration/op_registration.h>
4
5
6
7

// If we are in a Windows environment, we need to define
// initialization functions for the _custom_ops extension
#ifdef _WIN32
Nicolas Hug's avatar
Nicolas Hug committed
8
void* PyInit_image(void) {
9
  return nullptr;
10
11
12
}
#endif

13
14
15
namespace vision {
namespace image {

16
17
static auto registry =
    torch::RegisterOperators()
Nicolas Hug's avatar
Nicolas Hug committed
18
        .op("image::decode_gif", &decode_gif)
19
20
        .op("image::decode_png(Tensor data, int mode, bool allow_16_bits = False, bool apply_exif_orientation=False) -> Tensor",
            &decode_png)
21
        .op("image::encode_png", &encode_png)
22
23
        .op("image::decode_jpeg(Tensor data, int mode, bool apply_exif_orientation=False) -> Tensor",
            &decode_jpeg)
24
25
26
        .op("image::encode_jpeg", &encode_jpeg)
        .op("image::read_file", &read_file)
        .op("image::write_file", &write_file)
27
28
        .op("image::decode_image(Tensor data, int mode, bool apply_exif_orientation=False) -> Tensor",
            &decode_image)
29
30
31
        .op("image::decode_jpeg_cuda", &decode_jpeg_cuda)
        .op("image::_jpeg_version", &_jpeg_version)
        .op("image::_is_compiled_against_turbo", &_is_compiled_against_turbo);
32
33
34

} // namespace image
} // namespace vision