0 项目前言

从查阅 PaddleXPaddleSeg 文档,到后来的熟悉项目框架布局和各种接口;认为会用项目并不是目的,而能利用项目提供的各种组件助力自身实现既定任务才是重点,所以研究设计了本项目,在学习 Paddle 项目的过程中也清楚的对算法细节进行把握。

本项目数据来源为 常规赛:PALM眼底彩照视盘探测与分割 赛题数据,将利用 PaddleX 构建目标检测器,调用与自建 PaddleSeg API实现视盘定位与分割任务,详细实现代码均放在单元格中。

追求一个单元格完成一个步骤,代码无注释但流程清晰;如果本项目的思路和代码片段对你有帮助,不妨点个赞哦~~

项目亮点

  • 针对视杯视盘分割任务中常见的 ROI 区域裁剪步骤,提出了一种将语义分割数据集转为目标检测数据集的方法[3.2.1],然后训练一个用于检测视盘区域的 PP-YOLO Tiny 检测器[3.2.2]。
  • 设计了一个可扩展 ROI 区域裁剪与复原类,以满足自定义裁剪方法的轻松迁移[3.3]。
  • 通过 PaddleSeg API 自建符合 PaddleSeg 设计理念的分割模型进行图像分割[4.3]。

结束本项目之后,将会参考 PaddleSeg 的算法架构自己编写一个 视杯视盘分割 3 分类全流程,即利用飞桨基础 API 实现的任务。这对于从会用工具到会改工具大有裨益。


1 数据介绍

本常规赛提供的金标准由中山大学中山眼科中心的7名眼科医生手工进行视盘像素级标注,之后由另一位高级专家将它们融合为最终的标注结果。存储为 BMP 图像,与对应的眼底图像大小相同,标签为 0 代表视盘(黑色区域);标签为 255 代表其他(白色区域)。

  • 训练数据集

文件名称:Train:Train文件夹里有 fundus_images 文件夹和 Disc_Masks 文件夹。

  1. fundus_images文件夹内包含800张眼底彩照,分辨率为1444×1444,或2124×2056。命名形如H0001.jpg、N0001.jpg、P0001.jpg和V0001.jpg。
  2. Disc_Masks文件夹内包含fundus_images里眼底彩照的视盘分割金标准,大小与对应的眼底彩照一致。命名前缀和对应的fundus_images文件夹里的图像命名一致,后缀为bmp。
  • 测试数据集

文件名称:PALM-Testing400-Images:包含400张眼底彩照,命名形如T0001.jpg。

2 方案架构

Fu H, Cheng J, Xu Y, Wong DWK, Liu J, Cao X. Joint Optic Disc and Cup Segmentation Based on Multi-Label Deep Network and Polar Transformation. IEEE Trans Med Imaging. 2018;37(7):1597-1605. doi:10.1109/TMI.2018.2791488

本项目的方案思路即引文提供图片所示。

简要概括:首先,裁剪视盘 ROI 区域进行极坐标转换;然后利用多尺度输入和多损失深度监督的方法进行分割。

需要注意的是,模型输出设计会进行修改(Side-Ouput,Sigmoid),以适应 PaddleSeg 的 argmax 设计理念。

2.1 依赖模块

首先导入必要模块,清理数据目录结构。

!pip install paddlex
!pip install paddleseg
!pip install scikit-image
import warnings
warnings.filterwarnings('ignore')

import paddle
import paddlex
import paddleseg

import os
import shutil
import glob

import numpy as np
import pandas as pd

import cv2
import imghdr
from PIL import Image

import matplotlib.pyplot as plt
%matplotlib inline

2.2 数据清洗

!wget https://bj.bcebos.com/v1/dataset-bj/%E5%8C%BB%E7%96%97%E6%AF%94%E8%B5%9B/%E5%B8%B8%E8%A7%84%E8%B5%9B%EF%BC%9APALM%E7%9C%BC%E5%BA%95%E5%BD%A9%E7%85%A7%E8%A7%86%E7%9B%98%E6%8E%A2%E6%B5%8B%E4%B8%8E%E5%88%86%E5%89%B2.zip -O data/dataset.zip
!unzip -oq /home/aistudio/data/dataset.zip -d data

!rm -r data/__MACOSX
!mv data/常规赛:PALM眼底彩照视盘探测与分割 data/dataset
!mv data/dataset/Train/fundus_image data/dataset/Train/JPEGImages
!mv data/dataset/Train/Disc_Masks data/dataset/Train/Annotations_origin
!mv data/dataset/PALM-Testing400-Images data/dataset/Test

整理好的数据目录如下(其中的 data119773 是在小数据集上训练好的一个视盘检测器,因为整个数据训练时间比较久,而且小数据集上的评估指标也还行)。

