建设网站后怎么发布,上海建筑设计研究院,网站服务器租用平台,企业门户网站 源码来源#xff1a;投稿 作者#xff1a;AI浩 编辑#xff1a;学姐 摘要
segmentation_models_pytorch是一款非常优秀的图像分割库#xff0c;albumentations是一款非常优秀的图像增强库#xff0c;这篇文章将这两款优秀结合起来实现多类别的图像分割算法。数据集选用CamVid… 来源投稿 作者AI浩 编辑学姐 摘要
segmentation_models_pytorch是一款非常优秀的图像分割库albumentations是一款非常优秀的图像增强库这篇文章将这两款优秀结合起来实现多类别的图像分割算法。数据集选用CamVid数据集类别有sky,building,pole,road,pavement,tree,signsymbol,fence,car,pedestrian,bicyclist,unlabelled等12个类别。数据量不大下载地址mirrors/alexgkendall/segnet-tutorial·GitCode。
通过这篇文章你可以学习到
1、如何在图像分割使用albumentations增强算法
2、如何使用dice_loss和cross_entropy_loss
3、如何segmentation_models_pytorch构架UNET模型
4、如何对分割数据做one-hot编码
项目结构
项目的结构如下 训练
新建train.py,插入一下代码
import os
import numpy as np
import cv2
import albumentations as albu
import torch
import segmentation_models_pytorch as smp
from torch.utils.data import DataLoader
from torch.utils.data import Dataset as BaseDataset导入需要的安装包接下来编写数据载入部分。
# ---------------------------------------------------------------
### 加载数据
# CamVid数据集中用于图像分割的所有标签类别
CLASSES [sky, building, pole, road, pavement,tree, signsymbol, fence, car,pedestrian, bicyclist, unlabelled]
class Dataset(BaseDataset):CamVid数据集。进行图像读取图像增强增强和图像预处理.Args:images_dir (str): 图像文件夹所在路径masks_dir (str): 图像分割的标签图像所在路径class_values (list): 用于图像分割的所有类别数augmentation (albumentations.Compose): 数据传输管道preprocessing (albumentations.Compose): 数据预处理def __init__(self,images_dir,masks_dir,augmentationNone,preprocessingNone,):self.ids os.listdir(images_dir)self.images_fps [os.path.join(images_dir, image_id) for image_id in self.ids]self.masks_fps [os.path.join(masks_dir, image_id) for image_id in self.ids]# convert str names to class values on masksself.class_values list(range(len(CLASSES)))self.augmentation augmentationself.preprocessing preprocessingdef __getitem__(self, i):# read dataimage cv2.imread(self.images_fps[i])image cv2.cvtColor(image, cv2.COLOR_BGR2RGB)mask cv2.imread(self.masks_fps[i], 0)# 从标签中提取特定的类别 (e.g. cars)masks [(mask v) for v in self.class_values]mask np.stack(masks, axis-1).astype(float)# 图像增强应用if self.augmentation:sample self.augmentation(imageimage, maskmask)image, mask sample[image], sample[mask]# 图像预处理应用if self.preprocessing:sample self.preprocessing(imageimage, maskmask)image, mask sample[image], sample[mask]print(mask.shape)return image, maskdef __len__(self):return len(self.ids)定义类别。类别的顺序对应mask的类别。 self.images_fps和self.masks_fps是图片的list和对应的mask图片的list。 self.class_values类别对应的indexindex的值对应mask上的类别值。 self.augmentation数据增强使用albumentations增强。self.preprocessing数据的预处理包含归一化和标准化预处理的方法来自smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)。 接下来解释__getitem__函数的内容: 读取图片。 将图片转为RGBcv2读取图片默认是BGR所以需要做转化。 接下来两行代码实现将mask转为one-hot编码。输入的shape是360480输出是36048012 图像增强。 图像预处理。 然后返回预处理后的图片和mask。 接下来是图片增强的代码
def get_training_augmentation():train_transform [albu.HorizontalFlip(p0.5),albu.ShiftScaleRotate(scale_limit0.5, rotate_limit0, shift_limit0.1, p1, border_mode0),albu.PadIfNeeded(min_height384, min_width480, always_applyTrue, border_mode0),albu.IAAAdditiveGaussianNoise(p0.2),albu.IAAPerspective(p0.5),albu.OneOf([albu.CLAHE(p1),albu.RandomBrightness(p1),albu.RandomGamma(p1),],p0.9,),albu.OneOf([albu.IAASharpen(p1),albu.Blur(blur_limit3, p1),albu.MotionBlur(blur_limit3, p1),],p0.9,),albu.OneOf([albu.RandomContrast(p1),albu.HueSaturationValue(p1),],p0.9,),]return albu.Compose(train_transform)def get_validation_augmentation():调整图像使得图片的分辨率长宽能被32整除test_transform [albu.PadIfNeeded(384, 480)]return albu.Compose(test_transform)def to_tensor(x, **kwargs):return x.transpose(2, 0, 1).astype(float32)def get_preprocessing(preprocessing_fn):进行图像预处理操作Args:preprocessing_fn (callbale): 数据规范化的函数(针对每种预训练的神经网络)Return:transform: albumentations.Compose_transform [albu.Lambda(imagepreprocessing_fn),albu.Lambda(imageto_tensor, maskto_tensor),]return albu.Compose(_transform)首先我们一起查看get_training_augmentation里面的代码。这里比较复杂。这些需要注意的是PadIfNeeded方法。 由于UNet系列的文章经历了5次缩放所以图片必须被32整除。所以通过填充的方式将图片的尺寸改为384480。 同样在验证集也要做这样的操作。 to_tensor函数是将图片的值转为tensor并将维度做交换。由于cv2读取的图片和mask的onehot的维度都是WHC需要高改为CWH。 get_preprocessing是对数据做预处理有归一化和标准化然后将图片和mask转为to_tensor。 接下来将最重要的训练部分
# $# 创建模型并训练
# ---------------------------------------------------------------
if __name__ __main__:ENCODER efficientnet-b1ENCODER_WEIGHTS imagenetACTIVATION softmax # could be None for logits or softmax2d for multiclass segmentationDEVICE cuda# 使用unet模型model smp.UnetPlusPlus(encoder_nameENCODER,encoder_weightsENCODER_WEIGHTS,classeslen(CLASSES),activationACTIVATION,)preprocessing_fn smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)这部分的代码主要是定义模型。 模型选用unet解码器是efficientnet-b1预训练权重为imagenet。 定义类别。 preprocessing_fn获取 smp.encoders的预处理方法。 # 数据集所在的目录DATA_DIR ./data/CamVid/# 如果目录下不存在CamVid数据集则克隆下载if not os.path.exists(DATA_DIR):print(Loading data...)os.system(git clone https://github.com/alexgkendall/SegNet-Tutorial ./data)print(Done!)# 训练集x_train_dir os.path.join(DATA_DIR, train)y_train_dir os.path.join(DATA_DIR, trainannot)# 验证集x_valid_dir os.path.join(DATA_DIR, val)y_valid_dir os.path.join(DATA_DIR, valannot)# 加载训练数据集train_dataset Dataset(x_train_dir,y_train_dir,augmentationget_training_augmentation(),preprocessingget_preprocessing(preprocessing_fn))# 加载验证数据集valid_dataset Dataset(x_valid_dir,y_valid_dir,augmentationget_validation_augmentation(),preprocessingget_preprocessing(preprocessing_fn))# 需根据显卡的性能进行设置batch_size为每次迭代中一次训练的图片数num_workers为训练时的工作进程数如果显卡不太行或者显存空间不够将batch_size调低并将num_workers调为0train_loader DataLoader(train_dataset, batch_size2, shuffleTrue, num_workers0)valid_loader DataLoader(valid_dataset, batch_size1, shuffleFalse, num_workers0)这部分的代码主要是数据集加载。 定义数据集所在路径。 获取训练集和验证集的路径。 加载训练集和验证集。 将训练集和测试集放入DataLoader中根据显卡的大小定义batch_size训练集需要shuffle验证集不需要。 然后定义loss loss smp.utils.losses.DiceLoss() smp.utils.losses.CrossEntropyLoss()metrics [smp.utils.metrics.IoU(threshold0.5),smp.utils.metrics.Recall()]optimizer torch.optim.Adam([dict(paramsmodel.parameters(), lr0.0001),])loss是DiceLoss和CrossEntropyLoss组合。 评分标准为IoU和Recall。 优化器选用Adam。 # 创建一个简单的循环用于迭代数据样本train_epoch smp.utils.train.TrainEpoch(model,lossloss,metricsmetrics,optimizeroptimizer,deviceDEVICE,verboseTrue,)valid_epoch smp.utils.train.ValidEpoch(model,lossloss,metricsmetrics,deviceDEVICE,verboseTrue,)# 进行40轮次迭代的模型训练max_score 0for i in range(0, 40):print(\nEpoch: {}.format(i))train_logs train_epoch.run(train_loader)valid_logs valid_epoch.run(valid_loader)# 每次迭代保存下训练最好的模型if max_score valid_logs[iou_score]:max_score valid_logs[iou_score]torch.save(model, ./best_model.pth)print(Model saved!)if i 25:optimizer.param_groups[0][lr] 1e-5print(Decrease decoder learning rate to 1e-5!)创建TrainEpoch和ValidEpoch循环用来迭代数据集。 按照迭代次数循环保存最好的模型。 完成上面的工作后就可以开始训练了。 测试
完成训练后就开始测试部分。
import osimport albumentations as albu
import cv2
import matplotlib.pyplot as plt
import numpy as np
import segmentation_models_pytorch as smp
import torch
from torch.utils.data import Dataset as BaseDatasetos.environ[CUDA_VISIBLE_DEVICES] 0导入所需要的包
# ---------------------------------------------------------------
### 加载数据
# CamVid数据集中用于图像分割的所有标签类别
CLASSES [sky, building, pole, road, pavement,tree, signsymbol, fence, car,pedestrian, bicyclist, unlabelled]
class Dataset(BaseDataset):CamVid数据集。进行图像读取图像增强增强和图像预处理.Args:images_dir (str): 图像文件夹所在路径masks_dir (str): 图像分割的标签图像所在路径class_values (list): 用于图像分割的所有类别数augmentation (albumentations.Compose): 数据传输管道preprocessing (albumentations.Compose): 数据预处理def __init__(self,images_dir,masks_dir,augmentationNone,preprocessingNone,):self.ids os.listdir(images_dir)self.images_fps [os.path.join(images_dir, image_id) for image_id in self.ids]self.masks_fps [os.path.join(masks_dir, image_id) for image_id in self.ids]# convert str names to class values on masksself.class_values list(range(len(CLASSES)))self.augmentation augmentationself.preprocessing preprocessingdef __getitem__(self, i):# read dataimage cv2.imread(self.images_fps[i])image cv2.cvtColor(image, cv2.COLOR_BGR2RGB)mask cv2.imread(self.masks_fps[i], 0)# 从标签中提取特定的类别 (e.g. cars)masks [(mask v) for v in self.class_values]mask np.stack(masks, axis-1).astype(float)# 图像增强应用if self.augmentation:sample self.augmentation(imageimage, maskmask)image, mask sample[image], sample[mask]# 图像预处理应用if self.preprocessing:sample self.preprocessing(imageimage, maskmask)image, mask sample[image], sample[mask]return image, maskdef __len__(self):return len(self.ids)# ---------------------------------------------------------------
### 图像增强def get_validation_augmentation():调整图像使得图片的分辨率长宽能被32整除test_transform [albu.PadIfNeeded(384, 480)]return albu.Compose(test_transform)def to_tensor(x, **kwargs):return x.transpose(2, 0, 1).astype(float32)def get_preprocessing(preprocessing_fn):进行图像预处理操作Args:preprocessing_fn (callbale): 数据规范化的函数(针对每种预训练的神经网络)Return:transform: albumentations.Compose_transform [albu.Lambda(imagepreprocessing_fn),albu.Lambda(imageto_tensor, maskto_tensor),]return albu.Compose(_transform)上面的代码是数据加载和数据增强和训练集的代码一样。
# 图像分割结果的可视化展示
def visualize(**images):PLot images in one row.n len(images)plt.figure(figsize(16, 5))for i, (name, image) in enumerate(images.items()):plt.subplot(1, n, i 1)plt.xticks([])plt.yticks([])plt.title( .join(name.split(_)).title())plt.imshow(image)plt.show()可视化测试结果展示原图真实的mask预测的mask。
# ---------------------------------------------------------------
if __name__ __main__:DATA_DIR ./data/CamVid/# 测试集x_test_dir os.path.join(DATA_DIR, test)y_test_dir os.path.join(DATA_DIR, testannot)ENCODER efficientnet-b1ENCODER_WEIGHTS imagenetACTIVATION softmax # could be None for logits or softmax2d for multiclass segmentationDEVICE cudapreprocessing_fn smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)# ---------------------------------------------------------------# $# 测试训练出来的最佳模型# 加载最佳模型best_model torch.load(./best_model.pth)# 创建测试数据集test_dataset Dataset(x_test_dir,y_test_dir,augmentationget_validation_augmentation(),preprocessingget_preprocessing(preprocessing_fn),)# ---------------------------------------------------------------# $# 图像分割结果可视化展示# 对没有进行图像处理转化的测试集进行图像可视化展示test_dataset_vis Dataset(x_test_dir, y_test_dir)# 从测试集中随机挑选3张图片进行测试for i in range(3):n np.random.choice(len(test_dataset))image_vis test_dataset_vis[n][0].astype(uint8)image, gt_mask test_dataset[n]gt_mask (np.argmax(gt_mask, axis0) * 255 / (gt_mask.shape[0])).astype(np.uint8)x_tensor torch.from_numpy(image).to(DEVICE).unsqueeze(0)pr_mask best_model.predict(x_tensor)pr_mask (pr_mask.squeeze().cpu().numpy())pr_mask (np.argmax(pr_mask, axis0) * 255 / (pr_mask.shape[0])).astype(np.uint8)# 恢复图片原来的分辨率gt_mask cv2.resize(gt_mask, (480, 360))pr_mask cv2.resize(pr_mask, (480, 360))visualize(imageimage_vis,ground_truth_maskgt_mask,predicted_maskpr_mask)获取测试集的路径。 定义ENCODER 为 efficientnet-b1ENCODER_WEIGHTS为imagenetACTIVATION为softmax。 获取预训练参数。 加载模型。 加载数据集。 加载没有做处理的图片。 随机选择3张图片 从test_dataset_vis获取图片。 从test_dataset获取对应的图片和mask。 将mask放大255的范围。 预测图片生成预测的mask。 将预测的mask也对应的放到255的范围。 然后重新resize到原来的尺寸。 可视化结果。 运行结果 CVPR图像分割必读论文
关注下方卡片《学姐带你玩AI》
回复“CVPR”免费领取
码字不易欢迎大家点赞评论收藏 引用 https://blog.csdn.net/hhhhhhhhhhwwwwwwwwww/article/details/124086830
