Целью данной работы является разработка модели компьютерного зрения для распознавания и классификации составляющих микроструктуры стали.
В металловедении принято называть составляющие микроструктуры зернами. Важность определения типа зерен микроструктуры продиктована влиянием размера и соотношения объемных долей зерен микроструктуры на механические свойства стали. В основном определение типа микроструктуры выполняется экспертами «на глаз», что в ряде случаев приводит к разногласиям в оценке.
В данной работе модель обучалась классифицировать такие составляющие микроструктуры, как феррит, бейнит и перлит. Обучение модели проводилось на микроструктуре стали в состоянии после прокатки без проведения дополнительной термической обработки. Это важная оговорка, т.к. различие в специфическом «узоре» микроструктуры между зернами металла с термообработкой и без термообработки является существенным.
На рисунке 1 приведено изображение микроструктуры стали с выделенным зерном.
import torchvision.datasets as datasets
import torchvision.transforms as transforms
train = datasets.ImageFolder("path/train", transform = transformations)
val = datasets.ImageFolder("path/test", transform = transformations_s)
train_loader = torch.utils.data.DataLoader(train, batch_size=4, shuffle=True)
val_loader = torch.utils.data.DataLoader(val, batch_size =4, shuffle=False)С помощью функции DataLoader были созданы загрузчики данных отдельными партиями — batch.
С целью компенсации небольшого количества данных был применен подход «Transfer Learning», использующий модель, обученную на большом количестве изображений. Последний слой предобученной модели был заменен на классификатор для 3х классов микроструктур. В качестве предобученной модели выбрана densenet161, для инициализации предварительного обучения установлен режим pretrained = True.
num_labels = 3
classifier = nn.Sequential(nn.Linear(512, num_labels),
nn.LogSoftmax(dim=1))
# Заменим последний полносвязный слой модели на наш классификатор
model.fc = classifierИзображения перед загрузкой в модель трансформированы в тензоры и тензоры в свою очередь нормализованы.
transformations = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485], std=[0.229])])Для оценки точности классификации модели была выбрана метрика «accuracy». При обучении на baseline-модели получен разброс по «accuracy» в 0,1 и отмечено, что от эпохи к эпохе эта метрика не улучшается. Сделан вывод, что модель недообучается. С целью повышения обучаемости модели было решено разнообразить набор данных с помощью аугментаций и сократить количество признаков путем бинаризации изображений. Кроме того, решено было заменить предобученную модель с densenet161 на ResNet18, использующую остаточное соединение (переброс исходного тензора через слои).
На рисунке 2 представлены бинаризованные изображения, сконвертированные обратно из тензора после наложения перспективных искажений, вертикальных и горизонтальных отражений.
Бинаризация проведена методом Оцу. Метод Оцу (Otsu's Method) использует гистограмму изображения для расчета порога отнесения пикселя к 0 или 1.
Для выявления наиболее эффективного вида аугментаций они применялись по отдельности. Результаты применения различных видов аугментаций приведены в таблице 1.
Таблица 1. Применение аугментаций
Из таблицы 1 видно, что наилучшие результаты получены при применении вертикальных и горизонтальных отражений, а также от перспективных искажений. Поэтому далее для трансформации изображений будем применять только эти виды аугментаций.
Для определения оптимальной скорости обучения была применена функция lr_scheduler (далее планировщик), позволяющая менять скорость обучения каждые 2 эпохи с мультипликатором 0,1. Мультипликатор (gamma=0,1) означает, что скорость обучения будет снижаться в 10 раз.
# Зададим начальную скорость обучения = 0.01
optimizer = optim.Adam(model.fc.parameters(), lr=1e-2)
# Импортируем шедулер
from torch.optim import lr_scheduler
# Функция снижения скорости обучения встроена в данный шедулер.
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.1)Апробация планировщика производилась с использованием двух предобученных моделей ResNet18 и ResNet34 отличающихся количеством слоев. Результаты апробации представлены в таблице 2.
Таблица 2. Сводные результаты
Получаем оптимальные результаты по accuracy на значениях LR 1e-5 и 1e-6.
Проверим, насколько точно модель сможет предсказать составляющую микроструктуры на реальных изображениях. Были отобраны изображения трех составляющих микроструктуры в градациях серого, которые до этого не были ни в наборе данных для обучения, ни в наборе данных для тестирования.
Результаты применения модели, основанной на предобученных моделях ResNet18 и ResNet34, представлены на рисунках 3 и 4 соответственно.
Из рисунков 3 и 4 мы видим, что в одном случае модель, основанная на ResNet34, допустила ошибку, указав, что микроструктура, фактически представляющая из себя бейнит, с вероятностью 100% является перлитом.
Проверка модели еще на одном изображении бейнита показала, что для изображений с бейнитом работа модели, основанной на Resnet34, нестабильна. Возможно, это связано с небольшой выборкой для датасета. Кроме того, специфический «узор» бейнита имеет большее разнообразие в отличии от феррита и перлита, что также указывает на необходимость в увеличении датасета.
По результатам исследования можно сделать следующие выводы:
Предсказание составляющих микроструктуры с подачей на вход модели случайных изображений происходит с вероятностью 93 – 100%;
Учитывая, что модель, основанная на ResNet34, делает ошибки при определении типа зерна «бейнит», необходимо увеличить набор данных и поработать с изображениями бейнита.
В дальнейшем планируется на основе классификации составляющих микроструктуры разработать модель, определяющую тип микроструктуры на реальных снимках, состоящих из 2х и более составляющих.
Также интересной задачей выглядит обучение модели для классификации составляющих микроструктуры металла в термообработанном состоянии.
Автор выражает благодарность людям, которые дали возможность осуществить данное исследование, а именно Антону Витвицкому за помощь при подборе методов обработки изображений для компьютерного зрения и за работу над ошибками, а также Марию Тихонову, являющуюся моим первым Учителем по машинному обучению. По данной ссылке можно подробнее узнать о курсе, руководителем которого является Мария Тихонова.