mirror of
https://github.com/open-mmlab/mmdeploy.git
synced 2025-01-14 08:09:43 +08:00
51 lines
1.1 KiB
C
51 lines
1.1 KiB
C
|
// Copyright (c) OpenMMLab. All rights reserved.
|
||
|
|
||
|
#ifndef MMDEPLOY_IMAGE2TENSOR_H
|
||
|
#define MMDEPLOY_IMAGE2TENSOR_H
|
||
|
|
||
|
#include "core/tensor.h"
|
||
|
#include "transform.h"
|
||
|
|
||
|
namespace mmdeploy {
|
||
|
/**
|
||
|
* Convert image to `Tensor` by given keys.
|
||
|
*
|
||
|
* The dimension order of input image is (1, H, W, C). The pipeline will convert
|
||
|
* it to (1, C, H, W).
|
||
|
*
|
||
|
*/
|
||
|
class ImageToTensorImpl : public TransformImpl {
|
||
|
public:
|
||
|
ImageToTensorImpl(const Value& args);
|
||
|
~ImageToTensorImpl() = default;
|
||
|
|
||
|
Result<Value> Process(const Value& input) override;
|
||
|
|
||
|
protected:
|
||
|
virtual Result<Tensor> HWC2CHW(const Tensor& tensor) = 0;
|
||
|
|
||
|
protected:
|
||
|
struct to_img_tensor_arg_t {
|
||
|
std::vector<std::string> keys;
|
||
|
};
|
||
|
using ArgType = struct to_img_tensor_arg_t;
|
||
|
|
||
|
protected:
|
||
|
ArgType arg_;
|
||
|
};
|
||
|
|
||
|
class ImageToTensor : public Transform {
|
||
|
public:
|
||
|
explicit ImageToTensor(const Value& args, int version = 0);
|
||
|
~ImageToTensor() = default;
|
||
|
|
||
|
Result<Value> Process(const Value& input) override { return impl_->Process(input); }
|
||
|
|
||
|
private:
|
||
|
std::unique_ptr<ImageToTensorImpl> impl_;
|
||
|
};
|
||
|
|
||
|
} // namespace mmdeploy
|
||
|
|
||
|
#endif // MMDEPLOY_IMAGE2TENSOR_H
|