简介

使用雅虎开源的 TensorFlow 2 Open-NSFW 模型

NSFW：not safe for work，工作场所不宜

1.jpg

2.jpg

自己去下

安装

Python 3.7 及以上，本人尝试 Python 3.6 失败

pip install opennsfw2

图像

import opennsfw2 as n2# 单张预测
image_path = '1.jpg'
nsfw_probability = n2.predict_image(image_path)
print(nsfw_probability)
# 0.16282974183559418# 批量预测
image_paths = ['1.jpg', '2.jpg']
nsfw_probabilities = n2.predict_images(image_paths)
print(nsfw_probabilities)
# [0.16282965242862701, 0.8638442158699036]

将自动下载预训练模型 open_nsfw_weights.h5 到 C:\Users\Administrator\.opennsfw2\weights

视频

import opennsfw2 as n2video_path = '1.mp4'
elapsed_seconds, nsfw_probabilities = n2.predict_video_frames(video_path)
for second, probability in zip(elapsed_seconds, nsfw_probabilities):print(f'{second:.2f}s: {probability * 100:.0f} %')
# 0.03s: 1%
# ...
# 10.01s: 87.00%
# ...
# 10.64s: 69.00%

预加载

import numpy as np
from PIL import Image
from opennsfw2._model import make_open_nsfw_model
from opennsfw2._image import preprocess_image, Preprocessingimage_path = '1.jpg'
image = preprocess_image(Image.open(image_path), Preprocessing.YAHOO)
model = make_open_nsfw_model()
nsfw_probability = float(model.predict(np.expand_dims(image, 0), batch_size=1)[0][1])
print(nsfw_probability)
# 0.16282974183559418

车速检测器

import time
import numpy as np
import tkinter as tk
from pathlib import Path
from tkinter import filedialog
from tkinter import messagebox
from PIL import ImageTk, Imagefrom opennsfw2._model import make_open_nsfw_model
from opennsfw2._image import preprocess_image, Preprocessingbegin = time.time()
model = make_open_nsfw_model()  # 加载模型
elapsed = time.time() - begin  # 加载模型耗时
initialdir = Path.cwd()  # 初始化目录，可切换为图片Path.home() / 'Pictures'
img = None  # 当前打开的图片def scale(size, width=None, height=None):"""获取按比例缩放后的宽高"""if not width and not height:width, height = sizeif not width or not height:_width, _height = sizeheight = width * _height / _width if width else heightwidth = height * _width / _height if height else widthreturn int(width), int(height)def img_resize(event=None):"""显示图片"""global imgif img:_img = img.resize(scale(img.size, height=win.winfo_height()))_img = ImageTk.PhotoImage(_img)label.config(image=_img)label.image = _imgdef on_closing():"""关闭事件"""if messagebox.askokcancel('关闭', '是否退出程序？'):win.destroy()def open_file(event=None):"""打开图片"""global initialdirglobal imgfile_path = filedialog.askopenfilename(title='选择图片', initialdir=initialdir,filetypes=[('image files', ('.png', '.jpg', '.jpeg', '.gif'))])if file_path:statusbar.config(text='正在加载...')statusbar.update_idletasks()begin = time.time()path = Path(file_path)initialdir = path.parentimg = Image.open(file_path)img_resize()_img = preprocess_image(Image.open(file_path), Preprocessing.YAHOO)probability = float(model.predict(np.expand_dims(_img, 0), batch_size=1)[0][1])print(probability)end = time.time()statusbar.config(text=f'{path.name} 耗时: {end - begin:.2f}s 概率: {probability * 100:.2f} %')win = tk.Tk()
win.title('黄图检测')  # 标题
menu = tk.Menu(win)
menu.add_command(label='打开', command=open_file)
win.config(menu=menu)
win.bind('<Configure>', img_resize)
win.geometry('600x300+300+300')
win.minsize(200, 200)
win.protocol('WM_DELETE_WINDOW', on_closing)
statusbar = tk.Label(win, text=f'加载模型耗时: {elapsed:.2f}s', bd=1, relief=tk.SUNKEN, anchor=tk.W, name='statusbar')
statusbar.pack(side=tk.BOTTOM, fill=tk.X)
label = tk.Label(win, text='双击打开图片')
label.bind('<Double-Button-1>', open_file)
label.pack(fill=tk.BOTH, expand=True)
win.mainloop()

