kekxv 技术日志

kekxv

0%

NSFW图像检测

发表于 更新于 分类于

随着网络的发展，越来越多的场景被发掘，例如用户交互。用户交互除了使用文字之外，还能使用图片，这就涉及到图像的检测，毕竟不能把所有的图都接受并显示出来，
那和twitter有什么区别。

这里简单的介绍一下NSFW的图像检测，支持返回结果：

  • drawings - safe for work drawings (including anime)
  • hentai - hentai and pornographic drawings
  • neutral - safe for work neutral images
  • porn - pornographic images, sexual acts
  • sexy - sexually explicit images, not pornography

原始的模型在github上：https://github.com/GantMan/nsfw_model ；模型类型分别为：Mobilenet、Inception，有兴趣的可以查询相关信息。

Python 模式

如果是用python部署，则参考github上面的代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
from nsfw_detector import predict

model = predict.load_model('./nsfw_mobilenet2.224x224.h5')

# Predict single image
predict.classify(model, '2.jpg')
# {'2.jpg': {'sexy': 4.3454722e-05, 'neutral': 0.00026579265, 'porn': 0.0007733492, 'hentai': 0.14751932, 'drawings': 0.85139805}}

# Predict multiple images at once
predict.classify(model, ['2.jpg', '6.jpg'])
# {'2.jpg': {'sexy': 4.3454795e-05, 'neutral': 0.00026579312, 'porn': 0.0007733498, 'hentai': 0.14751942, 'drawings': 0.8513979}, '6.jpg': {'drawings': 0.004214506, 'hentai': 0.013342537, 'neutral': 0.01834045, 'porn': 0.4431829, 'sexy': 0.5209196}}

# Predict for all images in a directory
predict.classify(model, '/path/')

OpenCV DNN 模式

这种方式没什么问题，不过我们可能不太方便使用python，那么可以选择使用opencv的dnn模块，来进行推理：

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
//
// Created by caesar kekxv on 2024/9/3.
//
#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::cvtColor(image, image, cv::COLOR_BGR2RGB);
    // 推理
    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;
  }
};

模型地址在 https://github.com/GantMan/nsfw_model/releases 进行下载；使用loadModel进行加载，调用detect进行分类。

内容来源参考

NSFW Detection Machine Learning Model