!tree data -d
data
├── data119773
└── dataset
    ├── Test
    └── Train
        ├── Annotations_origin
        └── JPEGImages

6 directories

根据 PaddleX 的设计理念,类别编号是从 0 开始递增的,这里对 0-255 分别映射到 1-0,并进行格式的对齐(JPG-RGB,PNG-L)与清洗(尺寸统一),将新的标注格式PNG图像放入 data/dataset/Train/Annotations

!mkdir data/dataset/Train/Annotations
image_path_list = glob.glob('data/dataset/Train/JPEGImages/*.jpg')

for i in range(len(image_path_list)):
    img_path = image_path_list[i]
    jpg_name = str(img_path.split('/')[-1]).split('.')[0]
    gt_path = os.path.join('data/dataset/Train/Annotations_origin', f'{jpg_name}.bmp')
    assert imghdr.what(img_path) and imghdr.what(gt_path)

    img = Image.open(img_path).convert('RGB')
    gt = Image.open(gt_path).convert('L')
    if img.size != gt.size:
        img = img.resize(gt.size, Image.ANTIALIAS)
    img.save(img_path)

    cv2.imwrite(os.path.join('data/dataset/Train/Annotations', f'{jpg_name}.png'), np.array(gt, dtype='uint8'))

3 视盘探测

欲实现视盘 ROI 区域的裁剪首先需要定位。

原文利用 UNet 进行一个粗分割,对分割结果利用图形学后处理后提取质心坐标实现的。

而本项目里是通过目标检测的方法检测视盘中点坐标,然后扩展裁剪 512x512 的 ROI 区域。

3.1 图形学定位

以下是笔者使用图形学进行视盘定位的一点探索~,主要思路是通过滤波平均区域亮度,提取图像最亮的点坐标。

该方法在测试集上验证可以达到 98% 准确率(部分病变图像影响),但这并不是理想结果,我们希望最好是100%包含视盘区域。

from scipy.ndimage import grey_opening

def _get_optic_disc_location(image):
    _image = image.copy()

    KERNEL_SIZE = (15, 15)
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, ksize=KERNEL_SIZE)
    clahe_ = cv2.createCLAHE(3, (8, 8))
    B, G, R = cv2.split(_image)
    clahe = clahe_.apply(R)

    gauss = cv2.GaussianBlur(clahe, ksize=KERNEL_SIZE, sigmaX=0)
    gauss_open = grey_opening(gauss, structure=kernel)
    minVal, maxVal, minLoc, maxLoc = cv2.minMaxLoc(gauss_open)

    return maxLoc

找一张图片测试一下,KERNEL_SIZE 影响着程序运行时长。

image = cv2.imread('data/dataset/Test/T0002.jpg')

maxLoc = _get_optic_disc_location(image)
cv2.circle(image, center=maxLoc, radius=256, color=(101, 67, 254), thickness=15)

plt.figure(figsize=(6, 6))
plt.imshow(image[:, :, [2, 1, 0]])
plt.show()

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-wNNI4HhR-1639095699399)(output_18_0.png)]

下面是 PaddleX 目标检测方法定位视盘区域。

3.2 目标检测定位

这里笔者使用经典轻量级模型 PPYOLO Tiny 实现视盘定位(后面会提供一个其他数据集上 80+20 训练集的 bbox_map=99.567(320x320) 的 PP-YOLO Tiny 模型以供大家免去此步骤的运行)。

3.2.1 数据集制作

这里需要根据训练集的标注结果构建 PascalVOC 格式目标检测数据集。

数据清洗完成之后,就开始制作目标检测数据集。首先新建一个专门用于检测训练的文件目录 data/detection,把训练图片先拷贝过来。

!mkdir data/detection
!mkdir data/detection/Annotations
!mkdir data/detection/JPEGImages

