《人工智能导论：案例与实践》 基于U-Net实现宠物图像分割

~/work/Dataset
||--images|--trimaps

%cd /home/aistudio/work# 初次运行时将注释取消，以便解压文件# 如果已经解压过，不需要运行此段代码，否则由于文件已经存在，解压时会报错#!unzip -q /home/aistudio/data/data79141/Dataset.zip

# coding=utf-8# 导入环境import osimport randomimport cv2import numpy as npfrom PIL import Imagefrom paddle.io import Datasetimport matplotlib.pyplot as plt# 在notebook中使用matplotlib.pyplot绘图时，需要添加该命令进行显示%matplotlib inlineimport paddleimport paddle.nn.functional as Fimport paddle.nn as nn

IMAGE_SIZE = (160, 160)
train_images_path = "/home/aistudio/work/Dataset/images"label_images_path = "/home/aistudio/work/Dataset/trimaps"image_count = len([os.path.join(train_images_path, image_name) for image_name in os.listdir(train_images_path) if image_name.endswith('.jpg')])# 对文件夹内的图像按文件名排序def _sort_images(image_dir, image_type):
    files = []    for image_name in os.listdir(image_dir):        if image_name.endswith('.{}'.format(image_type)) and not image_name.startswith('.'):
            files.append(os.path.join(image_dir, image_name))    return sorted(files)# 生成实验需要的train、test标签文件，标签格式如下：‘images/image1.jpg labels/label1.png’def write_file(images, labels, eval_num):
    train_file = open('/home/aistudio/work/Dataset/train.txt', 'w')
    test_file = open('/home/aistudio/work/Dataset/test.txt', 'w')
    predict_file = open('/home/aistudio/work/Dataset/predict.txt', 'w')    # 随机打乱文件顺序并保存文件路径
    image_index = [i for i in range(image_count)]
    random.shuffle(image_index)    for n in range(image_count):        if n < eval_num:
            test_file.write('{} {}\n'.format(images[image_index[n]], labels[image_index[n]]))
            predict_file.write('{} {}\n'.format(images[image_index[n]], labels[image_index[n]]))        else:
            train_file.write('{} {}\n'.format(images[image_index[n]], labels[image_index[n]]))

images = _sort_images(train_images_path, 'jpg')
labels = _sort_images(label_images_path, 'png')
eval_num = int(image_count * 0.20)# 对数据集进行处理，划分训练集、测试集write_file(images, labels, eval_num)

# 使用PIL库读取图片，并转为numpy array的格式# （思考）您也可以尝试修改下面代码中jpg文件夹包含的图片标号image_000XX，可视化数据集中的其他图片。image = Image.open('/home/aistudio/work/Dataset/images/British_Shorthair_128.jpg')
image = np.array(image)
label = Image.open('/home/aistudio/work/Dataset/trimaps/British_Shorthair_128.png')
label = np.array(label)# 画出读取的图片plt.subplot(1,2,1),
plt.title('Train Image')
plt.imshow(image)
plt.axis('off')

plt.subplot(1,2,2),
plt.title('Label')
plt.imshow(label)
plt.axis('off')

# 查看原始图片以及标签图片的形状print(image.shape, label.shape)

# 使用opencv库缩放图像def resize(im, target_size=608, interp=cv2.INTER_LINEAR):
    if isinstance(target_size, list) or isinstance(target_size, tuple):
        w = target_size[0]
        h = target_size[1]    else:
        w = target_size
        h = target_size
    im = cv2.resize(im, (w, h), interpolation=interp)    return im

# 随机缩放图像class ResizeStepScaling:
    def __init__(self, min_scale_factor=0.5, max_scale_factor=2.0, scale_step_size=0.25):
        self.min_scale_factor = min_scale_factor
        self.max_scale_factor = max_scale_factor
        self.scale_step_size = scale_step_size    def __call__(self, im, label=None):
        # 在[min_scale_factor， max_scale_factor]范围内，以scale_step_size为步长生成尺度列表
        num_steps = int((self.max_scale_factor - self.min_scale_factor) / self.scale_step_size + 1)
        scale_factors = np.linspace(self.min_scale_factor, self.max_scale_factor, num_steps).tolist()        # 随机打乱尺度列表
        np.random.shuffle(scale_factors)        # 取出列表中第一个尺度值
        scale_factor = scale_factors[0]
        w = int(round(scale_factor * im.shape[1]))
        h = int(round(scale_factor * im.shape[0]))        # 缩放图像
        im = resize(im, (w, h), cv2.INTER_LINEAR)        # 如果传入了标签图像，则进行缩放
        if label is not None:
            label = resize(label, (w, h), cv2.INTER_NEAREST)        if label is None:            return (im, )        else:            return (im, label)

