PyTorch ~ Framework de segmentación semántica

Aquí hay 3 para compartir~~

El enlace al conjunto de datos VOC utilizado está abierto en el artículo. El modelo preentrenado se ha subido a Github. Yo uso el entorno Colab pro. Puede descargar el modelo para hacer predicciones.

Enlace de código:  https://github.com/lixiang007666/segmentation-learning-experiment-pytorch

Instrucciones:

1. Descargue el conjunto de datos de VOC y coloque JPEGImages SegmentationClasslas dos carpetas en la carpeta de datos.

2. Cambie la terminal al directorio de destino y ejecute python train.py -hpara ver el entrenamiento

 

Seleccione el modelo y el número de GPU para el entrenamiento, por ejemplo, ejecutarpython train.py -m Unet -g 0

1. La predicción requiere la modificación manual predict.pydel modelo en

Si conoce FCN muy bien, puede saltarse d2lla explicación (aprendizaje profundo práctico) hasta la última parte.

2 conjuntos de datos

El conjunto de datos VOC se utiliza generalmente para la detección de objetivos.En la versión 2012, se agregó la tarea de segmentación semántica.

El conjunto de datos básicos incluye: un conjunto de entrenamiento que contiene 1464 imágenes, un conjunto de validación de 1449 y un conjunto de prueba de 1456. Hay 21 categorías de objetos en total.

En la tarea de segmentación de PASCAL VOC, hay un total de 20 categorías de objetos, y otro contenido se utiliza como categoría de fondo, donde el rojo representa la categoría de la aeronave, el negro es el fondo y el límite de la aeronave se dibuja con beige (pareciendo blanco) líneas, que indican el área borrosa de segmentación.

Entre ellos, las etiquetas de segmentación son todas las imágenes en formato png, que en realidad es una imagen de índice de color de un solo canal. Además de una imagen de índice con un solo canal y del mismo tamaño que la imagen, la imagen también almacena una lista de 256 valores de color (paleta), cada valor de índice corresponde a un valor de color RGB en la paleta, por lo tanto, un mapa de índice de un solo canal + paleta puede representar un mapa de color.

Imagen original: inserte la descripción de la imagen aquí

Etiquetas: insertar descripción de la imagen aquí

Al seleccionar una imagen, se puede encontrar que hay más de dos categorías para una sola segmentación de imagen, y la categoría de cada imagen no es fija.

3 redes neuronales totalmente convolucionales

La segmentación semántica puede clasificar cada píxel de una imagen. La red totalmente convolucional (FCN) utiliza una red neuronal convolucional para realizar la transformación de píxeles de imagen a categorías de píxeles. A diferencia de las redes neuronales convolucionales que presentamos anteriormente en las secciones de clasificación de imágenes o detección de objetos, esto se logra a través de la capa de convolución transpuesta presentada全卷积网络将中间层特征图的高和宽变换回输入图像的尺寸 en . Por lo tanto, la predicción de clase de salida tiene una correspondencia uno a uno con la imagen de entrada a nivel de píxel: dada una posición en la dimensión espacial, la salida en la dimensión de canal es la predicción de clase para el píxel correspondiente a esa posición.

%matplotlib inline
import torch
import torchvision
from torch import nn
from torch.nn import functional as F
from d2l import torch as d2l

 3.1 Estructura de la red 

A continuación, usamos el modelo ResNet-18 entrenado previamente en el conjunto de datos de ImageNet para extraer las características de la imagen y denotamos esta instancia de red como pretrained_net. Las últimas capas del modelo incluyen agrupación promedio global y capas totalmente conectadas, mientras que no se requieren en redes totalmente convolucionales. Cree una instancia de red totalmente convolucional net. Reproduce la mayoría de las capas previas al entrenamiento en Resnet-18, pero elimina la capa de agrupación promedio global final y la capa totalmente conectada más cercana a la salida.  

 

num_classes = 21
net.add_module('final_conv', nn.Conv2d(512, num_classes, kernel_size=1))
net.add_module('transpose_conv', nn.ConvTranspose2d(num_classes, num_classes,
                                    kernel_size=64, padding=16, stride=32))

3.2 Inicializar la capa de convolución transpuesta

