利用已经训练好的 TensorFlow 模型(.pb)进行图片鉴黄和筛选
以前和同学一起做过一个有人工智能鉴黄功能的匿名论坛,现在想自己也整一个模型来玩
在CSDN上看见大佬 王小鹏鹏 的模型,看了一下,这不就是当初同学和我用的模型吗?
这个项目里文件很多,看得出来作者应该是一直有在更新
上面这个项目源码我从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下载源码,里面还有很多其他好玩的东西。