!cp data/dataset/Train/JPEGImages/*.jpg data/detection/JPEGImages

下面将实现程序:生成目标检测标注文件放置于 data/detection/Annotations/*.xml

具体实现的算法也是利用图形学提取质心和边框,需要注意的是二值图中 0 被认为是背景,所以需要将视盘部分转换为 255;并且,数据集中部分图像是没有标注的(视盘不完整的是没有标注图的),就需要删除这部分数据。

from skimage import measure, morphology
label_content = '''<?xml version="1.0" encoding="utf-8"?>
<annotation>
    <folder>JPEGImages</folder>
    <filename>{}</filename>
    <size>
        <width>{}</width>
        <height>{}</height>
        <depth>3</depth>
    </size>
    <object>
        <name>ROI</name>
        <pose>Unspecified</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>{}</xmin>
            <ymin>{}</ymin>
            <xmax>{}</xmax>
            <ymax>{}</ymax>
        </bndbox>
    </object>
</annotation>
'''

def _seg_to_det(label, save_path: str, file_name: str, label_content: str):
    H, W = label.shape[0], label.shape[1]

    label = measure.label(label)
    label = morphology.remove_small_objects(label > 0, min_size=25)
    label = label.astype(np.uint8)

    regions = measure.regionprops(label)
    try:
        assert len(regions) == 1
    except Exception as e:
        return False

    for region in regions:
        h, w = region.centroid
        x_min, y_min, x_max, y_max = region.bbox[1], region.bbox[0], \
                                        region.bbox[3], region.bbox[2]

        with open(save_path, mode='w', encoding='utf-8') as f:
            f.write(label_content.format(
                file_name+'.jpg',
                W, H,
                x_min, y_min, x_max, y_max))

    return True
label_list = sorted(glob.glob('data/dataset/Train/Annotations/*.png'))

for i in range(len(label_list)):
    label = cv2.imread(label_list[i], cv2.IMREAD_GRAYSCALE)
    label = (label == 0) * 255

    image_name = str(str(label_list[i].split('/')[-1]).split('.')[0])
    save_path = os.path.join('data/detection/Annotations', image_name+'.xml')

    has_gt = _seg_to_det(
        label,
        save_path=save_path,
        file_name=image_name,
        label_content=label_content)
    if not has_gt:
        os.remove(os.path.join('data/detection/JPEGImages', image_name+'.jpg'))

还是有40张图是未标注的,删除完成后选择 75%(570) 训练, 25%(190) 验证。

!paddlex --split_dataset --format VOC\
    --dataset_dir data/detection\
    --val_value 0.25\
    --test_value 0

3.2.2 模型训练与验证

PaddleX 全流程大家也是耳熟能详,下面就把性质相同的代码都放一个单元吧。

以下是笔者常用的目标检测流程。首先构建数据增强方法,其中的 RandomCrop 及以上部分可以注释掉,并且使用线性插值进行缩放都是因为训练比较快(训练时长可以相差几倍…),当然,因为是裁剪 512x512 尺寸区域,所以交并比的精度要求不用太高。

from paddlex import transforms as T

train_transforms = T.Compose([
    T.MixupImage(alpha=1.5, beta=1.5, mixup_epoch=int(550 * 25. / 27)),
    T.RandomDistort(
        brightness_range=0.5, brightness_prob=0.5,
        contrast_range=0.5, contrast_prob=0.5,
        saturation_range=0.5, saturation_prob=0.5,
        hue_range=18.0, hue_prob=0.5),
    T.RandomExpand(prob=0.5, upper_ratio=2.0,
        im_padding_value=[float(int(x * 255)) for x in [0.485, 0.456, 0.406]]),
    T.RandomCrop(),
    T.Resize(target_size=320, interp='LINEAR'),
    T.RandomHorizontalFlip(prob=0.5),
    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
eval_transforms = T.Compose([
    T.Resize(target_size=320, interp='LINEAR'),
    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

train_dataset = paddlex.datasets.VOCDetection(
    data_dir='data/detection',
    file_list='data/detection/train_list.txt',
    label_list='data/detection/labels.txt',
    transforms=train_transforms,
    shuffle=True)
eval_dataset = paddlex.datasets.VOCDetection(
    data_dir='data/detection',
    file_list='data/detection/val_list.txt',
    label_list='data/detection/labels.txt',
    transforms=eval_transforms)

然后对边框进行聚类,构建 PP-TOLO Tiny 模型。

anchors = paddlex.tools.YOLOAnchorCluster(
    num_anchors=9,
    dataset=train_dataset,
    image_size=320)()

model = paddlex.det.PPYOLOTiny(
    num_classes=len(train_dataset.labels),
    backbone='MobileNetV3',
    anchors=anchors)

Warmup + Piecewise-decay 的训练策略。

learning_rate = 0.0015
warmup_steps = 95
warmup_start_lr = 0.0
train_batch_size = 32
step_each_epoch = train_dataset.num_samples // train_batch_size

lr_decay_epochs = [130, 540]
boundaries = [b * step_each_epoch for b in lr_decay_epochs]
values = [learning_rate * (0.1**i) for i in range(len(lr_decay_epochs) + 1)]

lr = paddle.optimizer.lr.PiecewiseDecay(
    boundaries=boundaries,
    values=values)

lr = paddle.optimizer.lr.LinearWarmup(
    learning_rate=lr,
    warmup_steps=warmup_steps,
    start_lr=warmup_start_lr,
    end_lr=learning_rate)

optimizer = paddle.optimizer.Momentum(
    learning_rate=lr,
    momentum=0.9,
    weight_decay=paddle.regularizer.L2Decay(0.0005),
    parameters=model.net.parameters())

开始训练 550 轮次。

model.train(
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,

    num_epochs=550,
    train_batch_size=train_batch_size,
    optimizer=optimizer,

    save_interval_epochs=30,
    log_interval_steps=step_each_epoch * 5,
    save_dir='output/PPYOLOTiny',
    pretrain_weights='IMAGENET',
    use_vdl=False)

3.3 区域裁剪与复原

训练好的轻量级视盘检测模型已经保存在 data/data119773/PPYOLOTiny.zip,解压之后,利用下面的代码块调用即可。

主要入口说明:

  • __init__(model_path, image_dir, label_dir):传入模型文件夹路径,原图文件夹路径,标签文件夹路径(可选),以及裁剪大小。
  • crop(write_file=True):执行裁剪,将 ROI 保存到同级目录下(例如原路径是 ../../images,则会在该路径下新建 ../../images_patch
  • paste_origin(patch_dir):将指定的裁剪下来的图像还原成原来的形状保存于 ../../images_patch_pasted
  • transform(image)/inv_transform(image):裁剪图像执行极坐标变换与逆变换,无需手动调用。

自定义定位模型:

  • get_bbox():里面可添加自己模型,构建返回 [x_min, y_min, x_max, y_max] 即可。
class DiscCrop:
    def __init__(self,
                 model_path: str,
                 image_dir: str,
                 label_dir=None,
                 crop_size=512):

        self.model = paddlex.load_model(model_path)
        self.model_type = self.model.get_model_info()['_Attributes']['model_type']

        self.crop_size = crop_size
        self.crop_logs = {}

        self.image_path_list = sorted(glob.glob(os.path.join(image_dir, '*.jpg')))
        self.image_save_dir = ''.join([image_dir, '_patch'])

        if label_dir is not None:
            self.label_path_list = sorted(glob.glob(os.path.join(label_dir, '*.png')))
            self.label_save_dir = ''.join([label_dir, '_patch'])

    @staticmethod
    def transform(image):
        center = (image.shape[1] // 2, image.shape[0] // 2)
        _image = cv2.linearPolar(
            src=image,
            center=center,
            maxRadius=center[0],
            flags=cv2.WARP_FILL_OUTLIERS)
        return cv2.rotate(_image, cv2.ROTATE_90_CLOCKWISE)

    @staticmethod
    def inv_transform(image):
        center = (image.shape[1] // 2, image.shape[0] // 2)
        _image = cv2.linearPolar(
            src=cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE),
            center=center,
            maxRadius=center[0],
            flags=cv2.WARP_FILL_OUTLIERS + cv2.WARP_INVERSE_MAP)
        return _image

    def get_bbox(self, image):
        result = self.model.predict(image)
        if self.model_type == 'detector':
            for item in result:
                x_min, y_min, w, h = item['bbox']

                x, y = x_min + w / 2., y_min + h / 2.
                x_min, x_max = max(0, x - self.crop_size // 2), min(image.shape[1], x + self.crop_size // 2)
                y_min, y_max = max(0, y - self.crop_size // 2), min(image.shape[0], y + self.crop_size // 2)

                return np.uint32([x_min, y_min, x_max, y_max])
        else:
            return NotImplemented

    def crop(self, write_file=True):
        if write_file:
            if not os.path.exists(self.image_save_dir):
                os.mkdir(self.image_save_dir)
            if hasattr(self, 'label_save_dir'):
                if not os.path.exists(self.label_save_dir):
                    os.mkdir(self.label_save_dir)

        for i in range(len(self.image_path_list)):
            image_path = self.image_path_list[i]
            image_name = image_path.split(os.sep)[-1]
            flag_name = image_name.split('.')[0]

            image = cv2.imread(image_path)
            x_min, y_min, x_max, y_max = self.get_bbox(image)
            self.crop_logs[flag_name] = [x_min, y_min, x_max, y_max, image.shape]
            if not write_file:
                continue

            crop_image = self.transform(image[y_min: y_max, x_min:x_max])
            image_save_path = os.path.join(self.image_save_dir, image_name)
            cv2.imwrite(image_save_path, crop_image)

            if hasattr(self, 'label_save_dir'):
                label_path = self.label_path_list[i]
                label = cv2.imread(label_path)
                crop_label = self.transform(label[y_min: y_max, x_min:x_max])
                crop_label = cv2.cvtColor(crop_label, cv2.COLOR_BGR2GRAY)
                label_save_path = os.path.join(self.label_save_dir, label_path.split(os.sep)[-1])
                cv2.imwrite(label_save_path, crop_label)

    def paste_origin(self, patch_dir):
        patch_path_list = sorted(glob.glob(os.path.join(patch_dir, '*')))
        pasted_save_dir = ''.join([patch_dir, '_pasted'])
        if not os.path.exists(pasted_save_dir):
            os.mkdir(pasted_save_dir)

        for i in range(len(patch_path_list)):
            patch_path = patch_path_list[i]
            patch_name = patch_path.split(os.sep)[-1]
            flag_name = patch_name.split('.')[0]
            x_min, y_min, x_max, y_max, origin_shape = self.crop_logs[flag_name]

            patch = cv2.imread(patch_path)
            patch = self.inv_transform(patch)

            origin_shape_file = np.zeros(shape=origin_shape, dtype=patch.dtype)
            origin_shape_file[y_min:y_max, x_min:x_max] = patch
            if patch_name.endswith('.png'):
                origin_shape_file = cv2.cvtColor(origin_shape_file, cv2.COLOR_BGR2GRAY)

            label_save_path = os.path.join(pasted_save_dir, patch_name)
            cv2.imwrite(label_save_path, origin_shape_file)

3.4 流程可视化

!unzip -oq data/data119773/PPYOLOTiny.zip

定义好训练集裁剪模型执行裁剪,最后还原以可视化(最后的测试集预测同理使用),这里重命名一下原/标注图的文件夹名称,加上 _origin 后缀。

!mv data/dataset/Train/JPEGImages data/dataset/Train/JPEGImages_origin
!rm -r data/dataset/Train/Annotations_origin  # 删除源标签,保留清洗之后的标签
!mv data/dataset/Train/Annotations data/dataset/Train/Annotations_origin
train_crop = DiscCrop(
    model_path='PPYOLOTiny',
    image_dir='data/dataset/Train/JPEGImages_origin',
    label_dir='data/dataset/Train/Annotations_origin')

然后执行裁剪(write_file=False 表示只记录裁剪的信息而不写入文件,为了防止内核重启导致的变量丢失而不想覆盖文件问题),结束之后会在相应同级目录下多出来两个 _patch 后缀文件夹。

train_crop.crop(write_file=True)

之后我们又将裁剪下来的 _patch 复原,生成 _patch_pasted 后缀文件夹。

train_crop.paste_origin(
    patch_dir='data/dataset/Train/JPEGImages_origin_patch')
train_crop.paste_origin(
    patch_dir='data/dataset/Train/Annotations_origin_patch')

当前训练集中的目录情况如下。

!tree data/dataset/Train -d
data/dataset/Train
├── Annotations_origin
├── Annotations_origin_patch
├── Annotations_origin_patch_pasted
├── JPEGImages_origin
├── JPEGImages_origin_patch
└── JPEGImages_origin_patch_pasted

6 directories

选择一张图片进行展示。

file_name = 'H0012'
image_list = [
    f'data/dataset/Train/JPEGImages_origin/{file_name}.jpg',
    f'data/dataset/Train/JPEGImages_origin_patch/{file_name}.jpg',
    f'data/dataset/Train/JPEGImages_origin_patch_pasted/{file_name}.jpg']
label_list = [
    f'data/dataset/Train/Annotations_origin/{file_name}.png',
    f'data/dataset/Train/Annotations_origin_patch/{file_name}.png',
    f'data/dataset/Train/Annotations_origin_patch_pasted/{file_name}.png']

plt.figure(figsize=(8, 6))
for i in range(6):
    plt.subplot(int('23'+str(i+1)))
    plt.axis('off')
    if i <= 2:
        plt.title(image_list[i].split('/')[-1])
        img = cv2.imread(image_list[i])
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    else:
        plt.title(label_list[i-3].split('/')[-1])
        img = cv2.imread(label_list[i-3], cv2.IMREAD_GRAYSCALE)
    plt.imshow(img)
plt.tight_layout()
plt.show()

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-39S7WCNu-1639095699400)(output_51_0.png)]

至此已经实现了视盘 ROI 区域的裁剪与复原步骤,后续就是构建模型进行分割。

4 视盘分割

此处将要基于 PaddleSeg 实现自定义分割网络 M-Net(修改版)进行视盘 ROI 区域的分割。

4.1 数据准备

因为 PaddleX 的数据划分比较方面,所以我们对 _patch 文件夹重命名,以满足 PaddleX 的格式。

所以裁剪下来的图片名称就变成了 JPEGImagesAnnotations

!mv data/dataset/Train/JPEGImages_origin_patch data/dataset/Train/JPEGImages
!mv data/dataset/Train/Annotations_origin_patch data/dataset/Train/Annotations

还需要注意的是,类别序号是递增的,所以需要覆盖写原来的标签,将 0-255 分别映射到 1(视盘)- 0(其他)。

注意:在步骤 [3.2.1] 构建目标检测数据集的时候,发现存在 40 张图片是没有标注的,这可能会影响训练,但此处是没有删除它们的。

label_path_list = sorted(glob.glob('data/dataset/Train/Annotations/*.png'))

for label_path in label_path_list:
    label = cv2.imread(label_path, cv2.IMREAD_GRAYSCALE)
    label = np.uint8((label == 0) * 1)
    cv2.imwrite(label_path, label)

80%(640)训练 + 20%验证(160);生成了三个文件:data/dataset/Train/{train_list, val_list, labels}.txt

!paddlex --split_dataset --format SEG\
    --dataset_dir data/dataset/Train\
    --val_value 0.2\
    --test_value 0

4.2 数据增强与读取

裁剪尺寸默认是 512x512,但如果整个区域超出原图边界,则会仅裁取原图部分而不 Padding,所以这里执行 Padding,然后水平翻转,归一化到 [-1, 1]。

import paddleseg.transforms as T

train_transforms = [
    T.Padding(target_size=(512, 512)),
    T.RandomHorizontalFlip(0.5),
    T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
]

eval_transforms = [
    T.Padding(target_size=(512, 512)),
    T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
]
train_dataset = paddleseg.datasets.Dataset(
    mode='train',
    num_classes=2,
    dataset_root='data/dataset/Train',
    train_path='data/dataset/Train/train_list.txt',
    transforms=train_transforms)

val_dataset = paddleseg.datasets.Dataset(
    mode='val',
    num_classes=2,
    dataset_root='data/dataset/Train',    
    val_path='data/dataset/Train/val_list.txt',
    transforms=eval_transforms)

4.3 构建自定义网络

常规常见下都是直接使用 PaddleSeg 内置网络,而很少自己去构建,但出于学习目的此处将自己构建网络并融入 PaddleSeg 中。

通过观察 PaddleSeg 源码可以发现,PaddleSeg 中的分类机制是 Softmax+Argmax,即输出通道数等于类别数,而引文使用的方法为每个通道进行 Sigmoid 预测两类,这导致我们需要对模型输出进行修改。

img

注意到上图的引文输出通道是 2,刚好符合本项目的 Softamx 2 分类的目的,所以直接按照网络架构搭建即可(但本项目的代码对每个卷积层都增加了批归一化层)。

import paddle.nn.functional as F
from paddle.nn import Layer, AvgPool2D, MaxPool2D, Conv2DTranspose, Upsample, Sequential, Conv2D
from paddleseg.cvlibs import manager
from paddleseg.models.layers import ConvBNReLU


@manager.MODELS.add_component
class MNet(Layer):
    def __init__(self, num_classes=2):
        super(MNet, self).__init__()

        self.input_down_sample = AvgPool2D(kernel_size=2)
        self.down_sample = MaxPool2D(kernel_size=2)

        self.cv1 = Sequential(
            ConvBNReLU(in_channels=3, out_channels=32, kernel_size=3),
            ConvBNReLU(in_channels=32, out_channels=32, kernel_size=3))

        self.cv2_head = Sequential(
            ConvBNReLU(in_channels=3, out_channels=64, kernel_size=3))
        self.cv2 = Sequential(
            ConvBNReLU(in_channels=96, out_channels=64, kernel_size=3),
            ConvBNReLU(in_channels=64, out_channels=64, kernel_size=3))

        self.cv3_head = Sequential(
            ConvBNReLU(in_channels=3, out_channels=128, kernel_size=3))
        self.cv3 = Sequential(
            ConvBNReLU(in_channels=192, out_channels=128, kernel_size=3),
            ConvBNReLU(in_channels=128, out_channels=128, kernel_size=3))

        self.cv4_head = Sequential(
            ConvBNReLU(in_channels=3, out_channels=256, kernel_size=3))
        self.cv4 = Sequential(
            ConvBNReLU(in_channels=384, out_channels=256, kernel_size=3),
            ConvBNReLU(in_channels=256, out_channels=256, kernel_size=3))

        self.cv_bottom = Sequential(
            ConvBNReLU(in_channels=256, out_channels=512, kernel_size=3),
            ConvBNReLU(in_channels=512, out_channels=512, kernel_size=3))

        self.cv4_ = Sequential(
            ConvBNReLU(in_channels=512, out_channels=256, kernel_size=3),
            ConvBNReLU(in_channels=256, out_channels=256, kernel_size=3))

        self.cv3_ = Sequential(
            ConvBNReLU(in_channels=256, out_channels=128, kernel_size=3),
            ConvBNReLU(in_channels=128, out_channels=128, kernel_size=3))

        self.cv2_ = Sequential(
            ConvBNReLU(in_channels=128, out_channels=64, kernel_size=3),
            ConvBNReLU(in_channels=64, out_channels=64, kernel_size=3))

        self.cv1_ = Sequential(
            ConvBNReLU(in_channels=64, out_channels=32, kernel_size=3),
            ConvBNReLU(in_channels=32, out_channels=32, kernel_size=3),
            ConvBNReLU(in_channels=32, out_channels=2, kernel_size=3))

        self.de_cv5 = Conv2DTranspose(in_channels=512, out_channels=256, kernel_size=2, stride=2)
        self.de_cv4 = Conv2DTranspose(in_channels=256, out_channels=128, kernel_size=2, stride=2)
        self.de_cv3 = Conv2DTranspose(in_channels=128, out_channels=64, kernel_size=2, stride=2)
        self.de_cv2 = Conv2DTranspose(in_channels=64, out_channels=32, kernel_size=2, stride=2)

        self.up_sample2 = Upsample(scale_factor=2)
        self.up_sample3 = Upsample(scale_factor=4)
        self.up_sample4 = Upsample(scale_factor=8)

        self.cv4_1x1 = Conv2D(in_channels=256, out_channels=num_classes, kernel_size=1)
        self.cv3_1x1 = Conv2D(in_channels=128, out_channels=num_classes, kernel_size=1)
        self.cv2_1x1 = Conv2D(in_channels=64, out_channels=num_classes, kernel_size=1)
        self.cv1_1x1 = Conv2D(in_channels=2, out_channels=num_classes, kernel_size=1)

    def encoder(self, x1, x2, x3, x4):
        o1 = self.cv1(x1)

        o2 = paddle.concat([self.cv2_head(x2), self.down_sample(o1)], axis=1)
        o2 = self.cv2(o2)

        o3 = paddle.concat([self.cv3_head(x3), self.down_sample(o2)], axis=1)
        o3 = self.cv3(o3)

        o4 = paddle.concat([self.cv4_head(x4), self.down_sample(o3)], axis=1)
        o4 = self.cv4(o4)

        bottom = self.cv_bottom(self.down_sample(o4))

        return o1, o2, o3, o4, bottom

    def decoder(self, o1, o2, o3, o4, bottom):
        o4 = paddle.concat([o4, self.de_cv5(bottom)], axis=1)
        o4 = self.cv4_(o4)

        o3 = paddle.concat([o3, self.de_cv4(o4)], axis=1)
        o3 = self.cv3_(o3)

        o2 = paddle.concat([o2, self.de_cv3(o3)], axis=1)
        o2 = self.cv2_(o2)

        o1 = paddle.concat([o1, self.de_cv2(o2)], axis=1)
        o1 = self.cv1_(o1)

        o2 = self.up_sample2(o2)
        o3 = self.up_sample3(o3)
        o4 = self.up_sample4(o4)

        return o1, o2, o3, o4

    def forward(self, x):
        x1 = x
        x2 = self.input_down_sample(x1)
        x3 = self.input_down_sample(x2)
        x4 = self.input_down_sample(x3)

        o1, o2, o3, o4, bottom = self.encoder(x1, x2, x3, x4)
        o1, o2, o3, o4 = self.decoder(o1, o2, o3, o4, bottom)

        o1 = self.cv1_1x1(o1)
        o2 = self.cv2_1x1(o2)
        o3 = self.cv3_1x1(o3)
        o4 = self.cv4_1x1(o4)
        o = paddle.add_n([o1, o2, o3, o4]) / 4.

        if self.training:
            return [o4, o3, o2, o1, o]
        else:
            return [o]

以上,通过装饰器实现自定义 PaddleSeg 模型,其余部分和飞桨基础 API 相同。

训练时的输出有5个,以利于梯度回传;推理阶段就只采用多尺度平均的结果。

model = MNet(num_classes=2)
paddle.summary(model, input_size=(None, 3, 512, 512))

4.4 配置与训练

装配 Warmup + Cosine-decay 学习率衰减策略的 Momentum 优化器,更新 10k 次参数,L2(0.0005)。

iters = 10000
train_batch_size = 4
learning_rate = 0.005

decayed_lr = paddle.optimizer.lr.CosineAnnealingDecay(
    learning_rate=learning_rate,
    T_max=iters)

decayed_lr = paddle.optimizer.lr.LinearWarmup(
    learning_rate=decayed_lr,
    warmup_steps=500,
    start_lr=0.0,
    end_lr=learning_rate)

optimizer = paddle.optimizer.Momentum(
    learning_rate=decayed_lr,
    momentum=0.9,
    weight_decay=paddle.regularizer.L2Decay(5e-4),
    parameters=model.parameters())

定义训练的损失函数,这里选择 DiceLoss。

from paddleseg.models import DiceLoss

losses = {
    'types': [DiceLoss()] * 5,
    'coef': [0.05, 0.1, 0.2, 0.4, 1]
}

装载配置信息,开始训练(训练40分钟,日志过长已清空)。

from paddleseg.core import train, evaluate, predict
train(
    train_dataset=train_dataset,
    val_dataset=val_dataset,

    model=model,
    optimizer=optimizer,
    losses=losses,

    iters=iters,
    batch_size=train_batch_size,

    save_interval=500,
    log_iters=100,
    num_workers=0,
    save_dir='output/MNet_512x512_B4_Dice_10k',
    use_vdl=False)

5 评估与预测

上一步已经完成的模型的训练。这里首先初始化之前构建的模型 MNet,然后载入训练好的参数(.pdparams)。

model = MNet(num_classes=2)

params_path = 'output/MNet_512x512_B4_Dice_10k/best_model/model.pdparams'
model_state_dict = paddle.load(params_path)
model.set_dict(model_state_dict)

复原模型之后先评估一下当前最优模型,使用水平翻转的增强评估。

evaluate(
    model,
    val_dataset,
    aug_eval=True,
    flip_horizontal=True)
  1. mIoU: 0.8790 Acc: 0.9361 Kappa: 0.8712
  2. Class IoU: [0.8894 0.8686]
  3. Class Acc: [0.9629 0.9055]

进行单张图片的预测结果可视化,这里挑选 1 张测试集的,若预测多张的话可以利用 glob.glob('../*.jpg') 生成列表。

这里要注意数据变换需要套上 T.Compose

test_crop = DiscCrop(
    model_path='PPYOLOTiny',
    image_dir='data/dataset/Test')

test_crop.crop()
image_list = ['data/dataset/Test_patch/T0006.jpg']

predict(
    model=model,
    model_path=params_path,
    transforms=T.Compose(eval_transforms),
    image_list=image_list,
    save_dir='demo')

将 PaddleSeg 的两种预测结果都还原至原来尺寸下的笛卡尔坐标图像。

test_crop.paste_origin(patch_dir='demo/added_prediction')
test_crop.paste_origin(patch_dir='demo/pseudo_color_prediction')

部分图像的预测结果和复原图的可视化情况。

image_list = [
    'demo/added_prediction/T0006.jpg',
    'demo/added_prediction_pasted/T0006.jpg',
    'demo/pseudo_color_prediction/T0006.png',
    'demo/pseudo_color_prediction_pasted/T0006.png']

plt.figure(figsize=(6, 6))
for i in range(4):
    plt.subplot(int('22'+str(i+1)))
    plt.axis('off')
    img = cv2.imread(image_list[i])
    plt.imshow(img[:, :, [2, 1, 0]])
plt.tight_layout()
寸下的笛卡尔坐标图像。


```python
test_crop.paste_origin(patch_dir='demo/added_prediction')
test_crop.paste_origin(patch_dir='demo/pseudo_color_prediction')

部分图像的预测结果和复原图的可视化情况。

image_list = [
    'demo/added_prediction/T0006.jpg',
    'demo/added_prediction_pasted/T0006.jpg',
    'demo/pseudo_color_prediction/T0006.png',
    'demo/pseudo_color_prediction_pasted/T0006.png']

plt.figure(figsize=(6, 6))
for i in range(4):
    plt.subplot(int('22'+str(i+1)))
    plt.axis('off')
    img = cv2.imread(image_list[i])
    plt.imshow(img[:, :, [2, 1, 0]])
plt.tight_layout()
plt.show()

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-KnYlQvtS-1639095699400)(output_88_0.png)]

如果需要用于比赛的提交,就使用 cv2.imread(path, cv2.IMREAD_GRAYSCALE)读入,然后根据 np.unique() 查看分析每个类别映射到哪个值,进行比赛提交标签的数值映射即可。

6 项目总结

本项目使用了 PaddleX - PP-YOLO Tniy 构建视盘区域检测模型,设计了一个极坐标变换与还原方案;通过装饰器自建 PaddleSeg 分割模型,实现如引文中所示的视盘分割架构方案。


我的 AI Studio 主页

Logo

学大模型,用大模型上飞桨星河社区!每天8点V100G算力免费领!免费领取ERNIE 4.0 100w Token >>>

更多推荐