87 lines
3.1 KiB
C++
87 lines
3.1 KiB
C++
// Copyright (c) OpenMMLab. All rights reserved.
|
|
|
|
#include "codebase/mmseg/mmseg.h"
|
|
#include "core/tensor.h"
|
|
#include "core/utils/device_utils.h"
|
|
#include "core/utils/formatter.h"
|
|
#include "preprocess/cpu/opencv_utils.h"
|
|
#include "preprocess/transform/transform.h"
|
|
|
|
namespace mmdeploy::mmseg {
|
|
|
|
class ResizeMask : public MMSegmentation {
|
|
public:
|
|
explicit ResizeMask(const Value &cfg) : MMSegmentation(cfg) {
|
|
try {
|
|
classes_ = cfg["params"]["num_classes"].get<int>();
|
|
} catch (const std::exception &e) {
|
|
ERROR("no ['params']['num_classes'] is specified in cfg: {}", cfg);
|
|
throw_exception(eInvalidArgument);
|
|
}
|
|
}
|
|
|
|
Result<Value> operator()(const Value &preprocess_result, const Value &inference_result) {
|
|
DEBUG("preprocess: {}\ninference: {}", preprocess_result, inference_result);
|
|
|
|
auto mask = inference_result["output"].get<Tensor>();
|
|
DEBUG("tensor.name: {}, tensor.shape: {}, tensor.data_type: {}", mask.name(), mask.shape(),
|
|
mask.data_type());
|
|
if (!(mask.shape().size() == 4 && mask.shape(0) == 1 && mask.shape(1) == 1)) {
|
|
ERROR("unsupported `output` tensor, shape: {}", mask.shape());
|
|
return Status(eNotSupported);
|
|
}
|
|
|
|
auto height = (int)mask.shape(2);
|
|
auto width = (int)mask.shape(3);
|
|
auto input_height = preprocess_result["img_metas"]["ori_shape"][1].get<int>();
|
|
auto input_width = preprocess_result["img_metas"]["ori_shape"][2].get<int>();
|
|
Device host{"cpu"};
|
|
OUTCOME_TRY(auto host_tensor, MakeAvailableOnDevice(mask, host, stream_));
|
|
OUTCOME_TRY(stream_.Wait());
|
|
if (mask.data_type() == DataType::kINT64) {
|
|
// change kINT64 to 2 INT32
|
|
TensorDesc desc{.device = host_tensor.device(),
|
|
.data_type = DataType::kINT32,
|
|
.shape = {1, 2, height, width},
|
|
.name = host_tensor.name()};
|
|
Tensor _host_tensor(desc, mask.buffer());
|
|
return MaskResize(_host_tensor, input_height, input_width);
|
|
} else if (mask.data_type() == DataType::kINT32) {
|
|
return MaskResize(host_tensor, input_height, input_width);
|
|
} else {
|
|
ERROR("unsupported `output` tensor, dtype: {}", (int)mask.data_type());
|
|
return Status(eNotSupported);
|
|
}
|
|
}
|
|
|
|
private:
|
|
Result<Value> MaskResize(Tensor &tensor, int dst_height, int dst_width) {
|
|
auto channel = tensor.shape(1);
|
|
auto height = tensor.shape(2);
|
|
auto width = tensor.shape(3);
|
|
|
|
// reshape tensor to convert it to cv::Mat
|
|
tensor.Reshape({1, height, width, channel});
|
|
auto mat = cpu::Tensor2CVMat(tensor);
|
|
auto dst = cpu::Resize(mat, dst_height, dst_width, "nearest");
|
|
if (channel == 1) {
|
|
auto output_tensor = cpu::CVMat2Tensor(dst);
|
|
SegmentorOutput output{output_tensor, dst_height, dst_width, classes_};
|
|
return to_value(output);
|
|
} else {
|
|
cv::Mat _dst;
|
|
cv::extractChannel(dst, _dst, 0);
|
|
auto output_tensor = cpu::CVMat2Tensor(_dst);
|
|
SegmentorOutput output{output_tensor, dst_height, dst_width, classes_};
|
|
return to_value(output);
|
|
}
|
|
}
|
|
|
|
protected:
|
|
int classes_{};
|
|
};
|
|
|
|
REGISTER_CODEBASE_COMPONENT(MMSegmentation, ResizeMask);
|
|
|
|
} // namespace mmdeploy::mmseg
|