Experimental image decoder based on imgcodec module (#4223)

New operators:
* experimental.decoders.image
* experimental.decoders.image_crop
* experimental.decoders.image_slice
* experimental.decoders.image_random_crop
* experimental.peek_image_shape

These are drop-in replacements for the existing operators.
The difference is that they now handle EXIF and TIFF orientation tags as well as support other output bit depths.

Missing features:
* Hybrid GPU backend for NVJPEG decoder
* Hardware backend for NVJPEG decoder

author Szymon Karpiński <skarpinski@nvidia.com>
committer Michal Zientkiewicz <michalz@nvidia.com>

Signed-off-by: Szymon Karpiński <skarpinski@nvidia.com>
Co-authored-by: Michal Zientkiewicz <michalz@nvidia.com>

Author: szkarpinski

dali/imgcodec/decoders/memory_pool.cc

@@ -21,6 +21,7 @@
 #include <unordered_map>
 #include "dali/imgcodec/decoders/memory_pool.h"
 #include "dali/core/cuda_error.h"
+#include "dali/core/cuda_stream_pool.h"
 #include "dali/core/mm/malloc_resource.h"
 #include "dali/pipeline/data/buffer.h"
 

dali/imgcodec/decoders/nvjpeg2k/nvjpeg2k.cc

@@ -22,6 +22,7 @@
 #include "dali/imgcodec/util/convert_gpu.h"
 #include "dali/core/static_switch.h"
 #include "dali/imgcodec/registry.h"
+#include "dali/pipeline/util/for_each_thread.h"
 
 namespace dali {
 namespace imgcodec {
@@ -63,8 +64,10 @@ NvJpeg2000DecoderInstance::NvJpeg2000DecoderInstance(
   nvjpeg2k_handle_ = NvJpeg2kHandle(&nvjpeg2k_dev_alloc_, &nvjpeg2k_pin_alloc_);
   DALI_ENFORCE(nvjpeg2k_handle_, "NvJpeg2kHandle initalization failed");
 
-  for (auto &res : per_thread_resources_)
-    res = {nvjpeg2k_handle_, device_memory_padding, device_id_};
+  ForEachThread(*tp_, [&](int tid) noexcept {
+    CUDA_CALL(cudaSetDevice(device_id));
+    per_thread_resources_[tid] = {nvjpeg2k_handle_, device_memory_padding, device_id_};
+  });
 
   for (const auto &thread_id : tp_->GetThreadIds()) {
     if (device_memory_padding > 0) {
@@ -81,9 +84,18 @@ NvJpeg2000DecoderInstance::NvJpeg2000DecoderInstance(
 }
 
 NvJpeg2000DecoderInstance::~NvJpeg2000DecoderInstance() {
+  tp_->WaitForWork();
   for (const auto &res : per_thread_resources_)
     CUDA_CALL(cudaStreamSynchronize(res.cuda_stream));
-  for (const auto &thread_id : tp_->GetThreadIds())
+
+  ForEachThread(*tp_, [&](int tid) {
+      auto &res = per_thread_resources_[tid];
+      res.tile_dec_res.clear();
+      res.nvjpeg2k_decode_state.reset();
+      res.intermediate_buffer.free();
+    });
+
+  for (auto thread_id : tp_->GetThreadIds())
     nvjpeg_memory::DeleteAllBuffers(thread_id);
 }
 
@@ -157,8 +169,9 @@ bool NvJpeg2000DecoderInstance::DecodeJpeg2000(ImageSource *in, void *out, const
                               ctx.nvjpeg2k_stream, &output_image, ctx.cuda_stream);
     return check_status(ret, in);
   } else {
-    // Decode tile by tile: nvjpeg2kDecodeImage seems to be bugged
     auto &image_info = ctx.image_info;
+
+    // Decode tile by tile: nvjpeg2kDecodeImage doesn't work properly with ROI
     auto &roi = ctx.roi;
     std::array tile_shape = {image_info.tile_height, image_info.tile_width};
 
@@ -185,11 +198,10 @@ bool NvJpeg2000DecoderInstance::DecodeJpeg2000(ImageSource *in, void *out, const
 
         if (begin_x < end_x && begin_y < end_y) {
           const TileDecodingResources &per_tile_ctx = ctx.tile_dec_res[state_idx];
-          state_idx = (state_idx + 1) % ctx.tile_dec_res.size();
 
           CUDA_CALL(cudaEventSynchronize(per_tile_ctx.decode_event));
 
-          NvJpeg2kDecodeParams params;
+          auto &params = per_tile_ctx.params;
           CUDA_CALL(nvjpeg2kDecodeParamsSetDecodeArea(params, begin_x, end_x, begin_y, end_y));
 
           auto output_image = PrepareOutputArea(out, pixel_data, pitch_in_bytes, output_offset_x,
@@ -208,6 +220,7 @@ bool NvJpeg2000DecoderInstance::DecodeJpeg2000(ImageSource *in, void *out, const
             return check_status(ret, in);
 
           CUDA_CALL(cudaEventRecord(per_tile_ctx.decode_event, ctx.cuda_stream));
+          state_idx = (state_idx + 1) % ctx.tile_dec_res.size();
         }
       }
     }
@@ -225,11 +238,11 @@ DecodeResult NvJpeg2000DecoderInstance::DecodeImplTask(int thread_idx,
   Context ctx(opts, roi, res);
   DecodeResult result = {false, nullptr};
 
+  CUDA_CALL(cudaEventSynchronize(ctx.decode_event));
+
   if (!ParseJpeg2000Info(in, ctx))
     return result;
 
-  CUDA_CALL(cudaEventSynchronize(ctx.decode_event));
-
   const int64_t channels = ctx.shape[0];
   DALIImageType format = channels == 1 ? DALI_GRAY : DALI_RGB;
   bool is_processing_needed =
@@ -241,6 +254,10 @@ DecodeResult NvJpeg2000DecoderInstance::DecodeImplTask(int thread_idx,
   auto decode_out = out;
   if (is_processing_needed) {
     int64_t type_size = dali::TypeTable::GetTypeInfo(ctx.pixel_type).size();
+    size_t new_size = volume(ctx.shape) * type_size;
+    if (new_size > res.intermediate_buffer.capacity()) {
+      CUDA_CALL(cudaStreamSynchronize(ctx.cuda_stream));
+    }
     res.intermediate_buffer.resize(volume(ctx.shape) * type_size);
     decode_out = {res.intermediate_buffer.data(), ctx.shape, ctx.pixel_type};
   }

dali/imgcodec/decoders/nvjpeg2k/nvjpeg2k.h

@@ -46,7 +46,7 @@ class DLL_PUBLIC NvJpeg2000DecoderInstance : public BatchParallelDecoderImpl {
                               const ROI &roi) override;
 
   FutureDecodeResults ScheduleDecode(DecodeContext ctx,
-                                     span<SampleView<CPUBackend>> out,
+                                     span<SampleView<GPUBackend>> out,
                                      cspan<ImageSource *> in,
                                      DecodeParams opts,
                                      cspan<ROI> rois = {}) override {
@@ -79,15 +79,16 @@ class DLL_PUBLIC NvJpeg2000DecoderInstance : public BatchParallelDecoderImpl {
   struct TileDecodingResources {
     NvJpeg2kDecodeState state;
     CUDAEvent decode_event;
+    NvJpeg2kDecodeParams params;
 
-    explicit TileDecodingResources(const NvJpeg2kHandle &nvjpeg2k_handle, int device_id,
-                                   cudaStream_t cuda_stream)
+    explicit TileDecodingResources(const NvJpeg2kHandle &nvjpeg2k_handle, int device_id)
         : state(nvjpeg2k_handle), decode_event(CUDAEvent::Create(device_id)) {
-      CUDA_CALL(cudaEventRecord(decode_event, cuda_stream));
     }
   };
 
   struct PerThreadResources {
+    static constexpr int kNumParallelTiles = 2;
+
     PerThreadResources() = default;
     PerThreadResources(const NvJpeg2kHandle &nvjpeg2k_handle,
                        size_t device_memory_padding, int device_id)
@@ -97,12 +98,10 @@ class DLL_PUBLIC NvJpeg2000DecoderInstance : public BatchParallelDecoderImpl {
     , decode_event(CUDAEvent::Create(device_id))
     , cuda_stream(CUDAStreamPool::instance().Get(device_id)) {
       intermediate_buffer.resize(device_memory_padding / 8);
-      CUDA_CALL(cudaEventRecord(decode_event, cuda_stream));
 
-      constexpr int kNumParallelTiles = 10;
       tile_dec_res.reserve(kNumParallelTiles);
       for (int i = 0; i < kNumParallelTiles; i++) {
-        tile_dec_res.emplace_back(nvjpeg2k_handle, device_id, cuda_stream);
+        tile_dec_res.emplace_back(nvjpeg2k_handle, device_id);
       }
     }
 

dali/imgcodec/decoders/nvjpeg2k/nvjpeg2k_test.cc

@@ -142,12 +142,9 @@ class NvJpeg2000DecoderTest : public NumpyDecoderTestBase<GPUBackend, OutputType
       AssertEqualSatNorm(img, ref);
   }
 
-  void RunTest(const ImageTestingData &data, std::optional<float> eps = std::nullopt,
-               DALIImageType format = DALI_RGB) {
+  void RunTest(const ImageTestingData &data, std::optional<float> eps = std::nullopt) {
     ImageBuffer image(data.img_path);
-    auto params = this->GetParams();
-    params.format = format;
-    auto decoded = this->Decode(&image.src, params, data.roi);
+    auto decoded = this->Decode(&image.src, this->GetParams(), data.roi);
     auto ref = this->ReadReferenceFrom(data.ref_path);
     AssertEqual(decoded, ref, eps);
   }
@@ -187,13 +184,6 @@ TYPED_TEST(NvJpeg2000DecoderTest, DecodeSingleRoi) {
     this->RunTest(from_regular_file(name, roi));
 }
 
-TYPED_TEST(NvJpeg2000DecoderTest, DecodeSingleAnyData) {
-  for (const auto &name : images) {
-    SCOPED_TRACE(name);
-    this->RunTest(from_regular_file(name), {}, DALI_ANY_DATA);
-  }
-}
-
 TYPED_TEST(NvJpeg2000DecoderTest, DecodeBatchSingleThread) {
   std::vector<ImageTestingData> data;
   for (const auto &name : images)

dali/imgcodec/decoders/opencv_fallback.cc

@@ -89,8 +89,11 @@ DecodeResult OpenCVDecoderInstance::DecodeImplTask(int thread_idx,
       cvimg = cv::imdecode(cv::_InputArray(raw, in->Size()), flags);
     }
 
-    // TODO(michalz): correct the orientation of images loaded with IMREAD_UNCHANGED
-    (void)adjust_orientation;
+    Orientation orientation = {};
+    if (adjust_orientation) {
+      auto info = ImageFormatRegistry::instance().GetImageFormat(in)->Parser()->GetInfo(in);
+      orientation = info.orientation;
+    }
 
     res.success = cvimg.ptr(0) != nullptr;
     if (res.success) {
@@ -110,14 +113,16 @@ DecodeResult OpenCVDecoderInstance::DecodeImplTask(int thread_idx,
 
       int in_channels = cvimg.channels();
       auto out_format = opts.format;
-      // OpenCV uses BGR by default. Here we avoid outputting BGR when requesting ANY_DATA
-      if (out_format == DALI_ANY_DATA && in_format == DALI_ANY_DATA && in_channels == 3) {
-        in_format = DALI_BGR;
-        out_format = DALI_RGB;
+      if (out_format == DALI_ANY_DATA && in_format == DALI_ANY_DATA) {
+        if (in_channels == 3) {
+          // OpenCV uses BGR by default. Here we avoid outputting BGR when requesting ANY_DATA
+          in_format = DALI_BGR;
+          out_format = DALI_RGB;
+        }  // TODO(michalz): support RGBA in DALI
       }
       TensorLayout layout = cvimg.dims == 3 ? "DHWC" : "HWC";
 
-      Convert(out, layout, out_format, in, layout, in_format, roi);
+      Convert(out, layout, out_format, in, layout, in_format, roi, orientation);
     }
   } catch (...) {
     res.exception = std::current_exception();

dali/imgcodec/parsers/jpeg.cc

@@ -78,7 +78,7 @@ ImageInfo JpegParser::GetInfo(ImageSource *encoded) const {
       stream->Read(exif_block.data(), size - 8);
       cv::ExifReader reader;
       if (!reader.parseExif(exif_block.data(), exif_block.size()))
-        DALI_FAIL(make_string("Couldn't parse EXIF data in: ", encoded->SourceInfo()));
+        continue;
       auto entry = reader.getTag(cv::ORIENTATION);
       if (entry.tag != cv::INVALID_TAG) {
         info.orientation = FromExifOrientation(static_cast<ExifOrientation>(entry.field_u16));

dali/operators/CMakeLists.txt

@@ -21,6 +21,7 @@ add_subdirectory(debug)
 add_subdirectory(decoder)
 add_subdirectory(generic)
 add_subdirectory(image)
+add_subdirectory(imgcodec)
 add_subdirectory(math)
 add_subdirectory(random)
 add_subdirectory(reader)
@@ -84,7 +85,7 @@ adjust_source_file_language_property("${DALI_OPERATOR_SRCS}")
 add_library(dali_operators ${LIBTYPE} ${DALI_OPERATOR_SRCS} ${DALI_OPERATOR_OBJ})
 set_target_properties(dali_operators PROPERTIES
     LIBRARY_OUTPUT_DIRECTORY "${DALI_LIBRARY_OUTPUT_DIR}")
-target_link_libraries(dali_operators PUBLIC dali dali_kernels dali_core)
+target_link_libraries(dali_operators PUBLIC dali dali_kernels dali_core dali_imgcodec)
 target_link_libraries(dali_operators PRIVATE dynlink_cuda ${DALI_LIBS})
 if (BUILD_NVML)
   target_link_libraries(dali_operators PRIVATE dynlink_nvml)

dali/operators/imgcodec/CMakeLists.txt

@@ -0,0 +1,17 @@
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+collect_headers(DALI_INST_HDRS PARENT_SCOPE)
+collect_sources(DALI_OPERATOR_SRCS PARENT_SCOPE)
+collect_test_sources(DALI_OPERATOR_TEST_SRCS PARENT_SCOPE)