mirror of https://github.com/alibaba/MNN.git
131 lines
4.2 KiB
C++
131 lines
4.2 KiB
C++
//
|
|
// segment.cpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2019/07/01.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include <stdio.h>
|
|
#include <MNN/ImageProcess.hpp>
|
|
#define MNN_OPEN_TIME_TRACE
|
|
#include <algorithm>
|
|
#include <fstream>
|
|
#include <functional>
|
|
#include <memory>
|
|
#include <sstream>
|
|
#include <vector>
|
|
#include <MNN/expr/Expr.hpp>
|
|
#include <MNN/expr/ExprCreator.hpp>
|
|
#include <MNN/AutoTime.hpp>
|
|
#include <MNN/Interpreter.hpp>
|
|
#define STB_IMAGE_IMPLEMENTATION
|
|
#include "stb_image.h"
|
|
#define STB_IMAGE_WRITE_IMPLEMENTATION
|
|
#include "stb_image_write.h"
|
|
|
|
using namespace MNN;
|
|
using namespace MNN::CV;
|
|
using namespace MNN::Express;
|
|
|
|
int main(int argc, const char* argv[]) {
|
|
if (argc < 4) {
|
|
MNN_PRINT("Usage: ./segment.out model.mnn input.jpg output.jpg\n");
|
|
return 0;
|
|
}
|
|
std::shared_ptr<Interpreter> net;
|
|
net.reset(Interpreter::createFromFile(argv[1]));
|
|
if (net == nullptr) {
|
|
MNN_ERROR("Invalid Model\n");
|
|
return 0;
|
|
}
|
|
ScheduleConfig config;
|
|
auto session = net->createSession(config);
|
|
auto input = net->getSessionInput(session, nullptr);
|
|
auto shape = input->shape();
|
|
if (shape[0] != 1) {
|
|
shape[0] = 1;
|
|
net->resizeTensor(input, shape);
|
|
net->resizeSession(session);
|
|
}
|
|
{
|
|
int size_w = 0;
|
|
int size_h = 0;
|
|
int bpp = 0;
|
|
bpp = shape[1];
|
|
size_h = shape[2];
|
|
size_w = shape[3];
|
|
if (bpp == 0)
|
|
bpp = 1;
|
|
if (size_h == 0)
|
|
size_h = 1;
|
|
if (size_w == 0)
|
|
size_w = 1;
|
|
MNN_PRINT("input: w:%d , h:%d, bpp: %d\n", size_w, size_h, bpp);
|
|
|
|
auto inputPatch = argv[2];
|
|
int width, height, channel;
|
|
auto inputImage = stbi_load(inputPatch, &width, &height, &channel, 4);
|
|
if (nullptr == inputImage) {
|
|
MNN_ERROR("Can't open %s\n", inputPatch);
|
|
return 0;
|
|
}
|
|
MNN_PRINT("origin size: %d, %d\n", width, height);
|
|
Matrix trans;
|
|
// Set scale, from dst scale to src
|
|
trans.setScale((float)(width-1) / (size_w-1), (float)(height-1) / (size_h-1));
|
|
ImageProcess::Config config;
|
|
config.filterType = CV::BILINEAR;
|
|
// float mean[3] = {103.94f, 116.78f, 123.68f};
|
|
// float normals[3] = {0.017f, 0.017f, 0.017f};
|
|
float mean[3] = {127.5f, 127.5f, 127.5f};
|
|
float normals[3] = {0.00785f, 0.00785f, 0.00785f};
|
|
::memcpy(config.mean, mean, sizeof(mean));
|
|
::memcpy(config.normal, normals, sizeof(normals));
|
|
config.sourceFormat = RGBA;
|
|
config.destFormat = RGB;
|
|
|
|
std::shared_ptr<ImageProcess> pretreat(ImageProcess::create(config));
|
|
pretreat->setMatrix(trans);
|
|
pretreat->convert((uint8_t*)inputImage, width, height, 0, input);
|
|
stbi_image_free(inputImage);
|
|
}
|
|
// Run model
|
|
net->runSession(session);
|
|
|
|
// Post treat by MNN-Express
|
|
{
|
|
/* Create VARP by tensor Begin*/
|
|
auto outputTensor = net->getSessionOutput(session, nullptr);
|
|
// First Create a Expr, then create Variable by the 0 index of expr
|
|
auto output = Variable::create(Expr::create(outputTensor));
|
|
if (nullptr == output->getInfo()) {
|
|
MNN_ERROR("Alloc memory or compute size error\n");
|
|
return 0;
|
|
}
|
|
/* Create VARP by tensor End*/
|
|
|
|
// Turn dataFormat to NHWC for easy to run TopKV2
|
|
output = _Convert(output, NHWC);
|
|
auto width = output->getInfo()->dim[2];
|
|
auto height = output->getInfo()->dim[1];
|
|
auto channel = output->getInfo()->dim[3];
|
|
MNN_PRINT("output w = %d, h=%d\n", width, height);
|
|
|
|
const int humanIndex = 15;
|
|
output = _Reshape(output, {-1, channel});
|
|
auto kv = _TopKV2(output, _Scalar<int>(1));
|
|
// Use indice in TopKV2's C axis
|
|
auto index = kv[1];
|
|
// If is human, set 255, else set 0
|
|
auto mask = _Select(_Equal(index, _Scalar<int>(humanIndex)), _Scalar<int>(255), _Scalar<int>(0));
|
|
|
|
//If need faster, use this code
|
|
//auto mask = _Equal(index, _Scalar<int>(humanIndex)) * _Scalar<int>(255);
|
|
|
|
mask = _Cast<uint8_t>(mask);
|
|
stbi_write_png(argv[3], width, height, 1, mask->readMap<uint8_t>(), width);
|
|
}
|
|
return 0;
|
|
}
|