mirror of
https://github.com/open-mmlab/mmdeploy.git
synced 2025-01-14 08:09:43 +08:00
c99379cc76
* fix ncnn float bugs for segmentor sdk * fix segment for all cases * fix ut Co-authored-by: RunningLeon <mnsheng@yeah.net>
94 lines
2.9 KiB
C++
94 lines
2.9 KiB
C++
// Copyright (c) OpenMMLab. All rights reserved.
|
|
|
|
#include "codebase/mmseg/mmseg.h"
|
|
#include "core/logger.h"
|
|
#include "core/tensor.h"
|
|
#include "core/utils/device_utils.h"
|
|
#include "core/utils/formatter.h"
|
|
#include "opencv_utils.h"
|
|
#include "preprocess/transform/transform.h"
|
|
|
|
namespace mmdeploy::mmseg {
|
|
|
|
// TODO: resize masks on device
|
|
// TODO: when network output is on device, cast it to a smaller type (e.g. int16_t or int8_t
|
|
// according to num classes) to reduce DtoH footprint
|
|
class ResizeMask : public MMSegmentation {
|
|
public:
|
|
explicit ResizeMask(const Value &cfg) : MMSegmentation(cfg) {
|
|
try {
|
|
classes_ = cfg["params"]["num_classes"].get<int>();
|
|
little_endian_ = IsLittleEndian();
|
|
} catch (const std::exception &e) {
|
|
MMDEPLOY_ERROR("no ['params']['num_classes'] is specified in cfg: {}", cfg);
|
|
throw_exception(eInvalidArgument);
|
|
}
|
|
}
|
|
|
|
Result<Value> operator()(const Value &preprocess_result, const Value &inference_result) {
|
|
MMDEPLOY_DEBUG("preprocess: {}\ninference: {}", preprocess_result, inference_result);
|
|
|
|
auto mask = inference_result["output"].get<Tensor>();
|
|
MMDEPLOY_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)) {
|
|
MMDEPLOY_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());
|
|
|
|
OUTCOME_TRY(auto cv_type, GetCvType(mask.data_type()));
|
|
cv::Mat mask_mat(height, width, cv_type, host_tensor.data());
|
|
|
|
if (mask_mat.channels() > 1) {
|
|
cv::extractChannel(mask_mat, mask_mat, little_endian_ ? 0 : mask_mat.channels() - 1);
|
|
}
|
|
if (mask_mat.type() != CV_32S) {
|
|
mask_mat.convertTo(mask_mat, CV_32S);
|
|
}
|
|
|
|
cv::Mat resized_mask = cpu::Resize(mask_mat, input_height, input_width, "nearest");
|
|
|
|
SegmentorOutput output{cpu::CVMat2Tensor(resized_mask), input_height, input_width, classes_};
|
|
return to_value(output);
|
|
}
|
|
|
|
private:
|
|
static Result<int> GetCvType(DataType type) {
|
|
switch (type) {
|
|
case DataType::kFLOAT:
|
|
return CV_32F;
|
|
case DataType::kINT64:
|
|
return CV_32SC2;
|
|
case DataType::kINT32:
|
|
return CV_32S;
|
|
default:
|
|
return Status(eNotSupported);
|
|
}
|
|
}
|
|
|
|
static bool IsLittleEndian() {
|
|
union Un {
|
|
char a;
|
|
int b;
|
|
} un;
|
|
un.b = 1;
|
|
return (int)un.a == 1;
|
|
}
|
|
|
|
protected:
|
|
int classes_{};
|
|
bool little_endian_;
|
|
};
|
|
|
|
REGISTER_CODEBASE_COMPONENT(MMSegmentation, ResizeMask);
|
|
|
|
} // namespace mmdeploy::mmseg
|