# 随机裁剪图像class RandomPaddingCrop:
    def __init__(self, crop_size=(512, 512), im_padding_value=(127.5, 127.5, 127.5), label_padding_value=255):
        self.crop_size = crop_size
        self.im_padding_value = im_padding_value
        self.label_padding_value = label_padding_value    def __call__(self, im, label=None):
        crop_width = self.crop_size[0]
        crop_height = self.crop_size[1]

        img_height = im.shape[0]
        img_width = im.shape[1]        # 如果剪裁尺寸等于原图像尺寸，则直接返回原图像
        if img_height == crop_height and img_width == crop_width:            if label is None:                return (im, )            else:                return (im, label)        else:            # 如果目标裁切尺寸大于原始图像，则将添加右下角的填充
            pad_height = max(crop_height - img_height, 0)
            pad_width = max(crop_width - img_width, 0)            if (pad_height > 0 or pad_width > 0):
                im = cv2.copyMakeBorder(im, 0, pad_height, 0, pad_width, cv2.BORDER_CONSTANT, value=self.im_padding_value)                if label is not None:
                    label = cv2.copyMakeBorder(label, 0, pad_height, 0, pad_width, cv2.BORDER_CONSTANT, value=self.label_padding_value)
                img_height = im.shape[0]
                img_width = im.shape[1]            # 随机裁剪图像
            if crop_height > 0 and crop_width > 0:
                h_off = np.random.randint(img_height - crop_height + 1)
                w_off = np.random.randint(img_width - crop_width + 1)

                im = im[h_off:(crop_height + h_off), w_off:(w_off + crop_width), :]                if label is not None:
                    label = label[h_off:(crop_height + h_off), w_off:(w_off + crop_width)]        if label is None:            return (im, )        else:            return (im, label)

# 水平翻转图像def horizontal_flip(im):
    if len(im.shape) == 3:
        im = im[:, ::-1, :]    elif len(im.shape) == 2:
        im = im[:, ::-1]    return im

# 随机翻转图像class RandomHorizontalFlip:
    def __init__(self, prob=0.5):
        self.prob = prob    def __call__(self, im, label=None):
        # 以0.5的概率水平翻转图像
        if random.random() < self.prob:
            im = horizontal_flip(im)            if label is not None:
                label = horizontal_flip(label)        if label is None:            return (im, )        else:            return (im, label)

# 归一化class Normalize:
    def __init__(self, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)):
        self.mean = mean
        self.std = std        from functools import reduce        if reduce(lambda x, y: x * y, self.std) == 0:            raise ValueError('{}: std is invalid!'.format(self))    def __call__(self, im, label=None):
        # 图像归一化
        mean = np.array(self.mean)[np.newaxis, np.newaxis, :]
        std = np.array(self.std)[np.newaxis, np.newaxis, :]
        im = im.astype(np.float32, copy=False) / 255.0
        im -= mean
        im /= std        if label is None:            return (im, )        else:            return (im, label)

# 图像预处理方法汇总def transforms(im, label=None, mode='train'):
    # 如果是训练集，需要进行数据增广
    if mode == 'train':
        transform_list = [ResizeStepScaling, RandomPaddingCrop, RandomHorizontalFlip, Normalize]    # 如果是测试集或验证集，仅需要归一化图像
    else:
        transform_list = [Normalize]    if isinstance(im, str):
        im = cv2.imread(im).astype('float32')    if isinstance(label, str):        # 将标签'1,2,3'调整为'0,1,2'
        label = np.asarray(Image.open(label))-1
    im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)    for transform_class in transform_list:
        op = transform_class()
        outputs = op(im, label)
        im = outputs[0]        if len(outputs) == 2:
            label = outputs[1]
    im = np.transpose(im, (2, 0, 1))    return (im, label)

