使用UNet进行图像分割(Pytorch搭建)

使用UNet进行图像分割(利用Pytorch搭建)文章目录使用UNet进行图像分割(利用Pytorch搭建)简述环境准备代码数据集模型训练结果简述这里介绍一下如何使用Pytorch搭建一个UNet的图像分割模型,并训练出效果,论文中的一些trick这里没有使用。只包含简单的几个模块,并且大部分代码都有注释。环境平台:Windowspython版本:3.7Pytorch版本:torch:1.3.0,torchvision:0.4.0准备在搭建模型之前,我们还需要做些准备工作,那就是搜集数据
摘要由CSDN通过智能技术生成

使用UNet进行图像分割(利用Pytorch搭建)

简述

这里介绍一下如何使用Pytorch搭建一个UNet的图像分割模型,并训练出效果,论文中的一些trick这里没有使用。

只包含简单的几个模块,并且大部分代码都有注释。

环境

平台:Windows

python版本:3.7

Pytorch版本:torch:1.3.0,torchvision:0.4.0

准备

在搭建模型之前,我们还需要做些准备工作,那就是搜集数据,这里我提供一份眼球毛细血管数据集和一份VOC2012数据集来进行训练,如果有特殊的需求还是要自己搜集数据。

VOC数据集是 PASCAL VOC 挑战赛这个比赛使用的数据,里面包含了 目标分类、目标检测、目标分割、姿态识别、行为分类 所需要的数据与标签,我这里只使用分割的部分就可以了。

VOC数据集下载地址

官方: https://pjreddie.com/projects/pascal-voc-dataset-mirror/

或者:https://pan.baidu.com/s/1yfUILB185VvlgQ8bXk536w 提取码:geir

图像样式

原始图片





标签





毛细血管数据集下载地址

https://pan.baidu.com/s/1C06ERcImDpXlTneTrVuXPg 提取码:vmq0

图片样式

原始图片



标签



毛细血管数据集与VOC的用法差不多,就只是读取方式有点区别。

代码

总共三个py文件,文件名分别为:dataset(数据集)、unet(网络模型)、train(训练模块)

数据集

下载好数据集之后,解压完毕,可以看到这些文件:



这里我们只使用 JPEGImages 和 SegmentationClass 下的图片来进行语义分割,总共2913张图片。

首先是数据集部分的代码,没有使用数据增强,VOC部分:

import os
import cv2
import torchvision

from torch.utils.data import Dataset
from torchvision.utils import save_image

# 简单的数据集,没有进行数据增强
class Datasets(Dataset):

    def __init__(self, path):
        self.path = path
        # 语义分割需要的图片的加载进来,做标签,总共2913张图片
        self.name = os.listdir(os.path.join(path, "SegmentationClass"))
        self.trans = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])

    def __len__(self):
        return len(self.name)

    # 简单的正方形转换,把图片和标签转为正方形
    # 图片会置于中央,两边会填充为黑色,不会失真
    def __trans__(self, img, size):
        # 图片的宽高
        h, w = img.shape[0:2]
        # 需要的尺寸
        _w = _h = size
        # 不改变图像的宽高比例
        scale = min(_h / h, _w / w)
        h = int(h * scale)
        w = int(w * scale)
        # 缩放图像
        img = cv2.resize(img, (w, h), interpolation=cv2.INTER_CUBIC)
        # 上下左右分别要扩展的像素数
        top = (_h - h) // 2
        left = (_w - w) // 2
        bottom = _h - h - top
        right = _w - w - left
        # 生成一个新的填充过的图像,这里用纯黑色进行填充(0,0,0)
        new_img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(0, 0, 0))
        return new_img

    def __getitem__(self, index):
        # 拿到的图片
        name = self.name[index]
        # 把标签名的格式改成jpg,与原始图片一致
        name2jpg = name[:-3] + "jpg"

        # 所有的原始图片和标签
        img_path = [os.path.join(self.path, i) for i in ("JPEGImages", "SegmentationClass")]
        # 读取原始图片和标签,并转RGB
        img_o = cv2.imread(os.path.join(img_path[0], name2jpg))
        img_l = cv2.imread(os.path.join(img_path[1], name))
        img_o = cv2.cvtColor(img_o, cv2.COLOR_BGR2RGB)
        img_l = cv2.cvtColor(img_l, cv2.COLOR_BGR2RGB)

        # 转成网络需要的正方形
        img_o = self.__trans__(img_o, 
  • 83
    点赞
  • 594
    收藏
    觉得还不错? 一键收藏
  • 72
    评论
PyTorch是一个流行的深度学习框架,可以用于训练和部署图像分割模型。图像分割是将图像中的每个像素分类到特定的类别或分割区域的任务。在图像分割中,模型通过学习像素级别的特征来确定每个像素所属的类别或分割区域。 在PyTorch中,有许多用于图像分割的模型和库可供使用。例如,引用中提到的segmentation_models_pytorch库是一个优秀的图像分割库,它提供了许多常用的图像分割模型,如Unet、FPN、PSPNet等。这些模型在不同的数据集上进行了训练,并且可以根据实际需求进行微调或迁移学习。 具体来说,图像分割模型使用卷积神经网络(CNN)来提取图像中的特征,并使用这些特征来预测每个像素的类别或分割区域。常用的模型结构包括Unet、FCN、DeepLab等。这些模型通常由编码器和解码器组成。编码器负责提取图像的低级和高级特征,解码器则根据这些特征生成像素级别的预测结果。 训练图像分割模型通常需要大量的标注数据集,其中每个像素都有相应的标签。这些数据集可以包含多个类别,如引用中提到的21个类别,其中一个类别是背景类。模型通过将输入图像传递给网络,并使用与标签进行比较的损失函数来进行训练。训练过程中的优化器和学习率策略也会影响模型的性能。 在使用PyTorch进行图像分割时,通常需要进行数据预处理和增强操作,如引用中提到的albumentations库。这些操作可以提高模型的鲁棒性和泛化能力。 总结来说,图像分割是将图像中的每个像素分类到特定的类别或分割区域的任务。PyTorch是一款流行的深度学习框架,可以用于训练和部署图像分割模型。在PyTorch中,有许多优秀的图像分割库和模型可供使用,如segmentation_models_pytorch和albumentations。训练图像分割模型需要大量的标注数据集,并使用卷积神经网络来提取特征和预测每个像素的类别或分割区域。
评论 72
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值