Image decoder for data_pipeline

.github/workflows/_build_wheel-macos.yaml

@@ -37,7 +37,7 @@ jobs:
           HOMEBREW_NO_INSTALL_UPGRADE: 1
           HOMEBREW_NO_INSTALLED_DEPENDENTS_CHECK: 1
         run: |
-          brew install libsndfile python@${{ inputs.py }} || true
+          brew install libjpeg libpng libsndfile python@${{ inputs.py }} || true
       - name: Create the Python virtual environment
         run: |
           /usr/local/bin/python${{ inputs.py }} -m venv ~/venv
@@ -100,7 +100,7 @@ jobs:
           HOMEBREW_NO_INSTALL_UPGRADE: 1
           HOMEBREW_NO_INSTALLED_DEPENDENTS_CHECK: 1
         run: |
-          brew install libsndfile python@${{ inputs.py }} || true
+          brew install libjpeg libpng libsndfile python@${{ inputs.py }} || true
       - name: Download wheels and native tests from staging
         uses: actions/download-artifact@v3
         with:

INSTALL_FROM_SOURCE.md

@@ -53,21 +53,22 @@ reusing an existing one to avoid dependency conflicts.
 ## 3. Install Dependencies
 
 ### 3.1 System Dependencies
-fairseq2 has a dependency on
-[libsndfile](https://github.com/libsndfile/libsndfile) that can be installed via
-the system package manager on most Linux distributions, or via Homebrew on
+fairseq2 depends on [libjpeg](https://libjpeg.sourceforge.net),
+[libpng](http://www.libpng.org/pub/png/libpng.html), and
+[libsndfile](https://github.com/libsndfile/libsndfile), which can be installed
+via the system package manager on most Linux distributions, or via Homebrew on
 macOS.
 
 For Ubuntu-based systems, run:
 
 ```sh
-sudo apt install libsndfile-dev
+sudo apt install libjpeg8-dev libpng-dev libsndfile-dev
 ```
 
 Similarly, on Fedora, run:
 
 ```sh
-sudo dnf install libsndfile-devel
+sudo dnf install libjpeg-devel libpng-devel libsndfile-devel
 ```
 
 For other Linux distributions, please consult its documentation on how to
@@ -76,7 +77,7 @@ install packages.
 For macOS, you can use Homebrew:
 
 ```sh
-brew install libsndfile
+brew install libjpeb libpng libsndfile
 ```
 
 ### 3.2 PyTorch

README.md

@@ -59,19 +59,20 @@ fairseq2 is also used by various external projects such as:
 ## Installing on Linux
 
 ### System Dependencies
-fairseq2 has a dependency on
-[libsndfile](https://github.com/libsndfile/libsndfile) that can be installed via
-the system package manager on most Linux distributions. For Ubuntu-based
+fairseq2 depends on [libjpeg](https://libjpeg.sourceforge.net),
+[libpng](http://www.libpng.org/pub/png/libpng.html), and
+[libsndfile](https://github.com/libsndfile/libsndfile), which can be installed
+via the system package manager on most Linux distributions. For Ubuntu-based
 systems, run:
 
 ```sh
-sudo apt install libsndfile1
+sudo apt install libjpeg8 libpng16-16 libsndfile1
 ```
 
 Similarly, on Fedora, run:
 
 ```sh
-sudo dnf install libsndfile
+sudo dnf install libjpeg libpng libsndfile
 ```
 
 For other Linux distributions, please consult its documentation on how to
@@ -139,12 +140,13 @@ pip install fairseq2\
 ## Installing on macOS
 
 ### System Dependencies
-fairseq2 has a dependency on
-[libsndfile](https://github.com/libsndfile/libsndfile) that can be installed via
-Homebrew:
+fairseq2 depends on [libjpeg](https://libjpeg.sourceforge.net),
+[libpng](http://www.libpng.org/pub/png/libpng.html), and
+[libsndfile](https://github.com/libsndfile/libsndfile), which can be installed
+via Homebrew:
 
 ```sh
-brew install libsndfile
+brew install libjpeg libpng libsndfile
 ```
 
 ### pip

fairseq2n/CMakeLists.txt

@@ -87,6 +87,13 @@ option(FAIRSEQ2N_TREAT_WARNINGS_AS_ERRORS
         OFF
 )
 
+option(FAIRSEQ2N_SUPPORT_IMAGE
+    #DESCRIPTION
+        "Supports JPEG/PNG decoding."
+    #VALUE
+        ON
+)
+
 option(FAIRSEQ2N_USE_LIBTORCH
     #DESCRIPTION
         "Uses libtorch instead of PyTorch."
@@ -160,6 +167,12 @@ find_package(SndFile 1.0.25 REQUIRED)
 
 find_package(Threads REQUIRED)
 
+if(FAIRSEQ2N_SUPPORT_IMAGE)
+    find_package(JPEG REQUIRED)
+
+    find_package(PNG REQUIRED)
+endif()
+
 if(FAIRSEQ2N_THREAD_LIB STREQUAL "tbb")
     find_package(TBB 2021.8 REQUIRED)
 endif()

fairseq2n/cmake/summary.cmake

@@ -34,6 +34,7 @@ function(fairseq2n_print_project_summary)
         message(STATUS "  FAIRSEQ2N_SANITIZERS               : ${FAIRSEQ2N_SANITIZERS}")
     endif()
     message(STATUS "  FAIRSEQ2N_TREAT_WARNINGS_AS_ERRORS : ${FAIRSEQ2N_TREAT_WARNINGS_AS_ERRORS}")
+    message(STATUS "  FAIRSEQ2N_SUPPORT_IMAGE            : ${FAIRSEQ2N_SUPPORT_IMAGE}")
     message(STATUS "  FAIRSEQ2N_USE_LIBTORCH             : ${FAIRSEQ2N_USE_LIBTORCH}")
     message(STATUS "  FAIRSEQ2N_USE_CUDA                 : ${FAIRSEQ2N_USE_CUDA}")
     if(FAIRSEQ2N_USE_CUDA)
@@ -50,6 +51,10 @@ function(fairseq2n_print_project_summary)
     if(FAIRSEQ2N_THREAD_LIB STREQUAL "tbb")
         message(STATUS "    Intel oneTBB                     : ${TBB_VERSION}")
     endif()
+    if(FAIRSEQ2N_SUPPORT_IMAGE)
+        message(STATUS "    libjpeg                          : ${JPEG_VERSION}")
+        message(STATUS "    libpng                           : ${PNG_VERSION_STRING}")
+    endif()
     message(STATUS "    libsndfile                       : ${SndFile_VERSION}")
     message(STATUS "")
 endfunction()

fairseq2n/python/src/fairseq2n/__init__.py

@@ -132,6 +132,13 @@ def supports_cuda() -> bool:
     return _supports_cuda()  # type: ignore[no-any-return]
 
 
+def supports_image() -> bool:
+    """Return ``True`` if fairseq2n supports JPEG/PNG decoding."""
+    from fairseq2n.bindings import _supports_image  # type: ignore[attr-defined]
+
+    return _supports_image()  # type: ignore[no-any-return]
+
+
 def cuda_version() -> Optional[Tuple[int, int]]:
     """Return the version of CUDA that fairseq2n supports.
 

fairseq2n/python/src/fairseq2n/bindings/CMakeLists.txt

@@ -17,8 +17,10 @@ target_sources(py_bindings
         init.cc
         memory.cc
         data/audio.cc
+        data/image.cc
         data/data_pipeline.cc
         data/init.cc
+        data/image.cc
         data/string.cc
         data/text/converters.cc
         data/text/init.cc

fairseq2n/python/src/fairseq2n/bindings/data/image.cc

@@ -0,0 +1,45 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include "fairseq2n/bindings/module.h"
+
+#include <cstdint>
+#include <memory>
+
+#include <ATen/Device.h>
+#include <ATen/ScalarType.h>
+
+#include <fairseq2n/float.h>
+#include <fairseq2n/data/image/image_decoder.h>
+
+namespace py = pybind11;
+
+namespace fairseq2n {
+
+void
+def_image(py::module_ &data_module)
+{
+    py::module_ m = data_module.def_submodule("image");
+
+    // ImageDecoder
+    py::class_<image_decoder, std::shared_ptr<image_decoder>>(m, "ImageDecoder")
+        .def(
+            py::init([](
+                std::optional<at::Device> maybe_device,
+                bool pin_memory)
+            {
+                auto opts = image_decoder_options()
+                    .maybe_device(maybe_device).pin_memory(pin_memory);
+
+                return std::make_shared<image_decoder>(opts);
+            }),
+            py::arg("device") = std::nullopt,
+            py::arg("pin_memory") = false)
+        .def("__call__", &image_decoder::operator(), py::call_guard<py::gil_scoped_release>{});
+
+    map_functors().register_<image_decoder>();
+}
+}  // namespace fairseq2n

fairseq2n/python/src/fairseq2n/bindings/data/init.cc

@@ -40,6 +40,8 @@ def_data(py::module_ &base)
 
     def_audio(m);
 
+    def_image(m);
+
     def_data_pipeline(m);
 
     def_string(m);

fairseq2n/python/src/fairseq2n/bindings/init.cc

@@ -24,6 +24,13 @@ PYBIND11_MODULE(bindings, m)
           return supports_cuda;
         });
 
+    m.def(
+        "_supports_image",
+        []
+        {
+          return supports_image;
+        });
+
 // See https://github.com/llvm/llvm-project/issues/57123.
 #pragma clang diagnostic push
 #pragma clang diagnostic ignored "-Wunreachable-code-return"

fairseq2n/python/src/fairseq2n/bindings/module.h

@@ -22,6 +22,9 @@ namespace fairseq2n {
 void
 def_audio(pybind11::module_ &data_module);
 
+void
+def_image(pybind11::module_ &data_module);
+
 void
 def_data(pybind11::module_ &base_module);
 

fairseq2n/src/fairseq2n/CMakeLists.txt

@@ -55,6 +55,7 @@ target_sources(fairseq2n
         data/audio/detail/sndfile.cc
         data/detail/file.cc
         data/detail/file_system.cc
+        data/image/image_decoder.cc
         data/text/string_splitter.cc
         data/text/string_to_int_converter.cc
         data/text/string_to_tensor_converter.cc
@@ -69,6 +70,14 @@ target_sources(fairseq2n
         data/text/sentencepiece/sp_processor.cc
 )
 
+if(FAIRSEQ2N_SUPPORT_IMAGE)
+    target_sources(fairseq2n
+        PRIVATE
+            data/image/detail/jpeg_decompress_struct.cc
+            data/image/detail/png_read_struct.cc
+    )
+endif()
+
 if(FAIRSEQ2N_USE_CUDA)
     target_sources(fairseq2n
         PRIVATE
@@ -80,8 +89,8 @@ fairseq2n_set_compile_options(fairseq2n)
 
 target_compile_features(fairseq2n PUBLIC cxx_std_17)
 
-if(FAIRSEQ2N_THREAD_LIB STREQUAL "tbb")
-    target_compile_definitions(fairseq2n PRIVATE FAIRSEQ2N_USE_TBB)
+if(FAIRSEQ2N_SUPPORT_IMAGE)
+    target_compile_definitions(fairseq2n PRIVATE FAIRSEQ2N_SUPPORT_IMAGE)
 endif()
 
 if(FAIRSEQ2N_USE_CUDA)
@@ -90,6 +99,10 @@ if(FAIRSEQ2N_USE_CUDA)
     target_compile_definitions(fairseq2n PRIVATE FAIRSEQ2N_USE_CUDA)
 endif()
 
+if(FAIRSEQ2N_THREAD_LIB STREQUAL "tbb")
+    target_compile_definitions(fairseq2n PRIVATE FAIRSEQ2N_USE_TBB)
+endif()
+
 if(PROJECT_IS_TOP_LEVEL)
     set(system)
 else()
@@ -102,6 +115,9 @@ target_include_directories(fairseq2n ${system}
         $<BUILD_INTERFACE:${PROJECT_BINARY_DIR}/src>
 )
 
+find_package(PNG REQUIRED)
+find_package(JPEG REQUIRED)
+
 target_link_libraries(fairseq2n
     PRIVATE
         ${CMAKE_DL_LIBS}
@@ -114,18 +130,23 @@ target_link_libraries(fairseq2n
         Threads::Threads
         sentencepiece-static
         SndFile::sndfile
+        PNG::PNG
+        JPEG::JPEG
     PUBLIC
         torch
 )
 
-if(FAIRSEQ2N_THREAD_LIB STREQUAL "tbb")
-    target_link_libraries(fairseq2n PRIVATE TBB::tbb)
+if(FAIRSEQ2N_SUPPORT_IMAGE)
+    target_link_libraries(fairseq2n PRIVATE JPEG::JPEG PNG::PNG)
 endif()
 
 if(FAIRSEQ2N_USE_CUDA)
     target_link_libraries(fairseq2n PRIVATE CUDA::cudart)
 endif()
 
+if(FAIRSEQ2N_THREAD_LIB STREQUAL "tbb")
+    target_link_libraries(fairseq2n PRIVATE TBB::tbb)
+endif()
 
 fairseq2n_set_link_options(fairseq2n)
 
@@ -174,10 +195,20 @@ install(
 # Library Configuration
 # ------------------------------------------------------------
 
+if(FAIRSEQ2N_SUPPORT_IMAGE)
+    set(SUPPORTS_IMAGE "true")
+else()
+    set(SUPPORTS_IMAGE "false")
+endif()
+
 if(FAIRSEQ2N_USE_CUDA)
+    set(USES_CUDA "true")
+
     set(CUDA_VERSION_MAJOR "${CUDAToolkit_VERSION_MAJOR}")
     set(CUDA_VERSION_MINOR "${CUDAToolkit_VERSION_MINOR}")
 else()
+    set(USES_CUDA "false")
+
     set(CUDA_VERSION_MAJOR "std::nullopt")
     set(CUDA_VERSION_MINOR "std::nullopt")
 endif()

fairseq2n/src/fairseq2n/config.h.in

@@ -15,9 +15,11 @@ constexpr std::int32_t version_major = @PROJECT_VERSION_MAJOR@;
 constexpr std::int32_t version_minor = @PROJECT_VERSION_MINOR@;
 constexpr std::int32_t version_patch = @PROJECT_VERSION_PATCH@;
 
+constexpr bool supports_image = @SUPPORTS_IMAGE@;
+
+constexpr bool supports_cuda = @USES_CUDA@;
+
 constexpr std::optional<std::int32_t> cuda_version_major = @CUDA_VERSION_MAJOR@;
 constexpr std::optional<std::int32_t> cuda_version_minor = @CUDA_VERSION_MINOR@;
 
-constexpr bool supports_cuda = cuda_version_major.has_value();
-
 }  // namespace fairseq2n

fairseq2n/src/fairseq2n/data/image/detail/jpeg_decompress_struct.cc

@@ -0,0 +1,25 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.error_ptr
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include "fairseq2n/data/image/detail/jpeg_decompress_struct.h"
+
+namespace fairseq2n::detail {
+
+jpeg_decompress::jpeg_decompress() : cinfo() {
+    jpeg_create_decompress(&cinfo);
+}
+
+jpeg_decompress::~jpeg_decompress() {
+    if(cinfo.err != nullptr) {
+        jpeg_destroy_decompress(&cinfo);
+    }
+}
+
+jpeg_decompress_struct& jpeg_decompress::get() {
+    return cinfo;
+}
+
+} // namespace fairseq2n::detail