如何使用機器學習來檢測手機上的聊天屏幕截圖

作者 | 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,則聊天否則為「否」聊天」)。

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方法的參數傳遞,該方法有助於將數據饋送到模型。必須執行兩次此過程(一次用於訓練數據,另一次用於測試數據)。這裡要記住的重要一點是,僅應擴充訓練數據,而不應擴充測試數據。

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 =試驗實施例號。

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%20model%20and%20predict.py

來自作者收藏的圖片

在此處找到完整的代碼。

https://github.com/Suji04/Chat_ScreenShot_Classifier

要使用此模型對手機上某個文件夾的所有圖像進行分類, 只需要遍歷該文件夾並將一次圖像傳遞給該模型即可。

  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」)

還在好奇嗎?觀看最近製作的視頻