效果

视频车速检测器

无声音

安装

pip install imageio-ffmpeg

import time
import threading
import tkinter as tk
from pathlib import Path
from tkinter import filedialog
from tkinter import messagebox
from timeit import default_timer as timerimport imageio
import numpy as np
from PIL import ImageTk, Imagefrom opennsfw2._model import make_open_nsfw_model
from opennsfw2._image import preprocess_image, Preprocessingbegin = time.time()
model = make_open_nsfw_model()  # 加载模型
elapsed = time.time() - begin  # 加载模型耗时
initialdir = Path.cwd()  # 初始化目录，可切换为图片Path.home() / 'Pictures'
reader = None  # 视频读取器accum_time = 0
curr_fps = 0
last_fps = 0
prev_time = timer()def on_closing():"""关闭事件"""if messagebox.askokcancel('关闭', '是否退出程序？'):win.destroy()def play():global readerglobal prev_time, accum_time, curr_fpsfor image in reader:image = Image.fromarray(image)frame_image = ImageTk.PhotoImage(image)label.config(image=frame_image)label.image = frame_image_img = preprocess_image(image, Preprocessing.YAHOO)probability = float(model.predict(np.expand_dims(_img, 0), batch_size=1)[0][1])# FPScurr_time = timer()exec_time = curr_time - prev_timeprev_time = curr_timeaccum_time = accum_time + exec_timecurr_fps = curr_fps + 1if accum_time > 1:accum_time = accum_time - 1last_fps = curr_fpscurr_fps = 0statusbar.config(text=f'概率: {probability * 100:.2f} % FPS: {last_fps}')def open_file(event=None):"""打开视频"""global initialdirglobal readerfile_path = filedialog.askopenfilename(title='选择视频', initialdir=initialdir,filetypes=[('Select files', ('.mp4', '.mkv', '.avi', '.wmv'))])if file_path:statusbar.config(text='正在加载...')statusbar.update_idletasks()path = Path(file_path)initialdir = path.parentreader = imageio.get_reader(path)thread = threading.Thread(target=play, daemon=True)thread.start()win = tk.Tk()
win.title('黄图检测')  # 标题
menu = tk.Menu(win)
menu.add_command(label='打开', command=open_file)
win.config(menu=menu)
win.geometry('1280x720+300+300')
win.minsize(200, 200)
win.protocol('WM_DELETE_WINDOW', on_closing)
statusbar = tk.Label(win, text=f'加载模型耗时: {elapsed:.2f}s', bd=1, relief=tk.SUNKEN, anchor=tk.W, name='statusbar')
statusbar.pack(side=tk.BOTTOM, fill=tk.X)
label = tk.Label(win, text='双击打开视频')
label.bind('<Double-Button-1>', open_file)
label.pack(fill=tk.BOTH, expand=True)
win.mainloop()

TODO: 有声音

方案一：tkinter + FPS + 单独播放声音

代码

方案二：pyglet

安装

pip install pyglet

代码

import ioimport pyglet
import numpy as np
from PIL import Image
from opennsfw2._model import make_open_nsfw_model
from opennsfw2._image import preprocess_image, Preprocessingmodel = make_open_nsfw_model()
filename = '1.mp4'
source = pyglet.media.load(filename)
video_format = source.video_format
width, height = video_format.width, video_format.height
title = 'Video Player'
window = pyglet.window.Window(width, height, title)
player = pyglet.media.Player()
player.queue(source)
player.play()@window.event
def on_draw():window.clear()if player.source and player.source.video_format:player.get_texture().blit(0, 0, width=width, height=height)image_data = player.get_texture().get_image_data()pitch = -(image_data.width * len('RGB'))data = image_data.get_data('RGB', pitch)_img = preprocess_image(Image.frombytes('RGB', (width, height), data, 'raw'), Preprocessing.YAHOO)probability = float(model.predict(np.expand_dims(_img, 0), batch_size=1)[0][1])print(probability)pyglet.app.run()