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;
  }
};