# 定义数据读取类class Dataset(paddle.io.Dataset):
    def __init__(self, dataset_root, mode='train', train_path=None, val_path=None, ignore_index=255):
        self.dataset_root = dataset_root
        self.file_list = list()
        self.mode = mode.lower()
        self.num_classes = 3
        self.ignore_index = ignore_index
        self.dataset_root = dataset_root        # 获取train、test标签文件
        if mode == 'train':
            file_path = train_path   
        else:
            file_path = val_path        # 获取原始图片和标签图片列表
        with open(file_path, 'r') as f:            for line in f:
                items = line.strip().split(' ')
                image_path = os.path.join(self.dataset_root, items[0])
                label_path = os.path.join(self.dataset_root, items[1])
                self.file_list.append([image_path, label_path])    # 获取数据
    def __getitem__(self, idx):
        image_path, label_path = self.file_list[idx]        # 如果是验证集，仅对原始图片进行数据预处理
        if self.mode == 'eval':
            im, _ = transforms(im=image_path, mode='val')            # 将标签'1,2,3'调整为'0,1,2'
            label = np.asarray(Image.open(label_path))-1
            label = label[np.newaxis, :, :]            return im, label        # 如果是训练集，需要对原始图片和标签图片进行数据增广处理
        else:
            im, label = transforms(im=image_path, label=label_path, mode='train')            return im, label    def __len__(self):
        return len(self.file_list)

DATADIR = '/home/aistudio/work/Dataset'train_path = '/home/aistudio/work/Dataset/train.txt'eval_path = '/home/aistudio/work/Dataset/test.txt'# 创建数据读取类train_dataset = Dataset(DATADIR, mode='train', train_path=train_path)
eval_dataset = Dataset(DATADIR, mode='eval', val_path=eval_path)# 使用paddle.io.DataLoader创建数据读取器，并设置batchsize，进程数量num_workers等参数train_sampler = paddle.io.DistributedBatchSampler(train_dataset, batch_size=8, shuffle=True, drop_last=True)
train_loader = paddle.io.DataLoader(train_dataset, batch_sampler=train_sampler, num_workers=0, return_list=True)
valid_sampler = paddle.io.DistributedBatchSampler(eval_dataset, batch_size=1, shuffle=False, drop_last=False)
valid_loader = paddle.io.DataLoader(eval_dataset, batch_sampler=valid_sampler, num_workers=0, return_list=True)

class UNet(nn.Layer):
    # 继承paddle.nn.Layer定义网络结构
    def __init__(self, num_classes=3):
        # 初始化函数
        super().__init__()        # 定义编码器
        self.encode = Encoder()        # 定义解码器
        self.decode = Decoder()        # 分类模块
        self.cls = nn.Conv2D(in_channels=64, out_channels=num_classes, kernel_size=3, stride=1, padding=1)    def forward(self, x):
        # 前向计算
        logit_list = []        # 编码运算
        x, short_cuts = self.encode(x)        # 解码运算
        x = self.decode(x, short_cuts)        # 分类运算
        logit = self.cls(x)
        logit_list.append(logit)        return logit_list

class ConvBNReLU(nn.Layer):
    def __init__(self, in_channels, out_channels, kernel_size, padding='same'):
        # 初始化函数
        super().__init__()        # 定义卷积层
        self._conv = nn.Conv2D(in_channels, out_channels, kernel_size, padding=padding)        # 定义批归一化层
        self._batch_norm = nn.SyncBatchNorm(out_channels)    def forward(self, x):
        # 前向计算
        x = self._conv(x)
        x = self._batch_norm(x)
        x = F.relu(x)        return x

