image.cpp 1.0 KB

123456789101112131415161718192021222324252627282930313233343536
  1. #include "image.h"
  2. #include <ATen/core/op_registration/op_registration.h>
  3. #ifdef USE_PYTHON
  4. #include <Python.h>
  5. #endif
  6. // If we are in a Windows environment, we need to define
  7. // initialization functions for the _custom_ops extension
  8. #ifdef USE_PYTHON
  9. #ifdef _WIN32
  10. PyMODINIT_FUNC PyInit_image(void) {
  11. // No need to do anything.
  12. return NULL;
  13. }
  14. #endif
  15. #endif // USE_PYTHON
  16. namespace vision {
  17. namespace image {
  18. static auto registry =
  19. torch::RegisterOperators()
  20. .op("image::decode_png", &decode_png)
  21. .op("image::encode_png", &encode_png)
  22. .op("image::decode_jpeg", &decode_jpeg)
  23. .op("image::encode_jpeg", &encode_jpeg)
  24. .op("image::read_file", &read_file)
  25. .op("image::write_file", &write_file)
  26. .op("image::decode_image", &decode_image)
  27. .op("image::decode_jpeg_cuda", &decode_jpeg_cuda)
  28. .op("image::_jpeg_version", &_jpeg_version)
  29. .op("image::_is_compiled_against_turbo", &_is_compiled_against_turbo);
  30. } // namespace image
  31. } // namespace vision