利用已经训练好的 TensorFlow 模型(.pb)进行图片鉴黄和筛选

以前和同学一起做过一个有人工智能鉴黄功能的匿名论坛,现在想自己也整一个模型来玩

在CSDN上看见大佬 王小鹏鹏 的模型,看了一下,这不就是当初同学和我用的模型吗?

链接放出来:https://blog.csdn.net/kingroc/article/details/89519293?ops_request_misc=%257B%2522request%255Fid%2522%253A%2522158864595119725219950302%2522%252C%2522scm%2522%253A%252220140713.130102334..%2522%257D&request_id=158864595119725219950302&biz_id=0&utm_medium=distribute.pc_search_result.none-task-blog-2~all~first_rank_v2~rank_v25-3

这个项目里文件很多,看得出来作者应该是一直有在更新

上面这个项目源码我从Github上下载到本地运行,由于platform等等原因运行不了(源码利用本地socket和网站),就改成了对本地文件夹内图片进行遍历

(实际上我有写过爬取图片的爬虫,正好能从爬取图片的文件夹中导入,看看效果)

话不多说,上🐎:

main_classifier.py的代码如下:

import tensorflow as tf
import os
import numpy as np
from PIL import Image

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
lines = tf.io.gfile.GFile('./inception_model/output_labels.txt').readlines()
uid_to_human = {}
for uid,line in enumerate(lines):
    line = line.strip('\n')
    uid_to_human[uid] = line

with tf.io.gfile.GFile('./inception_model/output_graph.pb', 'rb') as f:
    graph_def = tf.compat.v1.GraphDef()
    graph_def.ParseFromString(f.read())
    tf.import_graph_def(graph_def, name='')

def getResult(picPath):
    print("正在对该图片进行分析:{}".format(picPath))
    with tf.compat.v1.Session() as sess:
        softmax_tensor = sess.graph.get_tensor_by_name('final_result:0')
        image_data = tf.io.gfile.GFile(picPath, 'rb').read()
        predictions = sess.run(softmax_tensor, {'DecodeJpeg/contents:0': image_data})
        predictions = np.squeeze(predictions)

        top_k = predictions.argsort()[::-1]
        result = ''
        for node_id in top_k:
            if node_id not in uid_to_human:
                human_string = ''
            human_string = uid_to_human[node_id]
            score = predictions[node_id]
            buf = ('%s (score = %.2f%%) \n' % (human_string, score*100))
            result = result + buf
        print(result)
        confirm = input("是否显示该图片(y\\n)?")
        if confirm.lower() == "y":
            im = Image.open(picPath)
            im.show()

targetting.py代码如下:

import os
import main_classifier

def readNames(sourcePath):
    names = os.listdir(sourcePath)
    return names

sourcePath = "D:\爬取图片\\"

for name in readNames(sourcePath):
    for pic in readNames(sourcePath + name):
        main_classifier.getResult(sourcePath + name + '\\'+pic)

D:\爬取图片\   是我图片爬虫的地址,里面按不同网站分文件夹存放.JPG格式的图片

使用和原作者完全一致的训练好的模型,存放在和上述两.py文件同级的inception_model文件夹中,有兴趣的朋友可以访问原作者的Github下载源码,里面还有很多其他好玩的东西。

posted @ 2020-05-05 14:47  neumy  阅读(612)  评论(0编辑  收藏  举报