В этой статье я расскажу о том, как создать модель классификации изображений, которая обнаруживает изображения NSFW (небезопасные для работы).
Изображения NSFW означают изображения с возрастными ограничениями, такие как обнаженные тела, или изображения, оскорбляющие общественность.
В этом проекте мы создадим модель, которая будет классифицировать изображения на рисованные, хентай, нейтральные, порно, сексуальные
рисунок — мультфильм или нарисованное от руки изображение
хентай — жанр японского аниме, характеризующийся явно сексуализированными персонажами и откровенно сексуальными изображениями.
нейтральные — обычные изображения, не являющиеся оскорбительными.
порно — порнографические изображения
сексуальность — частично откровенные изображения
Сбор данных
Из-за характера набора данных мы не можем получить все изображения с определенных веб-сайтов, таких как Kaggle и других.
Но мы можем использовать текстовый файл, содержащий все ссылки на изображения, для этого проекта я использовал текстовые файлы URL Alex Kim, в которых есть все URL-адреса изображений.
ссылка на репо https://github.com/alex000kim/nsfw_data_scraper
После получения текстовых файлов мы можем запустить приведенный ниже код, чтобы создать необходимую папку и загрузить каждое изображение в определенную папку.
folders = ['drawings','hentai','neutral','porn','sexy']
urls = ['urls_drawings.txt','urls_hentai.txt','urls_neutral.txt','urls_porn.txt','urls_sexy.txt']
names = ['d','h','n','p','s']
for i,j,k in zip(folders,urls,names):
try:
#Specify the path of the folder that has to be made
folder_path = os.path.join('your directory',i)
os.mkdir(folder_path)
except:
pass
#setup the path of url text file
url_path = os.path.join('Datasets_Urls',j)
my_file = open(url_path, "r")
data = my_file.read()
#create a list with all urls
data_into_list = data.split("\n")
my_file.close()
icount = 0
for ii in data_into_list:
try:
#create a unique image names for each images
image_name = 'image'+str(icount)+str(k)+'.png'
image_path = os.path.join(folder_path,image_name)
#download it using the library
urllib.request.urlretrieve(ii, image_path)
icount+=1
except Exception as e:
pass
#this below code is done to make the count of the image same for all the data
#you can use a big number if you are building a more complex model or if you have a good system
if icount == 2000:
break
Здесь переменная папки обозначает каждый класс и создает папки в нужном каталоге, переменная urls используется для получения текстового файла URL (вы можете изменить это в соответствии с именем вашего текстового файла), переменная имени используется для создания уникального имени для каждого изображение.
Вышеупомянутое позволит вам загрузить 2000 изображений для каждого класса, вы можете отредактировать последнее условие «если», чтобы изменить размер изображений.
Вы можете использовать другие источники для сбора данных, например, расширение «Загрузить все изображения» в Chrome, которое позволяет загружать все изображения, которые вы искали (не беспокойтесь о повторяющихся изображениях, мы удалим их позже).
Подготовка данных
Нежелательный тип изображений
Папки могут содержать файлы других типов, поэтому сначала мы должны удалить все эти типы файлов.
image_exts = ['jpeg','.jpg','bmp','png']
path_list = ['drawings','hentai','neutral','porn','sexy']
cwd = os.getcwd()
def remove_other_images(path_list):
for ii in path_list:
data_dir = os.path.join(cwd,'DataSet',ii)
for image in os.listdir(os.path.join(data_dir)):
image_path = os.path.join(data_dir,image_class,image)
try:
img = cv2.imread(image_path)
tip = imghdr.what(image_path)
if tip not in image_exts:
print('Image not in ext list {}'.format(image_path))
os.remove(image_path)
except Exception as e:
print("Issue with image {}".format(image_path))
remove_other_images(path_list)
Приведенный выше код удаляет изображения с расширениями, отличными от указанного формата, переменная path_list обозначает имя папки.
Повторяющиеся изображения
Изображения могут содержать много повторяющихся изображений, поэтому мы должны удалить дубликаты из каждой папки. Мы можем запустить приведенный ниже код, чтобы удалить повторяющиеся изображения.
cwd = os.getcwd()
path_list = ['drawings','hentai','neutral','porn','sexy']
def remove_dup_images(path_list):
for ii in path_list:
os.chdir(os.path.join(cwd,'DataSet',ii))
filelist = os.listdir()
duplicates = []
hash_keys = dict()
for index, filename in enumerate(filelist):
if os.path.isfile(filename):
with open(filename,'rb') as f:
filehash = hashlib.md5(f.read()).hexdigest()
if filehash not in hash_keys:
hash_keys[filehash] = index
else:
duplicates.append((index,hash_keys[filehash]))
for index in duplicates:
os.remove(filelist[index[0]])
print('{} duplicates removed from {}'.format(len(duplicates),ii))
remove_dup_images(path_list)
Здесь мы используем кодировку hashlib.md5 для поиска повторяющихся изображений в каждом классе.
md5 создает уникальное хеш-значение, похожее на отпечаток пальца, для каждого изображения, и если хеш-значение повторяется (дубликат изображения), то оно добавляется в список, который позже будет использоваться для их удаления.
Изображения, поддерживающие Tensorflow
Далее нам нужно найти файлы, которые не поддерживаются TensorFlow, если какое-либо изображение не поддерживается TensorFlow, это приведет к ошибке при обучении модели.
Чтобы проверить, будет ли поддерживаться изображение, мы можем попытаться закодировать это изображение с помощью Tensorflow, и если это приведет к ошибке, мы можем просто удалить ее, используя обработку исключений.
Используйте приведенный ниже код для удаления неподдерживаемых файлов.
import tensorflow as tf
os.chdir('{data-set} directory')
cwd = os.getcwd()
for ii in path_list:
os.chdir(os.path.join(cwd,ii))
filelist = os.listdir()
for image_file in filelist:
with open(image_file, 'rb') as f:
image_data = f.read()
# Check the file format
_, ext = os.path.splitext(image_file)
if ext.lower() not in ['.jpg', '.jpeg', '.png', '.gif', '.bmp']:
print('Unsupported image format:', ext)
os.remove(os.path.join(cwd,ii,image_file))
else:
# Decode the image
try:
image = tf.image.decode_image(image_data)
except:
print(image_file)
print("unspported")
os.remove(os.path.join(cwd,ii,image_file))
Разделение данных
После очистки набора данных мы должны разделить данные.
Чтобы разделить папки поезда, проверки и теста и вручную добавить изображения с папками внутри него, мы также используем 80% для обучения, 10% для проверки и 10% для теста.
Разделение выполняется перед моделированием, чтобы предотвратить утечку данных, потому что мы должны выполнить увеличение данных для лучшего результата модели.
Создание модели
Для этой модели мы можем использовать алгоритмы обучения с передачей, такие как ResNet, VGG16 или сеть Xception.
Трансферное обучение — это мощная техника, которая позволяет применять знания, полученные в одной задаче или области, к другой связанной задаче или области. Вместо того, чтобы обучать модель с нуля для каждой конкретной задачи, трансферное обучение использует уже существующие знания, полученные из другой, но связанной задачи.
Такой подход значительно сокращает количество данных, времени и вычислительных ресурсов, необходимых для обучения новых моделей.
Импорт модулей
import tensorflow as tf import os import numpy as np import matplotlib.pyplot as plt from sklearn.utils import shuffle import hashlib from imageio import imread import numpy as np from tensorflow.keras.preprocessing.image import ImageDataGenerator from tensorflow.keras.applications.vgg16 import VGG16 from tensorflow.keras.applications.vgg16 import preprocess_input from tensorflow.keras.layers import Flatten,Dense,Input from tensorflow.keras.models import Model,Sequential from keras import optimizers
Для изображений мы можем установить размер по умолчанию 224 224.
IMAGE_SIZE = [224,224]
Увеличение данных
Для увеличения данных мы можем использовать библиотеку ImageDataGenerator, которая генерирует разные изображения в соответствии с заданным нами параметром.
Увеличение данных выполняется для увеличения размера набора данных. ImageDataGenerator создает новые изображения на основе заданного параметра и использует его для обучения (примечание: при использовании ImageDataGenerator исходные данные не будут использоваться для обучения).
train_datagen = ImageDataGenerator(
rescale=1./255,
preprocessing_function=preprocess_input,
rotation_range=40,
width_shift_range=0.2,
height_shift_range=0.2,
shear_range=0.2,
zoom_range=0.2,
horizontal_flip=True,
fill_mode='nearest')
Для тестовых данных нам нужно только масштабировать их.
test_datagen = ImageDataGenerator(rescale=1./255)
Загрузить данные
train_set = train_datagen.flow_from_directory('DataSet/train',
target_size=(224,224),
batch_size=32,
class_mode='sparse')
val_set = train_datagen.flow_from_directory('DataSet/validation',
target_size=(224,224),
batch_size=32,
class_mode='sparse')
Здесь целевой размер и размер пакета установлены на 224 224 и 32 соответственно, а режим класса установлен на разреженный, потому что проблема заключается в классификации с несколькими метками.
Загружается модель
Поскольку мы используем модель vgg16, сначала нам нужно настроить ее параметры. Приведенный ниже код загружает и настраивает предварительно обученную модель vgg16.
vgg = VGG16(input_shape=IMAGE_SIZE+[3],weights='imagenet',include_top=False)
for layer in vgg.layers:
layer.trainable = False
Приведенный выше код ограничивает обновление весов во всех слоях, присутствующих в модели vgg16, это делается потому, что модель уже обучена определенным весам.
Теперь нам нужно настроить выходной слой и подключить vgg16 к выходному слою.
x = Flatten()(vgg.output) prediction = Dense(5,activation='softmax')(x) model = Model(inputs=vgg.input, outputs=prediction) model.summary()
и используйте активацию «softmax», поэтому у нас есть несколько классов
from tensorflow.keras.metrics import MeanSquaredError
from tensorflow.keras.metrics import CategoricalAccuracy
adam = optimizers.Adam()
model.compile(loss='sparse_categorical_crossentropy',
optimizer=adam,
metrics=['accuracy',MeanSquaredError(name='val_loss'),CategoricalAccuracy(name='val_accuracy')])
Для модели мы можем использовать оптимизатор Адама и потери как «sparse_categorical_crossentropy», чтобы мы могли получать метки, закодированные как целые числа, а не с горячим кодированием (также используйте sparse_categorical_crossentropy, потому что мы загружали данные в режиме разреженного класса).
from datetime import datetime
from keras.callbacks import ModelCheckpoint
log_dir = 'vg_log'
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir = log_dir)
start = datetime.now()
history = model.fit_generator(train_set,
validation_data=val_set,
epochs=100,
steps_per_epoch=len(train_set)// batch_size,
validation_steps=len(val_set)//batch_size,
callbacks=[tensorboard_callback],
verbose=1)
duration = datetime.now() - start
print("Time taken for training is ",duration)
Мы обучаем эту модель на 100 эпох.
После обучения модели она получила результат точности проверки 80%. и f1 оценка 93%
Приведенная ниже функция возьмет список изображений и сделает прогноз в соответствии с ним.
import numpy as np
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
from scipy.ndimage import gaussian_filter
def print_classes(images,model):
classes = ['Drawing','Hentai','Neutral','Porn','Sexual']
fig, ax = plt.subplots(ncols=len(images), figsize=(20,20))
for idx,img in enumerate(images):
img = mpimg.imread(img)
resize = tf.image.resize(img,(224,224))
result = model.predict(np.expand_dims(resize/255,0))
result = np.argmax(result)
if classes[result] == 'Porn':
img = gaussian_filter(img, sigma=6)
elif classes[result] == 'Sexual':
img = gaussian_filter(img, sigma=6)
elif classes[result] == 'Hentai':
img = gaussian_filter(img, sigma=6)
ax[idx].imshow(img)
ax[idx].title.set_text(classes[result])
li = ['test1.jpeg','test2.jpeg','test3.jpeg','test4.jpeg','test5.jpeg']
print_classes(li,model)
Вы также можете найти версию этого проекта, развернутую на Django, в моем репозитории GitHub.