class Encoder(nn.Layer):
    def __init__(self):
        # 初始化函数
        super().__init__()        # # 封装两个ConvBNReLU模块
        self.double_conv = nn.Sequential(ConvBNReLU(3, 64, 3), ConvBNReLU(64, 64, 3))        # 定义下采样通道数
        down_channels = [[64, 128], [128, 256], [256, 512], [512, 512]]        # 封装下采样模块
        self.down_sample_list = nn.LayerList([self.down_sampling(channel[0], channel[1]) for channel in down_channels])    
    # 定义下采样模块
    def down_sampling(self, in_channels, out_channels):
        modules = []        # 添加最大池化层
        modules.append(nn.MaxPool2D(kernel_size=2, stride=2))        # 添加两个ConvBNReLU模块
        modules.append(ConvBNReLU(in_channels, out_channels, 3))
        modules.append(ConvBNReLU(out_channels, out_channels, 3))        return nn.Sequential(*modules)    def forward(self, x):
        # 前向计算
        short_cuts = []        # 卷积运算
        x = self.double_conv(x)        # 下采样运算
        for down_sample in self.down_sample_list:
            short_cuts.append(x)
            x = down_sample(x)        return x, short_cuts

# 定义上采样模块class UpSampling(nn.Layer):
    def __init__(self, in_channels, out_channels):
        # 初始化函数
        super().__init__()
        in_channels *= 2
        # 封装两个ConvBNReLU模块
        self.double_conv = nn.Sequential(ConvBNReLU(in_channels, out_channels, 3), ConvBNReLU(out_channels, out_channels, 3))    def forward(self, x, short_cut):
        # 前向计算
        # 定义双线性插值模块
        x = F.interpolate(x, paddle.shape(short_cut)[2:], mode='bilinear')        # 特征图拼接
        x = paddle.concat([x, short_cut], axis=1)        # 卷积计算
        x = self.double_conv(x)        return x

# 定义解码器class Decoder(nn.Layer):
    def __init__(self):
        # 初始化函数
        super().__init__()        # 定义上采样通道数
        up_channels = [[512, 256], [256, 128], [128, 64], [64, 64]]        # 封装上采样模块
        self.up_sample_list = nn.LayerList([UpSampling(channel[0], channel[1]) for channel in up_channels])    def forward(self, x, short_cuts):
        # 前向计算
        for i in range(len(short_cuts)):            # 上采样计算
            x = self.up_sample_list[i](x, short_cuts[-(i + 1)])        return x

# 声明定义好的UNet模型model = UNet()

# 定义总的训练批次数iters = 10000# （思考）您也可以尝试增加训练批次数（如30000），观察训练批次数的变化对模型效果的影响# 定义学习率参数lr = paddle.optimizer.lr.PolynomialDecay(0.01, iters, end_lr=0.0, power=0.9)# 使用Momentum优化器，学习率设置为0.001optimizer = paddle.optimizer.Momentum(learning_rate=lr, momentum=0.9, parameters=model.parameters(), weight_decay=4.0e-5)#（思考）您也可以尝试使用SGD、Adam等其他优化器，观察训练结果

class CrossEntropyLoss(nn.Layer):
    def __init__(self, ignore_index=255):
        super(CrossEntropyLoss, self).__init__()
        self.ignore_index = ignore_index
        self.EPS = 1e-8

    def forward(self, logit, label):
        logit = paddle.transpose(logit, [0, 2, 3, 1])        # 计算交叉熵损失，不计算损失的平均值
        loss = F.cross_entropy(logit, label, ignore_index=self.ignore_index, reduction='none')        # 计算掩膜
        mask = label != self.ignore_index
        mask = paddle.cast(mask, 'float32')        # 计算损失值
        loss = loss * mask
        avg_loss = paddle.mean(loss) / (paddle.mean(mask) + self.EPS)

        label.stop_gradient = True
        mask.stop_gradient = True
        return avg_loss

