Прежде чем перейти к статье, хочу вам представить, экономическую онлайн игру Brave Knights, в которой вы можете играть и зарабатывать. Регистируйтесь, играйте и зарабатывайте!
Целью данной работы является разработка модели компьютерного зрения для распознавания и классификации составляющих микроструктуры стали.
В металловедении принято называть составляющие микроструктуры фазами. Важность определения типа микроструктурных фаз продиктована влиянием размера и соотношения объемных долей фаз микроструктуры на механические свойства стали. В основном определение типа микроструктуры выполняется экспертами «на глаз», что в ряде случаев приводит к разногласиям в оценке.
В данной работе модель обучалась классифицировать такие фазы, как феррит, бейнит и перлит. Обучение модели проводилось на микроструктуре стали в состоянии после прокатки без проведения дополнительной термической обработки. Это важная оговорка, т.к. различие в специфическом «узоре» микроструктуры между фазами с термообработкой и без термообработки является существенным.
На рисунке 1 приведено изображение микроструктуры стали с выделенной фазой.
![Рисунок 1. Общий вид микроструктуры стали Рисунок 1. Общий вид микроструктуры стали](https://habrastorage.org/r/w1560/getpro/habr/upload_files/ca2/382/426/ca238242679f2c275ea8d5595127fef5.png)
При подготовке данных для обучения на каждом изображении оставлено зерно только одной из фаз (феррит, бейнит или перлит). Для каждого класса микроструктуры отобрано по 100 изображений размером 224х224 в градациях серого (8 бит). Набор данных разделен на выборки train и test в соотношении 80/20. Изображения отсортированы по папкам, соответствующим каждой из фаз (классу). Функция “ImageFolder” позволяет присвоить метку для изображений, которой является название соответствующей папки.
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 представлены бинаризованные изображения, сконвертированные обратно из тензора после наложения перспективных искажений, вертикальных и горизонтальных отражений.
![Рисунок 2. Аугментации Рисунок 2. Аугментации](https://habrastorage.org/r/w1560/getpro/habr/upload_files/c29/f03/644/c29f03644b5e06c7e26bb769491ae7bc.png)
Бинаризация проведена методом Оцу. Метод Оцу (Otsu's Method) использует гистограмму изображения для расчета порога отнесения пикселя к 0 или 1.
Для выявления наиболее эффективного вида аугментаций они применялись по отдельности. Результаты применения различных видов аугментаций приведены в таблице 1.
Таблица 1. Применение аугментаций
![](https://habrastorage.org/r/w780q1/getpro/habr/upload_files/748/df5/ca7/748df5ca7a09080b613a267baa4f64f9.jpg)
Из таблицы 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. Сводные результаты
![](https://habrastorage.org/r/w780q1/getpro/habr/upload_files/28a/f27/c9a/28af27c9a01feed8f04b1546e41ee169.jpg)
Получаем оптимальные результаты по accuracy на значениях LR 1e-5 и 1e-6.
Проверим, насколько точно модель сможет предсказать составляющую микроструктуры на реальных изображениях. Были отобраны изображения трех составляющих микроструктуры в градациях серого, которые до этого не были ни в наборе данных для обучения, ни в наборе данных для тестирования.
Результаты применения модели, основанной на предобученных моделях ResNet18 и ResNet34, представлены на рисунках 3 и 4 соответственно.
![](https://habrastorage.org/r/w780q1/getpro/habr/upload_files/5dd/930/a56/5dd930a56d9aeb0e265bf9053b4d8188.jpg)
![Рисунок 4. Применение модели, основанной на ResNet34 Рисунок 4. Применение модели, основанной на ResNet34](https://habrastorage.org/r/w780q1/getpro/habr/upload_files/e18/918/efc/e18918efcb2f4dea03361f93a90449b0.jpg)
Из рисунков 3 и 4 мы видим, что в одном случае модель, основанная на ResNet34, допустила ошибку, указав, что микроструктура, фактически представляющая из себя бейнит, с вероятностью 100% является перлитом.
Проверка модели еще на одном изображении бейнита показала, что для изображений с бейнитом работа модели, основанной на Resnet34, нестабильна. Возможно, это связано с небольшой выборкой для датасета. Кроме того, специфический «узор» бейнита имеет большее разнообразие в отличии от феррита и перлита, что также указывает на необходимость в увеличении датасета.
![Рисунок 5. Дополнительное изображение с бейнитом для модели, основанной на ResNet34 Рисунок 5. Дополнительное изображение с бейнитом для модели, основанной на ResNet34](https://habrastorage.org/r/w1560/getpro/habr/upload_files/2d1/d66/696/2d1d666968d3d3b7f40c42648750ce03.png)
По результатам исследования можно сделать следующие выводы:
Предсказание составляющих микроструктуры с подачей на вход модели случайных изображений происходит с вероятностью 93 – 100%;
Учитывая, что модель, основанная на ResNet34, делает ошибки при определении фазы «бейнит», необходимо увеличить набор данных и поработать с изображениями бейнита.
В дальнейшем планируется на основе классификации составляющих микроструктуры разработать модель, определяющую тип микроструктуры на реальных снимках, состоящих из 2х и более составляющих.
Также интересной задачей выглядит обучение модели для классификации составляющих микроструктуры металла в термообработанном состоянии.
Автор выражает благодарность людям, которые дали возможность осуществить данное исследование, а именно Антону Витвицкому за помощь при подборе методов обработки изображений для компьютерного зрения и за работу над ошибками, а также Марию Тихонову, являющуюся моим первым Учителем по машинному обучению.