Сеть указателей — это вариация сетей семантической сегментации, которая может сегментировать изображение практически на все: объекты (люди/автомобили...), вещи (небо, земля), части (ноги/колеса) и материальные фазы ( жидкости/твердые вещества). В отличие от семантической сети, она может сегментировать экземпляры с неизвестными классами и, в отличие от Mask RCNN, не ограничена ограничивающей рамкой.
Сеть указателей получает изображение, точку (пиксель) на изображении и предсказывает область сегмента, содержащего точку (рис. 1).
В этом руководстве объясняется, как обучить такую сеть в 60 строках кода.
Поскольку сеть Pointer предсказывает один сегмент за раз, если мы хотим сегментировать полное изображение, нам нужно выбрать множество случайных точек, запустить сеть в каждой точке и сшить предсказанные сегменты вместе. Это означает, что время работы медленное и пропорционально количеству сегментов. Я нашел этот подход эффективным для сегментации объектов на части и материалов на фазы. Если мы знаем конкретную область интереса (RoI), мы также можем добавить еще один ввод, который является маской области на изображении, которое мы хотим сегментировать (маска области интереса).
В этом уроке цель будет состоять в том, чтобы сегментировать изображение прозрачного сосуда, содержащего различные материалы, на материалы и пустые области (рис. 2).
В целом обучение такой сети очень похоже на обучение сетей семантической сегментации типа (DeepLab или Unet). Единственное изменение состоит в том, что сеть теперь имеет два новых входа: точку указателя и маску сосуда (рис. 3 слева). Маска сосуда не является обязательной. Он действует как своего рода область интереса (ROI), которая позволяет нам ограничить область, которую мы хотим сегментировать, но мы можем запустить сеть без нее, если захотим (рис. 3).
Семантические сети ожидают на входе RGB-изображение. Что мы собираемся сделать, так это представить как маску сосуда, так и точку указателя в виде изображений (рис. 3). Например, точка указателя будет просто изображением, в котором все пиксели равны нулю (черные), за исключением пикселей рядом с целевой точкой, которые будут равны 1 (белые). Маска объекта/области интереса будет такой же, но в этом случае только пиксели, принадлежащие сосуду, будут иметь значение 1. Мы собираемся сложить маску указателя и маску объекта на изображении, поэтому вместо 3 изображение каналов (R, G, B) у нас будет изображение пяти каналов (R, G, B, маска указателя, маска сосуда).
Выход семантической сети сегментации представляет собой карту сегментации изображения (рисунок 3) с классом на пиксель. В нашем случае у нас будет два класса:
0: Пиксель не является частью сегмента, содержащего точку ввода.
1: Пиксель является частью сегмента, содержащего точку ввода.
Другое чем это, обучение будет в значительной степени идентично обучению семантической сегментации.
Первый шаг — загрузить набор данных LabPics для обучения:
https://zenodo.org/record/3697452/files/LabPicsV1.zip?download=1
Он содержит изображения и карты сегментации сосудов и материалов на каждом изображении.
Вам также потребуется установить Pytorch для обучения и OpenCV для чтения изображений.
OpenCV можно установить с помощью:
pip install opencv-python
Для установки Pytorch см.:
Обратите внимание, что для обучения сети требуется доступ к мощному графическому процессору.
Полный код этого проекта можно скачать здесь:
Приступим к написанию обучающего кода.
Для начала импортируем пакеты:
import os import numpy as np import cv2 import torchvision.models.segmentation import torch import torchvision.transforms as tf
Далее мы определяем основные параметры обучения:
Learning_Rate = 1e-5
width = height = 800
batchSize = 4
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
width,height: размеры изображения, используемого для обучения. Все изображения в процессе обучения будут уменьшены до этого размера.
batchSize: количество изображений, которые будут использоваться для каждой итерации обучения.
batchSize*width*height будет пропорционален требованиям к памяти для обучения. В зависимости от вашего оборудования может потребоваться использовать меньший размер пакета или образа, чтобы избежать проблем с нехваткой памяти.
Обратите внимание, что, поскольку мы тренируемся только с изображением одного размера, сеть после обучения, вероятно, будет ограничена работой только с изображениями этого размера. В большинстве случаев вам нужно изменить размер каждой обучающей партии.
устройство: автоматически устанавливает устройство, на котором будет работать сеть (GPU или CPU). На практике обучение без сильного графического процессора происходит крайне медленно.
Далее мы создаем список всех изображений в наборе данных:
TrainFolder=”LabPicsV1.2/Simple/Train/” ListImages = os.listdir(os.path.join(TrainFolder, “Image”))
Были TrainFolder и являются папкой поезда набора данных LabPics.
Изображения хранятся в подпапке «image» папки TrainFolder.
Теперь давайте определим набор преобразований, которые будут выполняться на изображении с помощью модуля преобразования TorchVision:
transformImg=tf.Compose([tf.ToPILImage(),tf.Resize((height,width)),tf.ToTensor(),tf.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))]) transformAnn=tf.Compose([tf.ToPILImage(),tf.Resize((height,width),tf.InterpolationMode.NEAREST),tf.ToTensor()])
Это определяет набор преобразований, которые будут применяться к изображению и карте аннотаций. Это включает в себя преобразование в формат PIL, который является стандартным форматом для преобразования. Изменение размера и преобразование в формат PyTorch. Для изображения мы также нормализуем интенсивность пикселей изображения, вычитая среднее значение и разделив его на стандартное отклонение интенсивности пикселей. Среднее значение и отклонение были рассчитаны заранее для большого набора изображений.
Далее мы создаем функцию, которая позволит нам загружать случайное изображение и соответствующие маски на изображении. По сути, мы собираемся выбрать случайное изображение. Затем мы выбираем случайную точку на этом изображении и используем маску материала, содержащую точку, в качестве целевого сегмента.
Полная функция:
def ReadRandomImage():
idx=np.random.randint(0,len(ListImages))
img=cv2.imread(TrainFolder+“/Image/”+ListImages[idx]))[:,:,0:3]
insMap = cv2.imread(os.path.join(TrainFolder, “Instance/”, ListImages[idx].replace(“jpg”,”png”)))
insMap[insMap == 254] = 0
img = transformImg(img)
insMap = transformAnn(insMap)
vesMap = insMap[2, :, :]
matMap = insMap[0, :, :]
if vesMap.max()<=0: return ReadRandomImage()
while(True):
x = np.random.randint(vesMap.shape[1])
y = np.random.randint(vesMap.shape[0])
if vesMap[y,x]>0:
vesMask = (vesMap == vesMap[y,x]).type(torch.float32)
matMask= (vesMask*(matMap == matMap[y,x])).type(torch.float32)
pointerMask = torch.zeros_like(matMask)
pointerMask[y — 4:y + 4, x — 4:x + 4] = 1
img=torch.cat([img,vesMask.unsqueeze(0), pointerMask.unsqueeze(0)],0)
return img,matMask
В первой части функции мы выбираем случайный индекс из списка изображений и загружаем изображение, соответствующее этому индексу, и карту аннотаций изображения.
idx=np.random.randint(0,len(ListImages)) img=cv2.imread(TrainFolder+“/Image/”+ListImages[idx]))[:,:,0:3] insMap = cv2.imread(os.path.join(TrainFolder, “Instance/”, ListImages[idx].replace(“jpg”,”png”)))
Затем мы удаляем несегментированные области (помеченные как 254) на карте аннотаций.
insMap[insMap == 254] = 0
Затем следует применить преобразование, которое мы определили ранее, к изображению и карте аннотаций:
img = transformImg(img) insMap = transformAnn(insMap)
Карта аннотацийinsMapсодержит три канала. Первый канал (insMap[0, :, :]) содержит карту сегментации всех материалов в изображении (все пиксели, принадлежащие к одной и той же фазе материала, имеют одинаковое значение, рис. 4). Третий канал (insMap[2, :, :]) содержит карту сегментации всех сосудов на изображении (все пиксели, принадлежащие одному и тому же сосуду, имеют одинаковое значение).
Мы собираемся извлечь карту судна и карту материала:
vesMap = insMap[2, :, :] matMap = insMap[0, :, :]
Если на изображении нет сосудов, мы собираемся загрузить новое случайное изображение.
if vesMap.max()<=0: return ReadRandomImage()
Далее мы собираемся выбрать случайную точку (пиксель) на изображении:
x = np.random.randint(vesMap.shape[1]) y = np.random.randint(vesMap.shape[0])
если пиксель находится внутри сосуда
if vesMap[y,x]>0:
мы собираемся загрузить маски материала и сосуда, содержащие этот пиксель (иначе мы собираемся продолжить и выбрать другой пиксель).
vesMask = (vesMap == vesMap[y,x]).type(torch.float32) matMask= (vesMask*(matMap == matMap[y,x])).type(torch.float32)
Оператор vesMap == vesMap[y,x] представляет собой логическую операцию между пикселем и матрицей. Это означает, что нужно взять все пиксели в матрице vesMap, которые имеют то же значение, что и пиксель в позиции x,y (vesMap[y,x]), и дать им значение 1 (истина), а остальная часть пикселя дает значение 0 (ложь). Поскольку все пиксели с одинаковым значением принадлежат одному и тому же экземпляру, это в основном означает создание маски материала или сосуда, содержащего пиксели x,y.
Обратите внимание, что маска сосуда в этом случае необязательна, так как мы действительно хотим предсказать область материала.
Наконец, давайте создадим маску указателя. Сначала мы создадим пустое изображение (все нули) размером с исходное изображение:
pointerMask = torch.zeros_like(matMask)
и установите значение пикселей вокруг целевой точки (x, y) равным 1.
pointerMask[y — 4:y + 4, x — 4:x + 4] = 1
Теперь у нас есть маска указателя изображения и маска сосуда. Мы собираемся объединить их в 5-канальное изображение (r, g, b, маска указателя, маска сосуда).
img=torch.cat([img,vesMask.unsqueeze(0),pointerMask.unsqueeze(0)],0)
и мы закончили загрузку данных.
Для обучения нам нужно использовать пакет изображений. Это означает, что несколько изображений накладываются друг на друга в 4D-матрице. Мы создаем пакет, используя функцию:
def LoadBatch():
images = torch.zeros([batchSize,5,height,width])
ann = torch.zeros([batchSize, height, width])
for i in range(batchSize):
images[i],ann[i]=ReadRandomImage()
return images, ann
Первая часть создает пустую матрицу 4d, в которой будут храниться изображения с размерами: [размер пакета, каналы, высота, ширина], где каналы — это количество слоев изображения; это 5 для изображения (r,g,b, маска указателя, маска сосуда) и 1 для маски аннотации материала.
Следующая часть загружает набор изображений и аннотацию в пустую матрицу, используя функцию ReadRandomImage(), которую мы определили ранее.
for i in range(batchSize):
images[i],ann[i]=ReadRandomImage()
Теперь, когда мы можем загрузить наши данные, пришло время загрузить нейронную сеть.
Мы загружаем сеть Deeplab для семантической сегментации:
Net = torchvision.models.segmentation.deeplabv3_resnet50(pretrained=True)
torchvision.models содержит много полезных моделей для семантической сегментации, таких как UNET и FCN . Мы выбираем Deeplabv3, поскольку это одна из лучших сетей семантической сегментации. Установив pretrained=True, мы загружаем сеть с предварительно обученным весом в наборе данных COCO. Всегда лучше начинать с предварительно обученной модели при изучении новой проблемы, поскольку она позволяет сети использовать предыдущий опыт и быстрее сходиться.
Deeplab и другие сети семантической сегментации получают на вход трехканальное RGB-изображение. Сеть указателей должна иметь пять входных каналов (r, g, b, маска указателя, маска сосуда). Поэтому мы собираемся изменить входной сверточный слой с 3 на пять каналов:
Net.backbone.conv1=torch.nn.Conv2d(5, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
Обратите внимание, что если вы не используете маску сосуда в качестве входных данных, это будет четыре канала.
Мы также хотим сохранить исходные веса, используемые для изображения RGB, поскольку это сэкономит время для обучения. Поэтому мы копируем веса из исходного входного слоя в новый входной слой:
weights=Net.backbone.conv1.weight.data Net.backbone.conv1=torch.nn.Conv2d(5, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) Net.backbone.conv1.weight.data[:,:3,:,:]=weights
С этого момента все так же, как и в обучающей стандартной сети семантической сегментации.
Последний слой сети Deeplab представляет собой слой свертки с 256 входными слоями и 21 выходным слоем. 21 представляет количество выходных классов (вы можете просмотреть слои, написав print(net))
У нас есть только 2 класса на пиксель:
0: не является частью сегмента, содержащего точку и
1: часть сегмента, содержащего).
Мы хотим заменить его новым свёрточным слоем с двумя выходами:
Net.classifier[4] = torch.nn.Conv2d(256, 3, kernel_size=(1, 1), stride=(1, 1))
Затем мы загружаем сеть в наше устройство GPU или CPU:
Net=Net.to(device)
Наконец, мы загружаем оптимизатор:
optimizer=torch.optim.Adam(params=Net.parameters(),lr=Learning_Rate) # Create adam optimizer
Оптимизатор будет контролировать скорость градиента на этапе обратного распространения в обучении. Оптимизатор AdamW — это новейший доступный оптимизатор.
Наконец, мы можем запустить тренировочный цикл. Мы собираемся тренироваться на 100000 шагов (но вы можете в любой момент остановиться и использовать обученную сеть):
for itr in range(0,100000):
На каждом этапе мы загружаем изображение с помощью функции LoadBatch, которую мы определили ранее, и преобразуем ее в GPU:
images,ann=LoadBatch() images=torch.autograd.Variable(images,requires_grad=False).to(device ann = torch.autograd.Variable(ann, requires_grad=False).to(device)
torch.autograd.Variable: преобразование данных в переменные градиента, которые могут использоваться сетью. Мы устанавливаем Requires_grad=False, поскольку мы не хотим применять градиент к изображению, а только к слоям сети. to(device)копирует тензор на то же устройство (GPU/CPU), что и сеть.
Наконец, мы вводим изображение в сеть и получаем прогноз.
Pred=Net(images)[‘out’]
После того, как мы сделали прогноз, мы можем сравнить прогнозируемый сегмент с реальной (наземной) аннотацией и рассчитать потери:
criterion = torch.nn.CrossEntropyLoss() # Set loss function Loss=criterion(Pred,ann.long()) # Calculate cross entropy loss Loss.backward() # Backpropogate loss Optimizer.step() # Apply gradient descent change to weight
Сначала определим функцию потерь. Мы используем стандартную кросс-энтропийную потерю:
criterion = torch.nn.CrossEntropyLoss()
Мы используем эту функцию для расчета потерь с использованием прогноза и реальной аннотации:
Loss=criterion(Pred,ann.long())
Как только мы вычислим потери, мы можем применить обратное распространение и изменить чистые веса.
Loss.backward() # Backpropogate loss Optimizer.step() # Apply gradient descent change to weight
Это охватывает весь этап обучения, но нам также необходимо сохранить обученную модель. В противном случае она будет потеряна после остановки программы.
Сохранение модели занимает много времени, поэтому мы хотим делать это примерно раз в 1000 шагов:
if itr % 1000 == 0: print(“Saving Model” +str(itr) + “.torch”) torch.save(Net.state_dict(), str(itr) + “.torch”)
После запуска этого скрипта около 30000 шагов сеть должна дать достойные результаты.
Полный код можно найти здесь:
Всего 60 строк кода без учета пробелов и 50 строк без учета импорта :-)
После того, как сеть обучена, мы хотим применить к сегменту реальное изображение и посмотреть результат. Мы делаем это с помощью отдельного сценария вывода, который использует сеть поездов для сегментации области одного судна на изображении.
Этот скрипт будет работать, загружая изображение и маску судна.
Затем сеть выберет набор случайных точек в сосуде, предскажет область материалов, содержащих каждую точку (или пустые области), и объединит их в одну карту сегментации.
Полный код доступен здесь:
Код теста очень похож на код поезда. Сначала мы загружаем изображение и маску сосуда:
def ReadImageAndMask(imagePath,maskPath): imgOrigin2=cv2.imread(imagePath)[:,:,0:3] cv2.imread(maskPath,0) img = transformImg(imgOrigin2) vesMask = transformAnn(vesMask)
Единственным отличием от обучения является преобразование маски сосуда из значений 0–255 в значения 0–1:
vesMask=(vesMask>0).type(torch.float32)
Далее мы собираемся определить функцию, которая выбирает случайную точку в маске и генерирует маску указателя:
def CreatePointerMask(mask):
pointerMask = torch.zeros_like(mask)
xy=torch.where(mask)
xyInd=np.random.randint(0,len(xy[0]))
y=xy[0][xyInd]x=xy[1][xyInd]
pointerMask[y — 4:y + 4, x — 4:x + 4] = 1
Эта часть почти такая же, как и в обучении, за исключением того, что мы используем функцию torch.where, чтобы перечислить все точки xy внутри маски. И выберите случайную координату из этого списка: xy=torch.where(mask).
Следующий шаг — загрузка и модификация сети:
Net=torchvision.models.segmentation.deeplabv3_resnet50(pretrained=True) Net.backbone.conv1=torch.nn.Conv2d(5, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) Net.classifier[4] = torch.nn.Conv2d(256, 2, kernel_size=(1, 1), stride=(1, 1)) Net = Net.to(device) Net.load_state_dict(torch.load(modelPath)) Net.eval()
Единственное отличие от обучающего сценария состоит в том, что мы также загружаем сеть, которую мы обучили и сохранили ранее:
Net.load_state_dict(torch.load(modelPath))
Где modelPath — это путь к файлу, содержащему сохраненную модель.
Следующая часть — создание маски сегментации. Мы начнем с определения пустой карты сегментации по размеру изображения:
segMap=torch.zeros_like(vesMask)[0]
Каждый новый сегмент, который мы собираемся предсказать, будет добавлен к этой карте.
Мы также определяем карту несегментированной области изображения:
unsegmentedMask=vesMask[0]+0
Затем мы загружаем изображение и маску, используя созданную ранее функцию:
img, vesMask, imgOrigin = ReadImageAndMask(imagePath,maskPath)
и начинаем сегментацию:
for i in range(1,100):
pointerMask=CreatePointerMask(unsegmentedMask)
input=torch.cat([img, vesMask, pointerMask.unsqueeze(0)], 0)
input=torch.autograd.Variable(input, requires_grad=False).to(device).unsqueeze(0)
with torch.no_grad():
Prd = Net(input)[‘out’]
seg = (Prd[0][1]>Prd[0][0])
segMap[seg]=i
unsegmentedMask[seg]=0
if unsegmentedMask.sum() / vesMask.sum() < 0.05: break
Цикл сегментации прост.
Во-первых, мы выбираем случайную точку в несегментированной области изображения, используя функцию, которую мы определили ранее:
pointerMask=CreatePointerMask(unsegmentedMask)
Как и прежде, мы накладываем изображение, маску указателя и маску сосуда на 5 входных каналов:
input=torch.cat([img, vesMask, pointerMask.unsqueeze(0)], 0)
и запустите прогноз на этом входе:
with torch.no_grad():
Prd = Net(input)[‘out’]
no.grad означает, что мы запускаем сеть без сбора градиента, это намного быстрее, и в любом случае нам нужен был только градиент для обучения.
Чистый вывод — это логит-карта, в которой для каждого пикселя указано, является ли он частью сегмента или нет. Мы преобразуем его в бинарную маску прогнозируемого сегмента, выбирая все пиксели, которые с большей вероятностью будут частью сегмента:
seg = (Prd[0][1]>Prd[0][0])
Теперь, когда у нас есть спрогнозированный сегмент (seg), мы добавим его на карту сегментации, установив для всех пикселей на карте сегментации, принадлежащих этому сегменту, уникальное значение (i).
segMap[seg]=i
мы также удаляем область сегмента из несегментированной области изображения:
несегментированная маска[seg]=0.
Каждый цикл мы предсказываем новый сегмент, добавляем его на карту сегментации и удаляем из несегментированного региона.
Когда сегментировано 95% региона, останавливаемся:
if unsegmentedMask.sum() / vesMask.sum() < 0.05: break
Мы можем наложить выходную маску сегментации на изображение, задав каждому сегменту случайный цвет:
segMap2=segMap.cpu().detach().numpy() segMap2=cv2.resize(segMap2,(imgOrigin.shape[1],imgOrigin.shape[0]),cv2.INTER_NEAREST) segImage=imgOrigin.copy() for i in np.unique(segMap2): # mark segmented regions on image if i==0: continue segImage[:, :, 0][segMap2 == i] = np.random.randint(0, 255) segImage[:, :, 1][segMap2 == i] = np.random.randint(0, 255) segImage[:, :, 2][segMap2 == i] = np.random.randint(0, 255) outImg=np.hstack([segImage,imgOrigin])[:,:,::-1] plt.imshow(outImg) plt.show()
Сначала мы преобразуем карту сегментации в numpy:
segMap2=segMap.cpu().detach().numpy()
Затем мы назначаем каждому сегменту случайный цвет:
for i in np.unique(segMap2): # mark segmented regions on image if i==0: continue segImage[:, :, 0][segMap2 == i] = np.random.randint(0, 255) segImage[:, :, 1][segMap2 == i] = np.random.randint(0, 255) segImage[:, :, 2][segMap2 == i] = np.random.randint(0, 255)
и отображать:
plt.imshow(outImg) plt.show()
Запустив это с тестовым изображением и маской:
Должен дать следующий результат:
Полный код доступен по адресу:
Все изображения в этом руководстве были взяты из набора данных LabPics (лицензия MIT) и извлечены со страницы YouTube NileRed.
Связанных с работой:
https://pubs.acs.org/doi/10.1021/acscentsci.0c00460
https://arxiv.org/abs/1902.07810
https://arxiv.org/abs/1909.07829