# 分别计算宠物预测结果和真实标签的相交区域、宠物预测结果区域、真实标签区域的面积def calculate_area(pred, label, num_classes=3, ignore_index=255):
    if len(pred.shape) == 4:
        pred = paddle.squeeze(pred, axis=1)    if len(label.shape) == 4:
        label = paddle.squeeze(label, axis=1)    # 生成图像区域的掩膜
    mask = label != ignore_index
    pred = pred + 1
    label = label + 1
    pred = pred * mask
    label = label * mask    # 将预测结果和标签转换为one hot格式
    pred = F.one_hot(pred, num_classes + 1)
    label = F.one_hot(label, num_classes + 1)
    pred = pred[:, :, :, 1:]
    label = label[:, :, :, 1:]

    pred_area = []
    label_area = []
    intersect_area = []    # 循环遍历每个类别
    for i in range(num_classes):
        pred_i = pred[:, :, :, i]
        label_i = label[:, :, :, i]        # 计算预测区域的面积
        pred_area_i = paddle.sum(pred_i)        # 计算真实区域的面积
        label_area_i = paddle.sum(label_i)        # 计算预测区域和真实区域相交的面积
        intersect_area_i = paddle.sum(pred_i * label_i)
        pred_area.append(pred_area_i)
        label_area.append(label_area_i)
        intersect_area.append(intersect_area_i)
    pred_area = paddle.concat(pred_area)
    label_area = paddle.concat(label_area)
    intersect_area = paddle.concat(intersect_area)    return intersect_area, pred_area, label_area

# 计算平均ioudef mean_iou(intersect_area, pred_area, label_area):
    intersect_area = intersect_area.numpy()
    pred_area = pred_area.numpy()
    label_area = label_area.numpy()    # 计算预测区域和真实区域的并集部分
    union = pred_area + label_area - intersect_area
    class_iou = []    # 计算iou
    for i in range(len(intersect_area)):        if union[i] == 0:
            iou = 0
        else:
            iou = intersect_area[i] / union[i]
        class_iou.append(iou)    # 计算各个类别的iou均值
    miou = np.mean(class_iou)    return np.array(class_iou), miou

# 计算平均准确率def accuracy(intersect_area, pred_area):
    intersect_area = intersect_area.numpy()
    pred_area = pred_area.numpy()
    class_acc = []    # 计算准确率
    for i in range(len(intersect_area)):        if pred_area[i] == 0:
            acc = 0
        else:
            acc = intersect_area[i] / pred_area[i]
        class_acc.append(acc)    # 计算平均准确率
    macc = np.sum(intersect_area) / np.sum(pred_area)    return np.array(class_acc), macc

# 定义训练过程# 调用GPU进行运算use_gpu = Truepaddle.set_device('gpu:0') if use_gpu else paddle.set_device('cpu')print('start training ... ')# 将模型切换到训练模式model.train()

