46 lines
1.0 KiB
C++
46 lines
1.0 KiB
C++
// Copyright (c) OpenMMLab. All rights reserved.
|
|
|
|
#ifndef MMDEPLOY_LOAD_H
|
|
#define MMDEPLOY_LOAD_H
|
|
|
|
#include "core/mat.h"
|
|
#include "core/tensor.h"
|
|
#include "transform.h"
|
|
|
|
namespace mmdeploy {
|
|
class PrepareImageImpl : public TransformImpl {
|
|
public:
|
|
explicit PrepareImageImpl(const Value& args);
|
|
~PrepareImageImpl() = default;
|
|
|
|
Result<Value> Process(const Value& input) override;
|
|
|
|
protected:
|
|
virtual Result<Tensor> ConvertToBGR(const Mat& img) = 0;
|
|
virtual Result<Tensor> ConvertToGray(const Mat& img) = 0;
|
|
|
|
protected:
|
|
struct prepare_image_arg_t {
|
|
bool to_float32{false};
|
|
std::string color_type{"color"};
|
|
};
|
|
using ArgType = struct prepare_image_arg_t;
|
|
|
|
ArgType arg_;
|
|
};
|
|
|
|
class PrepareImage : public Transform {
|
|
public:
|
|
explicit PrepareImage(const Value& args, int version = 0);
|
|
~PrepareImage() = default;
|
|
|
|
Result<Value> Process(const Value& input) override { return impl_->Process(input); }
|
|
|
|
private:
|
|
std::unique_ptr<PrepareImageImpl> impl_;
|
|
};
|
|
|
|
} // namespace mmdeploy
|
|
|
|
#endif // MMDEPLOY_LOAD_H
|