1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93
|
#include <iostream> #include <vector> #include <opencv2/opencv.hpp>
using namespace std; using namespace cv;
class NsfwDnn { public: enum ModelType { UNKNOWN, TENSORFLOW, ONNX, };
const std::vector<string> labels{"drawings", "hentai", "neutral", "porn", "sexy"};
private: cv::dnn::Net net; int size = 224; ModelType type = TENSORFLOW; public: explicit NsfwDnn() = default;
~NsfwDnn() = default;
bool loadModel(const std::string &model_path, ModelType type = UNKNOWN) { this->type = type; if (type == TENSORFLOW) { size = 224; net = cv::dnn::readNetFromTensorflow(model_path); } else if (type == ONNX) { size = 299; net = cv::dnn::readNetFromONNX(model_path); } return !empty(); }
bool loadModel(const std::vector<unsigned char> &model_buff, const ModelType type) { this->type = type; if (type == TENSORFLOW) { size = 224; net = cv::dnn::readNetFromTensorflow(model_buff); } else if (type == ONNX) { size = 299; net = cv::dnn::readNetFromONNX(model_buff); } return !empty(); }
bool empty() const { return net.empty(); }
std::vector<std::tuple<string, float> > detect(const std::vector<unsigned char> &bin) { if (bin.empty())return {}; const auto image = imdecode(bin, cv::IMREAD_COLOR); if (image.empty())return {}; return detect(image); }
std::vector<std::tuple<string, float> > detect(const cv::Mat &frame) { std::vector<std::tuple<string, float> > result; if (empty())return result; const int w = frame.cols; const int h = frame.rows; const int _max = std::max(h, w); cv::Mat image = cv::Mat::zeros(cv::Size(_max, _max), CV_8UC3); cv::Rect roi(0, 0, w, h); frame.copyTo(image(roi)); cv::Mat blob; cv::dnn::blobFromImage(image, blob, 1.0 / 255.0, cv::Size(size, size), cv::Scalar(), true, false); if (size == 299) { blob = blob.reshape(1, cv::dnn::MatShape({1, size, size, 3})); } net.setInput(blob);
std::vector<cv::Mat> outputs; net.forward(outputs, net.getUnconnectedOutLayersNames());
const int rows = outputs[0].size[1]; const auto rest = reinterpret_cast<float *>(outputs[0].data); result.reserve(rows); for (int i = 0; i < rows; i++) { result.emplace_back(labels[i], *(rest + i)); } return result; } };
|