El sobremuestreo generalmente se usa para ampliar una imagen. La interpolación bilineal es uno de los métodos de muestreo superior comúnmente utilizados y también se usa a menudo para inicializar capas convolucionales transpuestas.

Para explicar la interpolación bilineal, suponga que, dada una imagen de entrada, queremos calcular cada píxel en una imagen de salida sobremuestreada. El muestreo ascendente para la interpolación bilineal se puede lograr mediante la transposición de capas convolucionales, y el kernel se construye mediante la siguiente función bilinear_kernel. Debido a limitaciones de espacio, solo damos la implementación de la función bilinear_kernel sin discutir el principio del algoritmo. 

def bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size + 1) // 2
    if kernel_size % 2 == 1:
        center = factor - 1
    else:
        center = factor - 0.5
    og = (torch.arange(kernel_size).reshape(-1, 1),
          torch.arange(kernel_size).reshape(1, -1))
    filt = (1 - torch.abs(og[0] - center) / factor) * \
           (1 - torch.abs(og[1] - center) / factor)
    weight = torch.zeros((in_channels, out_channels,
                          kernel_size, kernel_size))
    weight[range(in_channels), range(out_channels), :, :] = filt
    return weight

Experimente con muestreo ascendente con interpolación bilineal que se implementa mediante capas convolucionales transpuestas. Construimos una capa de convolución transpuesta que duplica la altura y el ancho de la entrada e inicializamos su kernel de convolución con la función bilinear_kernel.

conv_trans = nn.ConvTranspose2d(3, 3, kernel_size=4, padding=1, stride=2,
                                bias=False)
conv_trans.weight.data.copy_(bilinear_kernel(3, 3, 4));

En redes totalmente convolucionales, inicializamos capas convolucionales transpuestas con muestreo ascendente con interpolación bilineal. Para capas convolucionales 1×1, usamos parámetros de inicialización de Xavier. ¿Qué software  es?  http://143ai.com 

W = bilinear_kernel(num_classes, num_classes, 64)
net.transpose_conv.weight.data.copy_(W);

3.3 Formación

La función de pérdida y el cálculo de precisión no difieren fundamentalmente de la clasificación de imágenes, ya que usamos los canales de la capa convolucional transpuesta para predecir la clase de píxeles, por lo que la dimensión del canal se especifica en el cálculo de pérdida. Además, el modelo calcula la precisión en función de si la clase predicha es correcta para cada píxel.

def loss(inputs, targets):
    return F.cross_entropy(inputs, targets, reduction='none').mean(1).mean(1)

num_epochs, lr, wd, devices = 5, 0.001, 1e-3, d2l.try_all_gpus()
trainer = torch.optim.SGD(net.parameters(), lr=lr, weight_decay=wd)
d2l.train_ch13(net, train_iter, test_iter, loss, trainer, num_epochs, devices)

4 Código fuente abierto y conjunto de datos

Dirección de descarga del conjunto de datos: http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar

Muestras de entrada: Muestras de salida:  Capacitación:

!python3 train.py -m Unet -g 0

Predicción: el código del modelo incluye la implementación de FCN, U-Net y Deeplab. Puede cambiar el entrenamiento y la predicción del modelo de manera más conveniente. Resultados de segmentación de DeeplabV3:  Resultados de segmentación de FCN: inserte la descripción de la imagen aquí

Resultados de la segmentación de U-Net:

5 resumen

Al comparar con la imagen estándar segmentada, se puede encontrar que la imagen segmentada de salida del modelo es casi consistente con la imagen estándar segmentada, y la imagen segmentada de salida del modelo también está bien integrada con la imagen original, lo que indica que el modelo tiene mejor precisión.

Además, desde la perspectiva del tamaño de la imagen de entrada, el modelo puede ingresar una imagen de cualquier tamaño y generar una imagen segmentada etiquetada del mismo tamaño. Dado que está segmentado para las imágenes del conjunto de datos PASCAL VOC, solo se admiten 20 categorías en el conjunto de datos PASCAL VOC (el fondo es la categoría 21), por lo que al segmentar, todo lo que no esté en las 20 categorías se marcará como fondo. .

Pero, en general, el modelo logra una alta tasa de precisión para la segmentación de imágenes del conjunto de datos PASCAL VOC.