avg_loss = 0.0iters_per_epoch = len(train_sampler)
best_miou = -1.0best_model_iter = -1# 训练模型 now_iter = 0while now_iter < iters:    for data in train_loader:
        now_iter += 1
        if now_iter > iters:            break
        images = data[0]
        labels = data[1].astype('int64')        # 运行模型前向计算，得到预测值
        logits = model(images)[0]        # 计算损失
        loss_func = CrossEntropyLoss()
        loss= loss_func(logits, labels)        # 反向传播，更新权重，清除梯度
        loss.backward()
        optimizer.step()
        optimizer._learning_rate.step()
        model.clear_gradients()
        avg_loss += loss.numpy()[0]        # 每隔100个iters打印当前指标
        if (now_iter) % 100 == 0 :
            avg_loss /= 100
            remain_iters = iters - now_iter
            lr_now = optimizer.get_lr()            print("[TRAIN] epoch={}, iter={}/{}, loss={:.4f}, lr={:.6f}" .format((now_iter - 1) // iters_per_epoch + 1, now_iter, iters, avg_loss, lr_now))        
            # 将avg_loss重置为0
            avg_loss = 0.0

        # 每1000个iters训练完成后，进行模型评估
        if (now_iter % 1000 == 0 or now_iter == iters) and (eval_dataset is not None):            # 将模型切换到评估模式
            model.eval()
            intersect_area_all = 0
            pred_area_all = 0
            label_area_all = 0
            with paddle.no_grad():                for eval_iter, (im, label) in enumerate(valid_loader):
                    label = label.astype('int64')                    # 运行模型前向计算，得到预测值
                    pred = model(im)[0]
                    pred = paddle.argmax(pred, axis=1, keepdim=True, dtype='int32')                    # 计算相交区域、预测区域、标签区域
                    intersect_area, pred_area, label_area = calculate_area(pred, label, num_classes=3, ignore_index=255)
                    intersect_area_all = intersect_area_all + intersect_area
                    pred_area_all = pred_area_all + pred_area
                    label_area_all = label_area_all + label_area            # 计算平均iou
            class_iou, miou = mean_iou(intersect_area_all, pred_area_all, label_area_all)            # 计算准确率
            class_acc, acc = accuracy(intersect_area_all, pred_area_all)            print("[EVAL] #Images={} mIoU={:.4f} Acc={:.4f}".format(len(eval_dataset), miou, acc))            # 如果当前的平均iou大于最佳的平均iou，保存最佳模型
            if miou > best_miou:
                best_miou = miou
                best_model_iter = now_iter
                paddle.save(model.state_dict(), 'best_model.pdparams')            print('[EVAL] The model with the best validation mIoU ({:.4f}) was saved at iter {}.' .format(best_miou, best_model_iter))
            model.train()

# 保存模型参数paddle.save(model.state_dict(), 'last_model.pdparams')# 保存优化器信息和相关参数，方便继续训练paddle.save(optimizer.state_dict(), 'last_model.pdopt')

# 开启0号GPUuse_gpu = Truepaddle.set_device('gpu:0') if use_gpu else paddle.set_device('cpu')print('start evaluation .......')# 实例化模型model = UNet()# 加载模型参数params_file_path="best_model.pdparams"model_state_dict = paddle.load(params_file_path)
model.load_dict(model_state_dict)# 将模型切换到评估模式model.eval()

intersect_area_all = 0pred_area_all = 0label_area_all = 0with paddle.no_grad():    for eval_iter, (im, label) in enumerate(valid_loader):
        label = label.astype('int64')        # 运行模型前向计算，得到预测值
        pred = model(im)[0]
        pred = paddle.argmax(pred, axis=1, keepdim=True, dtype='int32')        # 计算相交区域、预测区域、标签区域
        intersect_area, pred_area, label_area = calculate_area(pred, label, 3, ignore_index=255)
        intersect_area_all = intersect_area_all + intersect_area
        pred_area_all = pred_area_all + pred_area
        label_area_all = label_area_all + label_area# 计算平均iouclass_iou, miou = mean_iou(intersect_area_all, pred_area_all, label_area_all)# 计算准确率class_acc, acc = accuracy(intersect_area_all, pred_area_all)print("[EVAL] #Images={} mIoU={:.4f} Acc={:.4f}".format(len(eval_dataset), miou, acc))

# 开启0号GPUuse_gpu = Truepaddle.set_device('gpu:0') if use_gpu else paddle.set_device('cpu')# 实例化模型model = UNet()# 加载模型参数params_file_path="best_model.pdparams"model_state_dict = paddle.load(params_file_path)
model.load_dict(model_state_dict)

model.eval()# 读取test_list.txt中记录的第一张图片的路径及标签with open('/home/aistudio/work/Dataset/predict.txt') as test_file:
    img_path, label_path = test_file.readline().split(' ')# 读取图片并对图片进行预处理img = transforms(im=img_path, mode='val')[0]# 读取标签图片label = np.asarray(Image.open(label_path.split('\n')[0]))# 为图片添加batch维度batch_img = np.expand_dims(img, 0)
input_tensor = paddle.to_tensor(batch_img)# 运行模型前向计算，得到预测图片pred = model(input_tensor)[0]# 算法训练时，将标签统一减1，因此需要将最终结果恢复原样pred = paddle.argmax(pred, axis=1, keepdim=True, dtype='int32') + 1pred = pred[0].squeeze().numpy()# 可视化原始图片image = np.asarray(Image.open(img_path))
plt.subplot(3, 3, 1)
plt.imshow(image)
plt.title('Input Image')
plt.axis("off")# 可视化真实标签图片plt.subplot(3, 3, 2)
plt.imshow(label, cmap='gray')
plt.title('Label')
plt.axis("off")# 可视化预测结果图片plt.subplot(3, 3, 3)
plt.imshow(pred, cmap='gray')
plt.title('Predict')
plt.axis("off")