作者 | Sujan Dutta
来源 | Medium
编辑 | 代码医生团队
对某些即时通讯应用的朋友,最终采取这一聊天的截图,然后将其发送给他们。如果发送或接收了大量这些屏幕截图,那么最终手机的大部分内存都将被阻塞。在保留重要图像安全的同时查找和删除这些屏幕快照是一项非常耗时的任务。因此想用机器学习来完成这项工作
理念
从普通图像中检测聊天屏幕截图的任务可以表述为经典的二进制图像分类问题!可以使用卷积神经网络(CNN)来完成这项工作。CNN的输入层将是一幅图像,输出层将仅包含一个神经元,告诉输入图像是正常图像还是聊天屏幕截图。在接下来的部分中,将介绍构建模型的所有细节。
数据采集
在机器学习中,一切都始于数据。在此分类问题中,有两个类:'聊天'和'不聊天'。第一个表示聊天屏幕截图,另一个表示普通图像。因此从不同的消息传递应用程序(如WhatsApp,Messenger,Instagram等)中收集了与朋友聊天的屏幕截图。从手机和互联网上收集了一些人,地点,风景的随机图像。总共拍摄了660张图像。请注意对于许多更棘手的问题,这不是足够的数据量。
训练测试拆分
将80%的数据用于训练,其余的用于测试。为了能够在Keras中使用flow_from_directory函数,将数据整理成如下:
数据文件夹树
建立模型
每个CNN都由两个主要部分组成:卷积基础和完全连接网络。在卷积基础中,使用了两个卷积块,每个包含32个过滤器。内核大小为3 * 3。第一卷积层的输入尺寸为64 * 64 * 3(大小为64 px * 64 px的 RGB图像)。每个卷积块后跟一个大小为2 * 2的max_pooling层。Relu激活功能用于卷积层。卷积块的输出被展平为一个向量,以将其传递到完全连接的网络。隐藏层由128个组成神经元。该层的激活功能再次是Relu。输出层(即最后一层)仅包含一个神经元,它将告诉我们结果。由于这是一个二进制分类问题,因此我在这一层中使用了S形函数,该函数输出介于0到1之间的数字(p),表示输入图像属于“聊天”类别的概率(如果p≤0.5,则聊天否则为“否”聊天”)。
代码语言:javascript复制classifier = Sequential()
classifier.add(Convolution2D(32, 3, 3, input_shape = (64, 64, 3), activation = 'relu'))
classifier.add(MaxPooling2D(pool_size = (2, 2)))
classifier.add(Convolution2D(32, 3, 3, activation = 'relu'))
classifier.add(MaxPooling2D(pool_size = (2, 2)))
classifier.add(Flatten())
classifier.add(Dense(output_dim = 128, activation = 'relu'))
classifier.add(Dense(output_dim = 1, activation = 'sigmoid'))
模型的架构
馈送数据
由于数据是以上述特定方式组织的,因此现在可以使用ImageDataGenerator类和Keras的flow_from_directory方法来扩充数据并将其输入模型。首先创建一个ImageDataGenerator对象。在这个对象的帮助下,使用了缩放,剪切,翻转转换来增强数据。为了规范化像素值,应将图像重新缩放为1 / 255.0倍。现在目录路径,class_mode和target_size作为flow_from_directory方法的参数传递,该方法有助于将数据馈送到模型。必须执行两次此过程(一次用于训练数据,另一次用于测试数据)。这里要记住的重要一点是,仅应扩充训练数据,而不应扩充测试数据。
代码语言:javascript复制train_datagen = ImageDataGenerator(rescale = 1./255, shear_range = 0.2, zoom_range = 0.2,
horizontal_flip = True)
test_datagen = ImageDataGenerator(rescale=1./255)
training_set = train_datagen.flow_from_directory('training_set', target_size=(64, 64),
batch_size=32, class_mode='binary')
test_set = test_datagen.flow_from_directory('test_set', target_size=(64, 64),
batch_size=32, class_mode='binary')
训练
现在是模型学习的部分。这里需要一个优化器,因为学习无非就是通过更新模型的权重和偏差来优化成本函数。在这种情况下,选择了Adam优化器。成本函数是binary_crossentropy(因为这是二进制分类)。Keras提供了一个名为fit_generator的函数,可用于运行训练。在这里还可以设置时期数,steps_per_epoch和validation_steps。由于数据是比较小的,因此用steps_per_epoch =训练示例数和validation_steps =试验实施例号。
代码语言:javascript复制classifier.fit_generator(
training_set,
steps_per_epoch=528, #total number of images in training set
epochs=5,
validation_data=test_set,
validation_steps=132 #total number of images in test set
)
结果
仅5个星期后,该模型即可达到99%的训练准确度和98%的测试准确度。保存模型后,便可以根据需要多次使用它。为了能够使用此模型预测新图像,必须将图像重塑为64 * 64 * 3并标准化像素。该脚本完成了这项工作。
https://github.com/Suji04/Chat_ScreenShot_Classifier/blob/master/load model and predict.py
来自作者收藏的图片
在此处找到完整的代码。
https://github.com/Suji04/Chat_ScreenShot_Classifier
要使用此模型对手机上某个文件夹的所有图像进行分类, 只需要遍历该文件夹并将一次图像传递给该模型即可。
代码语言:javascript复制
import glob
for img_file in glob.iglob(“dir_name/*”):
new_image = load_image(img_file)
pred = classifier.predict(new_image)
if pred<.5 : print(“chat”)
else : print(“not chat”)
还在好奇吗?观看最近制作的视频