Implement cv::cuda::extractChannel

modules/cudaarithm/include/opencv2/cudaarithm.hpp

@@ -540,6 +540,16 @@ CV_EXPORTS void split(InputArray src, GpuMat* dst, Stream& stream = Stream::Null
 /** @overload */
 CV_EXPORTS_W void split(InputArray src, CV_OUT std::vector<GpuMat>& dst, Stream& stream = Stream::Null());
 
+/** @brief Extracts a plane of a multi-channel matrix into an single channel matrix.
+
+@param src Source matrix.
+@param dst Destination single-channel matrix.
+@param stream Stream for the asynchronous version.
+
+@sa split
+ */
+CV_EXPORTS void extractChannel(const GpuMat& src, GpuMat& dst, int channel_index, Stream& stream = Stream::Null());
+
 /** @brief Transposes a matrix.
 
 @param src1 Source matrix. 1-, 4-, 8-byte element sizes are supported for now.

modules/cudaarithm/src/cuda/extractchannel.cu

@@ -0,0 +1,95 @@
+#include "opencv2/cudev/util/vec_traits.hpp"
+#include "opencv2/core/cuda_types.hpp"
+
+namespace
+{
+template<typename T, int cn>
+constexpr T __CV_CUDA_HOST_DEVICE__ get_channel(const int index, const typename cv::cudev::MakeVec<T, cn>::type& value) noexcept
+{
+  return reinterpret_cast<const T*>(&value)[index];
+}
+
+template<typename T, int cn>
+__global__ void extract_channel_kernel(cv::cuda::PtrStepSz<typename cv::cudev::MakeVec<T, cn>::type> many_channel,
+                                       cv::cuda::PtrStepSz<T> single_channel, const int channel_index)
+{
+  const int x = blockIdx.x * blockDim.x + threadIdx.x;
+  const int y = blockIdx.y * blockDim.y + threadIdx.y;
+  if (x >= many_channel.cols || y >= many_channel.rows) {
+    return;
+  }
+
+  single_channel(y, x) = ::get_channel<T, cn>(channel_index, many_channel(y, x));
+}
+
+template<typename T, int cn>
+void extract_channel_impl(cv::cuda::PtrStepSz<typename cv::cudev::MakeVec<T, cn>::type> many_channel,
+                          cv::cuda::PtrStepSz<T> single_channel, const int channel_index, cv::cuda::Stream& stream)
+{
+
+  static constexpr dim3 block(32, 8);
+  const dim3 grid(cv::cudev::divUp(many_channel.cols, block.x), cv::cudev::divUp(many_channel.rows, block.y));
+  ::extract_channel_kernel<T, cn><<<grid, block, 0, cv::cuda::StreamAccessor::getStream(stream)>>>(many_channel, single_channel, channel_index);
+}
+
+template<typename T, int depth>
+void extract_channel_t(const cv::cuda::GpuMat input, cv::cuda::GpuMat& output, const int channel_index, cv::cuda::Stream& stream)
+{
+  static constexpr auto output_type = CV_MAKE_TYPE(depth, 1);
+  if (output.size() != input.size() || output.type() != output_type) {
+    output = cv::cuda::GpuMat(input.size(), output_type);
+  }
+
+  switch (input.channels()) {
+  case 1:
+    ::extract_channel_impl<T, 1>(input, output, channel_index, stream);
+    break;
+  case 2:
+    ::extract_channel_impl<T, 2>(input, output, channel_index, stream);
+    break;
+  case 3:
+    ::extract_channel_impl<T, 3>(input, output, channel_index, stream);
+    break;
+  case 4:
+    ::extract_channel_impl<T, 4>(input, output, channel_index, stream);
+    break;
+    default:
+    CV_Error(cv::Error::StsUnsupportedFormat, "Unsupported channel count");
+  }
+}
+}  // namespace
+
+namespace cv {
+namespace cuda {
+void extractChannel(const GpuMat input, GpuMat& output, const int channel_index, Stream& stream)
+{
+  switch (CV_MAT_DEPTH(input.type())) {
+  case CV_8U:
+    ::extract_channel_t<std::uint8_t, CV_8U>(input, output, channel_index, stream);
+    break;
+  case CV_8S:
+    ::extract_channel_t<std::int8_t, CV_8S>(input, output, channel_index, stream);
+    break;
+  case CV_16S:
+    ::extract_channel_t<std::int16_t, CV_16S>(input, output, channel_index, stream);
+    break;
+  case CV_16U:
+    ::extract_channel_t<std::uint16_t, CV_16U>(input, output, channel_index, stream);
+    break;
+  case CV_32F:
+    ::extract_channel_t<std::float_t, CV_32F>(input, output, channel_index, stream);
+    break;
+  case CV_32S:
+    ::extract_channel_t<std::int32_t, CV_32S>(input, output, channel_index, stream);
+    break;
+  case CV_64F:
+    ::extract_channel_t<std::double_t, CV_64F>(input, output, channel_index, stream);
+    break;
+  case CV_16F:
+    [[fallthrough]];
+  default:
+    CV_Error(Error::StsUnsupportedFormat, "Unsupported data type");
+  }
+}
+}  // namespace cuda
+}  // namespace cv