Прочитав несколько известных статей по сегментации спутниковых снимков земли, я решил попробовать создать и обучить свою модель нейросети для этой задачи. И конечно, в процессе возникало много вопросов, своими ответами на которые я решил поделиться в рамках этого туториала. Поделиться так подробно и просто, как это было бы понятно таким новичкам, как я.
Идея
Одним из главных источников идеи и подхода для создания и обучения моей нейросети стала статья LULC Segmentation of RGB Satellite Image Using FCN-8 [перевод на русский]: берем датасет спутниковых изображений земли Gaofen Image Dataset → разделяем изображения на подизображения, вместо уменьшения и аугментации исходных изображений: так она будет лучше обучаться → определяем модель нейросети с архитектурой Fully Convolutional Network (FCN, FCN-8) → обучаем ее. Единственное, вместо модели FCN-8 я решил выбрать доступную в PyTorch предобученную модель ResNet101, написанную по архитектуре FCN (FCN.ResNet101). Также вместо большой версии датасета (Large) я выбрал упрощенную версию (Fine) — она включает в себя 10 оригинальных и сегментированных по 15 классам изображений в формате tif и размером 7200x6800. Модель FCN.ResNet101 предполагает размер входных изображений 224x224, а так-как размер исходных изображений чуть больше, чем может вместиться целое кол-во подизображений размера 224x224, мы обрезаем исходные изображения до 7168x6720 (именно обрезаем, чтобы избежать потери детализации); таким образом, у нас получается 960 подизображений с одного исходного изображения.
Приступим к реализации этой идеи!
Импорты
Импортируем все необходимые библиотеки:
import torch import torchvision import numpy as np import matplotlib.pyplot as plt from PIL import Image from sklearn.model_selection import train_test_split from os import listdir from typing import Tuple, List, Dict, Generator, Any from IPython.display import clear_output
Датасет
Сайт датасета находится по этой ссылке, а две версии датасета (Fine и Large) на OneDrive. Но здесь я буду загружать его со своего Google Drive.
Загрузка
Если вы пользуетесь не Google Colab, вам может понадобиться установить модуль gdown:
pip3 install gdown
Загрузим архивы датасета:
%%capture # clear output !gdown "https://drive.google.com/uc?id=1--fNMFRXmBRDQPwdFqh54Sx7_RNE6z4a&confirm=t" # download image_RGB.zip !gdown "https://drive.google.com/uc?id=1-3y-bJK-QZapo3nwZEXdVLWDjkVO4AYN&confirm=t" # download label_15classes.zip
И разархивируем их:
%%capture !unzip image_RGB.zip !unzip label_15classes.zip
Организация
Чтобы не нагружать оперативную память всеми подизображениями, можно сохранить каждое как файл и организовать их подгрузку. Сделаем это, но перед этим определимся с классами.
Упрощенная версия датасета, как уже было сказано, содержит сегментированные по 15 классам изображения. Их описание можно посмотреть в readme.txt на OneDrive:
label information of 15 classes: industrial land RGB: 200, 0, 0 urban residential RGB: 250, 0, 150 rural residential RGB: 200, 150, 150 traffic land RGB: 250, 150, 150 paddy field RGB: 0, 200, 0 irrigated land RGB: 150, 250, 0 ...
И так-как в вышеупомянутой статье модели создавали и обучали для сегментации одного целевого класса, выберем один из них, например, Irrigated Land: он достаточно хорошо представлен на изображениях датасета, а потому с ним будет достаточно просто. Но также важно добавить еще один класс — Background, так-как он важен для более корректного обучения модели.
Инициализируем словари с классами — classes и classes_by_id (последний пригодится для одной из функций далее):
classes = { (0, 0, 0): (0, '__background__'), (150, 250, 0): (1, 'irrigated_land') } classes_by_id = dict() for rgb, (id, name) in classes.items(): classes_by_id[id] = (rgb, name)
Теперь подизображения. Напишем функцию, которая будет возвращать генератор подизображений из исходного изображения:
def get_subimages_generator( image: Image.Image, subimage_size: Tuple[int, int, int] ) -> Generator[Image.Image, None, None]: for r in range(image.size[1] // subimage_size[1]): for c in range(image.size[0] // subimage_size[0]): yield image.crop(box=( c * subimage_size[0], r * subimage_size[1], (c + 1) * subimage_size[0], (r + 1) * subimage_size[1] ) )
Создадим директорию dataset с директориями для каждого типа подизображений изображений датасета (оригинальные и сегментированные):
!mkdir dataset dataset/originals dataset/labeleds
Наконец, напишем функцию, которая будет сохранять все подизображения в соответствующие директории под строковыми id (они понадобятся нам далее). Функция будет проходится по каждым сегментированным подизображениям, фильтровать их по минимальным процентам представленности определенного цвета и, собственно, сохранять:
def save_dataset_subimages(classes_filter: Dict[Tuple[int, int, int], float]): for i, filename in enumerate(listdir('image_RGB/')): basename = filename[:filename.find('.tif')] image = Image.open(fp=f'image_RGB/{basename}.tif').crop(box=(16, 40, 7200 - 16, 6800 - 40)) image_labeled = Image.open(fp=f'label_15classes/{basename}_label.tif').crop(box=(16, 40, 7200 - 16, 6800 - 40)) subimages = get_subimages_generator(image=image, subimage_size=(224,224)) subimages_labeleds = get_subimages_generator(image=image_labeled, subimage_size=(224,224)) for si, subimage in enumerate(subimages): subimage_labeled = next(subimages_labeleds) # classes filter do_continue = False subimage_labeled_colors_dict = {rgb: count for count, rgb in subimage_labeled.getcolors()} for rgb, min_percent in classes_filter.items(): if rgb not in subimage_labeled_colors_dict \ or subimage_labeled_colors_dict[rgb] * 100 / 50176 < min_percent: # 50176 = subimage width * subimage height do_continue = True break if do_continue: continue subimage.save(fp=f'dataset/originals/i{i}si{si}.tif') subimage_labeled.save(fp=f'dataset/labeleds/i{i}si{si}_labeled.tif')
Задействуем ее:
save_dataset_subimages( classes_filter={ (150, 250, 0): 5 } )
Команда ниже поможет нам узнать, сколько всего оригинальных, а значит и сегментированных, подизображений было сохранено:
!ls -lR dataset/originals/*.tif | wc -l
Всего для целевого класса Irrigated Land сохранено 4166 оригинальных, а значит и сегментированных, подизображений.
Вспомогательные функции и класс датасета
Для подгрузки и использования подизображений необходимо определить несколько функций и класс датасета, которые будут их обрабатывать и представлять так, как это будет удобно для работы с моделью. А для этого необходимо представлять оригинальные и сегментированные подизображения в виде тензоров (типе torch.Tensor), причем для большинства моделей, в том числе и для FCN.ResNet101, размерность тензоров должна предполагать наличие партий (batches): (batches, width, height, channels), где width, height и channels — ширина, высота и кол-во каналов изображения соответственно, в нашем случае — 224, 224 и 3 (R, G, B) соответственно. В случае с сегментированными изображениями, оно должно быть представлено еще и в размерности (batches, classes, width, height), где classes — количество сегментируемых классов, в нашем случае — 2; то есть должно представлять из себя маску, где для каждого класса указывается отдельная маска 224x224, которая указывает 0 или 1, на каком пикселе присутствует определенный класс, а на каком нет. Собственно, эти значения (0 или 1) для каждого пикселя каждой маски модель для семантической сегментации и должна предсказывать, оперируя вероятностью от 0 до 1.
Наконец, приступим к коду. Определим функции чтобы…
собирать два соответствующих подизображения (оригинальное и сегментированное) по их строковым id, которые были определены при их сохранении:
def get_dataset_subimage(dataset_subimage_id: str) -> Tuple[Image.Image, Image.Image]: subimage = Image.open(fp=f'dataset/originals/{dataset_subimage_id}.tif') subimage_labeled = Image.open(fp=f'dataset/labeleds/{dataset_subimage_id}_labeled.tif') return subimage, subimage_labeled
получать маску из сегментированного изображения и наоборот:
В первой функции проходимся по каждому пикселю сегментированного изображения и присваиваем маске для определенного класса определенного пикселя маски значение 1.0, при этом, если для пикселя не находится класса, относим его к классу Background. Получаем массив размерностью (classes, width, height):
def get_image_mask_from_labeled( image_labeled: Image.Image, classes: Dict[Tuple[int, int, int], Tuple[int, str]] ) -> np.ndarray: image_mask = np.zeros(shape=(len(classes),image_labeled.size[0],image_labeled.size[1])) image_labeled_ndarray = np.array(object=image_labeled) for r in np.arange(stop=image_labeled_ndarray.shape[0]): for c in np.arange(stop=image_labeled_ndarray.shape[1]): class_rgb = tuple(image_labeled_ndarray[r][c]) class_value = classes.get(class_rgb) if class_value != None: image_mask[class_value[0]][r][c] = 1.0 else: image_mask[0][r][c] = 1.0 return image_mask
Во второй функции получаем массив из индексов классов с максимальными вероятностями среди других классов для каждого пикселя, а затем присваиваем каждому пикселю выходного изображения цвет в соответствии с его классом. То есть, обратно получаем массив размерностью (width, height, channels):
def get_image_labeled_from_mask( image_mask: np.ndarray, classes_by_id: Dict[Tuple[int, int, int], Tuple[int, str]] ) -> Image.Image: image_labeled_ndarray = np.zeros( shape=(image_mask.shape[1],image_mask.shape[2],3), dtype=np.uint8 ) image_mask_hot = image_mask.argmax(axis=0) for r in np.arange(stop=image_mask_hot.shape[0]): for c in np.arange(stop=image_mask_hot.shape[1]): class_id = image_mask_hot[r][c] class_by_id_value = classes_by_id.get(class_id) image_labeled_ndarray[r][c] = np.array(object=class_by_id_value[0]) image_labeled = Image.fromarray(obj=image_labeled_ndarray) return image_labeled
чтобы предобрабатывать изображения (переводить в тензор):
def image_preprocess(image: Image.Image) -> torch.Tensor: return torchvision.transforms.ToTensor()(pic=image)
переводить два соответствующих подизображения в тензоры:
def get_dataset_subimage_tensor( subimage: Image.Image, subimage_labeled: Image.Image, classes: Dict[Tuple[int, int, int], Tuple[int, str]], dtype: torch.FloatType = None, ) -> Tuple[torch.Tensor, torch.Tensor]: subimage_tensor = image_preprocess(image=subimage) subimage_mask_tensor = torch.tensor( data=get_image_mask_from_labeled( image_labeled=subimage_labeled, classes=classes ), dtype=dtype ) return subimage_tensor, subimage_mask_tensor
Проверим, как они работают. Загружаем два соответствующих подизображения по их строковому id:
subimage, subimage_labeled = get_dataset_subimage(dataset_subimage_id='i0si0')
Выводим их с помощью модуля matplotlib.pyplot:
fig, ax = plt.subplots(ncols=2) ax[0].imshow(subimage) ax[1].imshow(subimage_labeled) plt.show()
Проверим корректность обработки масок. Переводим сегментированное подизображение в маску:
subimage_mask = get_image_mask_from_labeled( image_labeled=subimage_labeled, classes=classes )
Переводим маску обратно в сегментированное подизображение:
subimage_labeled_from_mask = get_image_labeled_from_mask( image_mask=subimage_mask, classes_by_id=classes_by_id )
Сравниваем оригинальное подизображение с сегментированным из маски:
fig, ax = plt.subplots(ncols=2) ax[0].imshow(subimage) ax[1].imshow(subimage_labeled_from_mask) plt.show()
Работает отлично!
Осталось только определить специальный класс датасета (с помощью абстрактного класса torch.utils.data.Dataset), который вместе с даталоадером (типом torch.utils.data.DataLoader) помогут удобно подгружать подизображения. И чтобы все они не занимали много памяти, в самом классе датасета мы определим их строковые id, а по запросу на получение элемента (вызову перегруженного метода __getitem__) будем уже загружать их, переводить в тензоры и возвращать. При этом важно создать несколько датасетов — обучающий и тестовый.
Но перед этим, нам нужно собрать все id подизображений:
def get_dataset_subimages_id() -> List[str]: return [ filename[:filename.find('.tif')] for filename in listdir(path='dataset/originals/') ]
Теперь разобъем список id на обучающий и тестовый. Для этого используем функцию train_test_split из модуля sklearn, которая поможет разбить их на обучающую и тестовую выборки, где длина обучающей выборки будет составлять 90%, а тестовой — 10%:
train_dataset_subimages_id, test_dataset_subimages_id = train_test_split(get_dataset_subimages_id(), train_size=0.9)
Наконец, определим класс датасета. Дополнительно для даталоадера нам понадобится перегрузить метод __len__ (для получение длины датасета) — в качестве длины датасета в нашем случае будет выступать длина списка id:
class Dataset(torch.utils.data.Dataset): def __init__(self, dataset_subimages_id: List[str], classes: Dict[Tuple[int, int, int], Tuple[int, str]], ): self.dataset_subimages_id = dataset_subimages_id self.classes = classes def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]: subimage = Image.open(fp=f'dataset/originals/{self.dataset_subimages_id[idx]}.tif') subimage_labeled = Image.open(fp=f'dataset/labeleds/{self.dataset_subimages_id[idx]}_labeled.tif') subimage_tensor, subimage_mask_tensor = get_dataset_subimage_tensor( subimage=subimage, subimage_labeled=subimage_labeled, classes=self.classes ) return subimage_tensor, subimage_mask_tensor def __len__(self) -> int: return len(self.dataset_subimages_id)
Используем вышеописанный класс для создания датасетов:
train_dataset = Dataset( dataset_subimages_id=train_dataset_subimages_id, classes=classes, ) test_dataset = Dataset( dataset_subimages_id=test_dataset_subimages_id, classes=classes )
Даталоадеры удобны тем, что с помощью них можно удобно разбивать датасет на партии и перемешивать. Оба этих параметра оказываются достаточно важными при обучении. Для нашей модели наиболее подходящим для обучающей выборки, по результатам моих экспериментов, оказывается batch_size=32 и shuffle=True.
train_dataloader = torch.utils.data.DataLoader( dataset=train_dataset, batch_size=32, shuffle=True ) test_dataloader = torch.utils.data.DataLoader( dataset=test_dataset, batch_size=1 )
Модель
Можно было бы написать свой класс модели с помощью абстрактного класса torch.nn.Module, однако, как уже говорилось в начале, я решил использовать готовую модель с предобученными на датасете ImageNet1K весами из модуля torchvision — FCN.ResNet101. Причем, c предобученными весами не всей модели, а только слоев backbone’а (так называемых извлекателей признаков), которые изображены на рисунке архитектуры FCN ниже между изображением котика с собачкой и слоя pixelwise prediction.
Но для начала, нужно инициализировать переменную device, которая поможет нам автоматически определять, на какой оперативной памяти (CPU или GPU) она будет находится и на какой оперативной памяти будут проводится вычисления нейросети. При обучении нейросети, особенно с большой архитектурой, важно пользоваться преимуществами распараллеливания вычислений на GPU, чтобы модель быстрее обрабатывала входные данные:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
Наконец, инициализируем модель с помощью модуля torchvision и определим кол-во выходных классов и веса backbone’а:
%%capture model = torchvision.models.segmentation.fcn_resnet101( num_classes=len(classes), weights_backbone=torchvision.models.ResNet101_Weights.IMAGENET1K_V1 ).to(device=device)
Напишем функцию для предсказывания по входному изображению:
def predict( image: Image.Image, model: torch.nn.Module, device: torch.DeviceObjType, ) -> Image.Image: image_tensor = image_preprocess(image=image) with torch.no_grad(): output_image_mask = model(image_tensor.unsqueeze(0).to(device))['out'][0].cpu().numpy() predicted_image_labeled = get_image_labeled_from_mask( image_mask=output_image_mask, classes_by_id=classes_by_id ) return predicted_image_labeled
И проверим ее:
predicted_image_labeled = predict( image=subimage, model=model, device=device )
fig, ax = plt.subplots(ncols=3) ax[0].imshow(subimage) ax[1].imshow(predicted_image_labeled) ax[2].imshow(subimage_labeled) plt.show()
Хорошо, она на что-то способна.
Обучение и тестирование модели
Метрики
При обучении и тестировании модели важно пользоваться различными метриками, чтобы следить за ее развитием. И если метрику потерь мы получаем из функции потерь (ее мы определим после), то в модулях torch и torchvision таких важных для семантической сегментации метрик как Pixel Accuracy и Intersection over Union (IoU) не содержится. Хотя они есть в torchmetrics, но я решил написать свои версии Pixel Accuracy и IoU, чтобы понять, что они из себя представляют. Хороший и простой обзор по этим и другим метрикам для задач сегментации провел Джереми Джордан на своем сайте.
Pixel Accuracy
Альтернативной метрикой для оценки семантической сегментации является простой отчет о проценте пикселей на изображении, которые были правильно классифицированы. Точность пикселей обычно указывается для каждого класса отдельно, а также глобально для всех классов. (Джереми Джордан)
Мы можем вычислить ее с помощью формулы:
Напишем функцию для вычисления метрики Pixel Accuracy. Причем, определим классификацированные пиксели так, как мы бы их отрисовывали, то есть с помощью получения масива из индексов классов с максимальными вероятностями среди других классов для каждого пикселя:
def metric_pixel_accuracy( y_pred: torch.Tensor, y_true: torch.Tensor ) -> float: y_pred_argmax = y_pred.argmax(dim=1) y_true_argmax = y_true.argmax(dim=1) correct_pixels = (y_pred_argmax == y_true_argmax).count_nonzero() uncorrect_pixels = (y_pred_argmax != y_true_argmax).count_nonzero() result = (correct_pixels / (correct_pixels + uncorrect_pixels)).item() return result
IoU
Метрика Intersection over Union (IoU), также называемая индексом Жаккара, по сути, является методом количественной оценки процентного перекрытия между целевой маской и нашим прогнозируемым результатом. Эта метрика тесно связана с коэффициентом Дайса, который часто используется в качестве функции потерь во время обучения. (Джереми Джордан)
Ее мы можем вычислить с помощью формулы:
Наишем функцию для вычисления метрики IoU. Определим классифицированными только те пиксели, вероятность которых выше 0.51 (51%):
def metric_iou( y_pred: torch.Tensor, y_true: torch.Tensor ) -> float: y_pred_hot = y_pred >= 0.51 intersection = torch.logical_and(y_pred_hot, y_true).count_nonzero() union = torch.logical_or(y_pred_hot, y_true).count_nonzero() result = (intersection / union).item() return result
Обучение
Определим функцию для обучения, которая будет подгружать подизображения, подавать модели входные тензоры и получать выходные, вычислять на их основе потерю, оптимизировать обучение, вычислять и выводить историю метрик по каждой партии, и все это определенное кол‑во эпох. Причем, чтобы оперативная память не перегружалась мы будем удалять уже не нужные громоздкие переменные и очищать оперативную память графического процессора, если мы его используем:
def train( model: torch.nn.Module, device: torch.DeviceObjType, train_dataloader: torch.utils.data.DataLoader, loss_fn: Any, optim_fn: Any, epochs: int ) -> Dict[str, List[float]]: history_metrics = { 'loss': list(), 'pixel_accuracy': list(), 'iou': list() } for e in range(1, epochs + 1): for b, data in enumerate(train_dataloader, start=1): subimage_tensor, subimage_mask_tensor = data if device.type == 'cuda': subimage_tensor = subimage_tensor.to(device) subimage_mask_tensor = subimage_mask_tensor.to(device) optim_fn.zero_grad() output = model(subimage_tensor) loss = loss_fn(output['out'], subimage_mask_tensor) loss.backward() optim_fn.step() loss_item = loss.item() pixel_accuracy = metric_pixel_accuracy(output['out'], subimage_mask_tensor) iou = metric_iou(output['out'], subimage_mask_tensor) history_metrics['loss'].append(loss_item) history_metrics['pixel_accuracy'].append(pixel_accuracy) history_metrics['iou'].append(iou) # dynamic output clear_output() print( 'Epoch: {}. Batch: {}. Loss: {:.3f} | Pixel Accuracy: {:.3f} | IoU: {:.3f}'.format( e, b, loss, pixel_accuracy, iou ) ) # memory clear del subimage_tensor, subimage_mask_tensor, output, loss if device.type == 'cuda': torch.cuda.empty_cache() return history_metrics
Определим также функции потерь и оптимизации, а также кол‑во эпох. Для нашей модели и для большинства моделей нейросетей они являются необходимыми, ведь они помогают проводить и оптимизировать обучение. В качестве функции потерь я выбрал популярную для задач семантической сегментации CrossEntropyLoss (torch.nn.CrossEntropyLoss). В качестве функции оптимизации — AdamW (torch.optim.AdamW), и важно определиться с ее параметрами: по результатам моих экспериментов, наиболее подходящим начальным коэффициентом обучения (learning rate, lr) оказывается 0.0001 (1e-4). Наиболее подходящее кол‑во эпох, опять‑же, по результатам моих экспериментов, оказывается 8, ведь по моим наблюдениям дальше он перестает обучаться, хотя в упомянутом в начале исследовании использовали целых 100 эпох, но для своей версии датасета.
Все готово к обучению. Запустим его:
history_metrics = train( model=model, device=device, train_dataloader=train_dataloader, loss_fn=torch.nn.CrossEntropyLoss(), optim_fn=torch.optim.AdamW(params=model.parameters(), lr=1e-4), epochs=8 )
Этот процесс оказывается достаточно долгим даже при использовании GPU в Google Colab. Возможно, часть кода оказывается недостаточно эффективным, поэтому будет очень хорошо, если вы сможете оптимизировать его для своих задач.
Исследование обучения
Выведем историю метрик с помощью matplotlib.pyplot:
plt.plot( history_metrics['loss'], 'red', history_metrics['pixel_accuracy'], 'green', history_metrics['iou'], 'blue', ) plt.title('History Metrics in Training\nep=8, bs=32, loss_fn=CrossEntopyLoss(), optim_fn=AdamW(lr=1e-4)') plt.xlabel('Batch') plt.ylabel('Value') plt.legend(('Loss', 'Pixel Accuracy', 'IoU')) plt.show()
Хотя развитие модели видно достаточно явно: потери уменьшаются, метрики Pixel Accuracy и IoU увеличиваются, а значит точность и предсказания увеличивается. Однако показатели представляются для меня слишком колеблющимися. Это может значит, что модель после каждой партии модель проводит слишком резкие изменения. Все это можно регулировать с помощью дополнительной обработки датасета, изменения размера партий (batch_size), выбора определенных функций потерь, параметров функции оптимизации (learning rate и других) и т. п. Все это входит в спектр моего дальнейшего изучения, поэтому, если вы в этом хорошо разбираетесь, я буду рад, если вы поделитесь со мной своими знаниями.
Сохранение модели
Сохраним модель с помощью функции save из модуля torch в формате h5:
torch.save( model.state_dict(), 'remezova_fcn_resnet101_ep8bs32lr1e-4.h5' )
Тестирование
Напишем похожую на train функцию test, только без элементов обучения и c выводом медианы метрик из истории метрик с каждой партией:
def test( model: torch.nn.Module, device: torch.DeviceObjType, test_dataloader: torch.utils.data.DataLoader, loss_fn: Any, ) -> Dict[str, List[float]]: history_metrics = { 'pixel_accuracy': list(), 'iou': list() } for b, data in enumerate(test_dataloader, start=1): subimage_tensor, subimage_mask_tensor = data if device.type == 'cuda': subimage_tensor = subimage_tensor.to(device) subimage_mask_tensor = subimage_mask_tensor.to(device) with torch.no_grad(): output = model(subimage_tensor) pixel_accuracy = metric_pixel_accuracy(output['out'], subimage_mask_tensor) iou = metric_iou(output['out'], subimage_mask_tensor) history_metrics['pixel_accuracy'].append(pixel_accuracy) history_metrics['iou'].append(iou) clear_output() print( 'Batch: {}. median Pixel Accuracy: {:.3f} | median IoU: {:.3f}'.format( b, np.median(a=history_metrics['pixel_accuracy']), np.median(a=history_metrics['iou']) ) ) # memory clear del subimage_tensor, subimage_mask_tensor, output if device.type == 'cuda': torch.cuda.empty_cache() return history_metrics
Проведем тестирование:
test_history_metrics = test( model=model, device=device, test_dataloader=test_dataloader, loss_fn=torch.nn.CrossEntropyLoss() )
И тем не менее получаем достаточно хороший результат:
Batch: 417. median Pixel Accuracy: 0.703 | median IoU: 0.523
Визуализируем результаты предсказаний. Выведем предсказания по пяти тестовым подизображениям:
fig, ax = plt.subplots(nrows=3, ncols=5, figsize=(16,8)) for si, (subimage_tensor, subimage_mask_tensor) in enumerate(test_dataloader): subimage = torchvision.transforms.ToPILImage()(pic=subimage_tensor[0]) subimage_labeled = get_image_labeled_from_mask( image_mask=subimage_mask_tensor[0].cpu().numpy(), classes_by_id=classes_by_id ) predicted_subimage_labeled = predict( image=subimage, model=model, device=device ) ax[0][si].imshow(subimage) ax[1][si].imshow(predicted_subimage_labeled) ax[2][si].imshow(subimage_labeled) if si == 5 - 1: break plt.show()
Заключение
Мы подготовили датасет, создали, обучили и протестировали модель для сегментации спутниковых снимков. Не смотря на некоторые недостатки, модель показывает достаточно хорошие результаты. Однако этот туториал и весь код будут продолжать изменяться или дополняться с развитием моих знаний или ваших рекомендаций.
Этот туториал вышел достаточно большим, но я надеюсь, я смог достаточно хорошо ответить на те вопросы по нейросетям для семантической сегментации, которые, возможно, возникли и у вас; особенно, если вы такой же новичок, как и я.
Также вы можете воспользоваться более структурированным кодом этой модели в репозитории Remezova на GitHub. ❤️
