onnxruntime/winml/lib/Api.Image/inc/VideoFrameToTensorConverter.h

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#pragma once

#include "ImageConverter.h"
#include "ImageConversionHelpers.h"
#include "ImageConversionTypes.h"

namespace _winml {

class VideoFrameToTensorConverter : public ImageConverter {
 public:
  VideoFrameToTensorConverter() : shared_handle_(nullptr) {}

  // Function takes in a VideoFrame backed by either a SoftwareBitmap or D3DSurface,
  // and converts to a tensor DX12 Resource.
  // CommandQueue and commandlist should be a compute resource,
  // commandlist will be passed in open, closed and executing when function exits
  // User should pass in a BitmapBounds describing the region of interest, in the form of
  // {upperleft X, upperleft Y, width, height} to be turned into a tensor.
  // If the region of interest is the entire VideoFrame, the input BitmapBounds should describe the entire image.
  void VideoFrameToDX12Tensor(
      _In_ const UINT32 batch_index,
      _In_ winml::LearningModelSession& session,
      _In_ const wm::IVideoFrame& input_video_frame,
      _In_ const wgi::BitmapBounds& input_bounds,
      _In_ const ImageTensorDescription& tensor_description,
      _Inout_ ID3D12Resource* output_tensor);

  // Function takes in a VideoFrame backed by either a SoftwareBitmap or D3DSurface,
  // and converts to a tensor returned in a buffer.
  // User should pass in a BitmapBounds describing the region of interest, in the form of
  // {upperleft X, upperleft Y, width, height} to be turned into a tensor.
  // If the region of interest is the entire VideoFrame, the input BitmapBounds should describe the entire image.
  void VideoFrameToSoftwareTensor(
      _In_ const wm::IVideoFrame& input_video_frame,
      _In_ const wgi::BitmapBounds& input_bounds,
      _In_ const ImageTensorDescription& tensor_description,
      _Out_ BYTE* output_CPU_tensor);

  void ConvertBuffersToBatchedGPUTensor(
      _In_ const std::vector<wss::IBuffer>& buffers,
      _In_ size_t buffer_size_in_bytes,
      _In_ _winml::D3DDeviceCache& device_cache,
      _Inout_ ID3D12Resource* output_resource);

 private:
  GUID d3d11_texture_GUID_ = {0x485e4bb3, 0x3fe8, 0x497b, {0x85, 0x9e, 0xc7, 0x5, 0x18, 0xdb, 0x11, 0x2a}};  // {485E4BB3-3FE8-497B-859E-C70518DB112A}
  GUID handle_GUID_ = {0xce43264e, 0x41f7, 0x4882, {0x9e, 0x20, 0xfa, 0xa5, 0x1e, 0x37, 0x64, 0xfc}};
  ;  // CE43264E-41F7-4882-9E20-FAA51E3764FC
  Microsoft::WRL::ComPtr<ID3D12Resource> upload_heap_;
  Microsoft::WRL::ComPtr<ID3D12Resource> input_D3D12_resource_;
  HANDLE shared_handle_;

  Microsoft::WRL::ComPtr<ID3D12Resource> ShareD3D11Texture(ID3D11Texture2D* pTexture, ID3D12Device* pDevice);

  void ConvertSoftwareBitmapToGPUTensor(
      _In_ const UINT32 batch_index,
      _In_ const wm::IVideoFrame& videoFrame,
      _In_ _winml::D3DDeviceCache& device_cache,
      _In_ const wgi::BitmapBounds& input_bounds,
      _In_ const ImageTensorDescription& tensor_description,
      _Inout_ ID3D12Resource* pOutputResource);

  void ConvertDX12TextureToGPUTensor(
      _In_ const UINT32 batch_index,
      _In_ ID3D12Resource* pInputResource,
      _In_ _winml::D3DDeviceCache& device_cache,
      _In_ const ImageTensorDescription& tensor_description,
      _Inout_ ID3D12Resource* output_resource);

  static D3D12_UNORDERED_ACCESS_VIEW_DESC CreateUAVDescription(
      const UINT32 batch_index,
      const D3D12_RESOURCE_DESC& resource_description,
      const ImageTensorDescription& description);

  static void VideoFrameToTensorConverter::ConvertSoftwareBitmapToCPUTensor(
      _In_ const wgi::SoftwareBitmap& software_bitmap,
      _In_ const ImageTensorDescription& tensor_description,
      _In_ const wgi::BitmapBounds& input_bounds,
      _Inout_ void* CPU_tensor);
};
}  // namespace _winml