Machine Learning, Нейронные сети

Дообучение нейросети для поиска лиц в медицинских масках

Время прочтения: 5 мин.

Первое, что нам понадобится — это датасет (выборка фотографий) людей с масками и без для дообучения нейросети MobileNetV2, которая находится в открытом доступе. У меня был датасет в количестве 981 фотографий людей в масках и столько же без, одних и тех же людей.

Хотелось бы отметить один важный момент, что нейросеть MobileNetV2,  можно использовать практически для любой классификации, например, можно было дообучить её для определения пола, или попробовать автоматически определять в очках человек или нет, именно поэтому мы замораживаем все базовые слои модели, а в верхний слой подаём то, что нужно классифицировать. Но мы остановимся на поиске медицинской маски, как наиболее актуальной в настоящее время.

Итак, разместим наш датасет из 1962 фотографий в двух каталогах в папке dataset в масках в «WithMask» и без маски в «Withoutmask» соответственно. В каждой по 981 фотографии. Ещё одно важное замечание, это то, что дообучаем мы именно на лицах, а не просто, что человек на изображении в маске или без, хотя можно было и так.

Далее импортируем необходимые библиотеки:

from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.layers import AveragePooling2D
from tensorflow.keras.layers import Dropout
from tensorflow.keras.layers import Flatten
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import Input
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
from tensorflow.keras.preprocessing.image import img_to_array
from tensorflow.keras.preprocessing.image import load_img
from tensorflow.keras.utils import to_categorical
from sklearn.preprocessing import LabelBinarizer
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
from imutils import paths
import matplotlib.pyplot as plt
import numpy as np
import argparse
import os

# Указываем начальные гиперпараметры

Первый — это скорость обучения, он означает, с какой скоростью мы двигаемся по направлению к минимуму функции потерь

INIT_LR = 0,004

Второй — это количество эпох, одна эпоха это один  проход обучения на всем наборе данных

EPOCHS = 20

Третий — это размер пакета или батча, означает количество данных в одной партии.

BS = 32

Подбор гиперпараметров лежит за пределами данной статьи, скажем только то, что нет универсального способа, и в каждом случае надо их подбирать эмпирически.

imagePaths = list(paths.list_images (r'C:\dataset'))  # В этой папке хранятся два каталога с масками и без
data , labels = [] , []
for imagePath in imagePaths:
	# Извлечение класса из директории (с маской или без)
	label = imagePath.split(os.path.sep)[-2]
	# Загружам входное изображение 224х224 и обрабатываем его
	image = load_img(imagePath, target_size = (224, 224))
	image = img_to_array(image)
	image = preprocess_input(image)
	# Обновляем список файлов и классов
	data.append(image)
	labels.append(label)
# Переводим в NumPy массив
data = np.array(data, dtype="float32")
labels = np.array(labels)
# Переводим классы в бинарный вид, т.е. 0 без маски 1 с маской
lb = LabelBinarizer()
labels = lb.fit_transform(labels)
labels = to_categorical(labels)
# Разобьём датасет  на тренировочный и тестовый 80% на 20%;
(trainX, testX, trainY, testY) = train_test_split(data, labels,  test_size = 0.20, stratify = labels, random_state = 42)
# Аугментация датасета путем поворота изображений
aug = ImageDataGenerator(rotation_range = 20, zoom_range = 0.15,
	width_shift_range = 0.2, height_shift_range = 0.2, shear_range=0.15, horizontal_flip = True, fill_mode = "nearest")
# Загружаем базовую модель c предварительно обученными весами
path_weights = ‘mobilenet_v2_weights_tf_dim_ordering_tf_kernels_1.0_224_no_top.h5'    
baseModel = MobileNetV2(weights=path_weights, 
include_top=False, input_tensor=Input(shape=(224, 224, 3))
Запишем в нашу модель внешний слой из базовой модели
headModel = baseModel.output
headModel = AveragePooling2D(pool_size = (7, 7))(headModel)
headModel = Flatten(name = "flatten")(headModel)
headModel = Dense(128, activation = "relu")(headModel)
headModel = Dropout(0.5)(headModel)
headModel = Dense(2, activation = "softmax")(headModel)
model = Model(inputs = baseModel.input, outputs = headModel)
# Заморозка слоев базовой модели 
for layer in baseModel.layers:
	layer.trainable = False
# Скомпилируем нашу модель
opt = Adam(lr = INIT_LR, decay = INIT_LR / EPOCHS)
model.compile(loss = "binary_crossentropy", optimizer = opt, metrics = ["accuracy"])
# Тренируем нашу сеть
H = model.fit( aug.flow(trainX, trainY, batch_size = BS), steps_per_epoch = len(trainX) // BS,
	validation_data = (testX, testY), validation_steps = len(testX) // BS, epochs = EPOCHS)
# Делаем предсказание на тестовой выборке
predIdxs = model.predict(testX, batch_size = BS)
#  Для каждого изображения в тестовом наборе, ищем максимальную вероятность
predIdxs = np.argmax(predIdxs, axis=1)
# Показать отчет обучения 
print(classification_report(testY.argmax(axis = 1), predIdxs, target_names = lb.classes_))

# Сохраняем модель на диск и загружаем её

model.save('model_mask_FACE', save_format = "h5")
model_mask = tf.keras.models.load_model('model_mask_FACE)

Поиск маски на лице, на  примере

# Найдем лицо на изображении, используя библиотеку MTCNN

frame = cv2.cvtColor(cv2.imread(‘house.png'), cv2.COLOR_BGR2RGB)
frame_image = Image.fromarray(frame)
boxes, probs, landmarks = mtcnn.detect(frame_image, landmarks = True)
x1, y1, x2, y2 = [int(bx) for bx in boxes[0]]
image = Image.fromarray(frame[y1:y2, x1:x2]).resize((224,224))
face = img_to_array(image)

# Обработка изображения для загрузки в модель

face = preprocess_input(face)
face = np.expand_dims(face, axis=0)

# Загрузка лица  в нашу модель

(mask, withoutMask) = model_mask.predict(face)[0]
image = cv2.imread(‘house.png’)

# Прорисовка рамки и вероятности

if mask > withoutMask and max(mask, withoutMask) > 0.8: # уверенность
    label = "Mask" if mask > withoutMask else "No Mask"
    color = (0, 122, 0) if label == "Mask" else (0, 0, 122)
    label = "{}: {:.2f}%".format(label, max(mask, withoutMask) * 100)
    cv2.putText(image, label, (x1, y1 - 10),cv2.FONT_HERSHEY_SIMPLEX, 2, color, 5)
    cv2.rectangle(image, (x1, y1), (x2, y2), color, 5) 
    y = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

Итак, в данной статье мы показали, как дообучить нейросеть MobileNetV2, чтобы классифицировать изображения лиц людей на находящихся в медицинских масках и без них.

Советуем почитать