Сегментация изображений с глубоким обучением. Воспроизведение бумаги PSPNet (визуализация для обучения и тестирования).

1. Введение в PSPNet

PSPNet (сеть анализа сцены пирамиды) взята из статьи в CVPR2017. Китайский перевод — сеть анализа сцены пирамиды, которая в основном используется для сегментации изображений. В этой архитектуре представлен модуль Pyramid Pooling для сбора контекстной информации в различных масштабах. Pyramid Pooling может извлекать глобальную и локальную контекстную информацию в разных масштабах, помогая лучше понять семантическое содержание изображений и тем самым улучшая производительность сегментации.

1. Объяснение принципа

• (а) Входное изображение
• (б) Используйте предварительно обученную модель ResNet для получения карты объектов.
• (c) Используйте пирамидальный пул для получения представлений различных субрегионов и формируйте представления объектов, содержащие локальную и глобальную контекстную информацию, посредством повышения дискретизации и объединения.
• (d) Отправьте объекты на слой свертки, чтобы получить результаты прогнозирования на уровне пикселей.

2. Бумажное пояснение

Модуль объединения пирамид объединяет функции четырех разных масштабов пирамиды. Красным отмечено глобальное объединение в пул для получения выходных данных одного бункера. Следующие уровни пирамиды делят карту объектов на разные субрегионы и формируют объединенные представления в разных местах. Выходные данные разных уровней модуля объединения пирамид содержат карты объектов разных размеров. Чтобы сохранить вес глобальных функций, после каждого уровня пирамиды используется сверточный слой 1 × 1, чтобы уменьшить размерность представления контекста до 1/N исходного измерения (если размер уровня пирамиды равен N).

Затем низкоразмерная карта объектов подвергается повышающей дискретизации , и объекты того же размера, что и исходная карта объектов, получаются посредством билинейной интерполяции . Наконец, функции на разных уровнях соединяются как глобальные функции окончательного результата объединения пирамиды.

3. Сетевая модель

class PSPNet(BaseModel):
    def __init__(self, num_classes, in_channels=3, backbone='resnet152', pretrained=True, use_aux=True, freeze_bn=False, freeze_backbone=False):
        super(PSPNet, self).__init__()
        norm_layer = nn.BatchNorm2d  # 用于规范化的层类型

        # 使用getattr根据backbone参数选择合适的骨干网络模型，并可能加载预训练权重
        model = getattr(resnet, backbone)(pretrained, norm_layer=norm_layer)
        m_out_sz = model.fc.in_features  # 提取骨干网络的输出特征通道数

        self.use_aux = use_aux  # 是否使用辅助分割分支

        # 初始卷积层，根据in_channels来调整输入通道数
        self.initial = nn.Sequential(*list(model.children())[:4])
        if in_channels != 3:
            self.initial[0] = nn.Conv2d(in_channels, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.initial = nn.Sequential(*self.initial)

        # 骨干网络的不同层
        self.layer1 = model.layer1
        self.layer2 = model.layer2
        self.layer3 = model.layer3
        self.layer4 = model.layer4

        # 主要分割分支，包括特征融合和分割输出
        self.master_branch = nn.Sequential(
            PSPModule(m_out_sz, bin_sizes=[1, 2, 3, 6], norm_layer=norm_layer),  # 特征融合模块
            nn.Conv2d(m_out_sz // 4, num_classes, kernel_size=1)  # 分割输出卷积层
        )

        # 辅助分割分支，可选，用于训练时帮助主分割任务
        self.auxiliary_branch = nn.Sequential(
            nn.Conv2d(m_out_sz // 2, m_out_sz // 4, kernel_size=3, padding=1, bias=False),
            norm_layer(m_out_sz // 4),
            nn.ReLU(inplace=True),
            nn.Dropout2d(0.1),
            nn.Conv2d(m_out_sz // 4, num_classes, kernel_size=1)
        )

        # 初始化网络权重
        initialize_weights(self.master_branch, self.auxiliary_branch)

    def forward(self, x):
        input_size = (x.size()[2], x.size()[3])  # 记录输入图像的尺寸

        x = self.initial(x)  # 初始卷积层
        x = self.layer1(x)  # 第一层
        x = self.layer2(x)  # 第二层
        x_aux = self.layer3(x)  # 第三层，用于辅助分割分支
        x = self.layer4(x)  # 第四层

        output = self.master_branch(x)  # 主要分割分支
        output = F.interpolate(output, size=input_size, mode='bilinear')  # 插值操作，将分割输出大小调整为输入大小
        output = output[:, :, :input_size[0], :input_size[1]]  # 调整输出的尺寸以匹配输入

        # 如果在训练模式下且使用辅助分割分支，还生成辅助分割输出
        if self.training and self.use_aux:
            aux = self.auxiliary_branch(x_aux)
            aux = F.interpolate(aux, size=input_size, mode='bilinear')  # 调整辅助分割输出大小
            aux = aux[:, :, :input_size[0], :input_size[1]]  # 调整输出的尺寸以匹配输入
            return output, aux  # 返回主分割输出和辅助分割输出
        return output  # 只返回主分割输出

其中，PSPModule类的定义如下

class PSPModule(nn.Module):
    def __init__(self, in_channels, bin_sizes, norm_layer):
        super(_PSPModule, self).__init__()

        # 计算每个池化分支的输出通道数
        out_channels = in_channels // len(bin_sizes)

        # 创建池化分支，将它们存储在一个 ModuleList 中
        self.stages = nn.ModuleList([self._make_stages(in_channels, out_channels, b_s, norm_layer) 
                                                        for b_s in bin_sizes])

        # 创建特征融合模块（bottleneck）
        self.bottleneck = nn.Sequential(
            nn.Conv2d(in_channels + (out_channels * len(bin_sizes)), out_channels, 
                                    kernel_size=3, padding=1, bias=False),  # 卷积层
            norm_layer(out_channels),  # 规范化层
            nn.ReLU(inplace=True),  # ReLU激活函数
            nn.Dropout2d(0.1)  # 2D Dropout层
        )

    def _make_stages(self, in_channels, out_channels, bin_sz, norm_layer):
        # 创建池化分支的内部结构
        prior = nn.AdaptiveAvgPool2d(output_size=bin_sz)  # 自适应平均池化层
        conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)  # 卷积层
        bn = norm_layer(out_channels)  # 规范化层
        relu = nn.ReLU(inplace=True)  # ReLU激活函数
        return nn.Sequential(prior, conv, bn, relu)  # 返回池化分支的Sequential模块

    def forward(self, features):
        h, w = features.size()[2], features.size()[3]  # 获取输入特征的高度和宽度

        pyramids = [features]  # 存储原始特征到金字塔中

        # 遍历每个池化分支，对特征进行插值操作并存储在金字塔中
        pyramids.extend([F.interpolate(stage(features), size=(h, w), mode='bilinear', 
                                        align_corners=True) for stage in self.stages])

        # 将金字塔中的特征拼接在一起并通过特征融合模块
        output = self.bottleneck(torch.cat(pyramids, dim=1))
        return output  # 返回特征融合后的输出

Этот класс используется для выполнения пирамидального пула и операций объединения объектов и объединения их в представление объектов с более богатой семантической информацией.

Основная идея PSPNet заключается в использовании 4-уровневой пирамидальной структуры.Ядро пула может охватывать все, половину и небольшие части изображения, то есть

 self.stages = nn.ModuleList([self._make_stages(in_channels, out_channels, b_s, norm_layer) 
                                                  for b_s in bin_sizes])

В приведенном выше коде self.stages содержит несколько ветвей пула, а bin_sizes — это список, содержащий 4 элемента, соответствующие 4 различным ветвям пула. Проходя bin_sizes, с помощью метода make_stages создаются 4 ветки пула. Каждая ветвь пула состоит из слоя адаптивного среднего пула, сверточного слоя, уровня нормализации и уровня активации ReLU. Это формирует часть пирамидальной структуры, предназначенную для извлечения признаков.

Четыре ветви объединения имеют разные размеры восприимчивых полей для захвата информации изображения в разных масштабах.

Другими словами, каждый элемент в self.stages представляет уровень пирамиды, отражая 4-уровневую структуру пирамиды. Прямой метод будет проходить через эти ветви объединения и выполнять операции интерполяции над входными объектами, изменяя их размер до того же размера, что и исходные объекты для объединения объектов.

2. Развертывание и внедрение

Мое окружение

• Операционная система: Win11
• Язык: питон3.10
• IDE: PyCharm 2023
• Графический процессор: RTX 4060

1、ПАСКАЛЬ ЛОС 2012

Набор данных использует классический PASCAL VOC 2012, стандартный набор данных для исследований в области компьютерного зрения, который предоставляет данные изображений и соответствующие аннотации для различных задач. Он содержит 20 различных категорий объектов, таких как самолеты, велосипеды, автомобили, собаки, кошки, стулья и т. д., а также категорию «фон». Эти изображения собраны из реального мира, охватывают разные сцены и ракурсы и представляют собой обычные повседневные предметы. Что касается размера, он содержит 1464 обучающих изображения, 1449 проверочных изображений и 1456 тестовых изображений. Каждое изображение сопровождается подробной аннотационной информацией, включая ограничивающие рамки для каждого экземпляра объекта (задача обнаружения объекта) и метки семантической сегментации на уровне пикселей (задача семантической сегментации).

2. Модельное обучение

核心训练部分的代码如下：

def _train_epoch(self, epoch):
    self.logger.info('\n')  # 打印日志信息

    self.model.train()  # 设置模型为训练模式
    if self.config['arch']['args']['freeze_bn']:  
        if isinstance(self.model, torch.nn.DataParallel):
            self.model.module.freeze_bn() 
        else:
            self.model.freeze_bn()  
    self.wrt_mode = 'train'  # 设置写入模式为'train'

    tic = time.time()  # 记录当前时间
    self._reset_metrics()  # 重置度量指标
    tbar = tqdm(self.train_loader, ncols=130)  # 创建一个进度条以迭代训练数据集

    for batch_idx, (data, target) in enumerate(tbar):  # 遍历训练数据
        self.data_time.update(time.time() - tic)  # 更新数据加载时间

        self.lr_scheduler.step(epoch=epoch - 1)  # 根据当前训练的epoch调整学习率

        # LOSS & OPTIMIZE
        self.optimizer.zero_grad()  # 清零梯度
        output = self.model(data)  # 前向传播，获取模型输出
        if self.config['arch']['type'][:3] == 'PSP':  
            assert output[0].size()[2:] == target.size()[1:]  # 检查输出和目标的空间尺寸匹配
            assert output[0].size()[1] == self.num_classes  # 检查输出通道数与类别数匹配
            loss = self.loss(output[0], target)  # 计算损失
            loss += self.loss(output[1], target) * 0.4  # 添加辅助损失，加权为0.4
            output = output[0]  # 将主要输出作为最终输出
        else:
            assert output.size()[2:] == target.size()[1:]  
            assert output.size()[1] == self.num_classes  
            loss = self.loss(output, target) 

        if isinstance(self.loss, torch.nn.DataParallel):  
            loss = loss.mean()  # 计算损失的均值
        loss.backward()  # 反向传播，计算梯度
        self.optimizer.step()  # 更新模型参数
        self.total_loss.update(loss.item())  # 更新总损失

        # measure elapsed time
        self.batch_time.update(time.time() - tic)  # 更新批次处理时间
        tic = time.time()

        # LOGGING & TENSORBOARD
        if batch_idx % self.log_step == 0:  # 每隔一定步数记录一次日志和TensorBoard
            self.wrt_step = (epoch - 1) * len(self.train_loader) + batch_idx  # 当前步数
            self.writer.add_scalar(f'{
      
      self.wrt_mode}/loss', loss.item(), self.wrt_step)  # 记录损失到TensorBoard

        # FOR EVAL
        seg_metrics = eval_metrics(output, target, self.num_classes)  # 计算分割度量指标
        self._update_seg_metrics(*seg_metrics)  # 更新分割度量指标
        pixAcc, mIoU, _ = self._get_seg_metrics().values()  # 获取分割指标值

        # PRINT INFO
        tbar.set_description('TRAIN ({}) | Loss: {:.3f} | Acc {:.2f} mIoU {:.2f} | B {:.2f} D {:.2f} |'.format(
                                            epoch, self.total_loss.average, 
                                            pixAcc, mIoU,
                                            self.batch_time.average, self.data_time.average))  # 打印训练信息

    # METRICS TO TENSORBOARD
    seg_metrics = self._get_seg_metrics()
    for k, v in list(seg_metrics.items())[:-1]:  # 遍历分割度量指标并记录
        self.writer.add_scalar(f'{
      
      self.wrt_mode}/{
      
      k}', v, self.wrt_step)
    for i, opt_group in enumerate(self.optimizer.param_groups):  # 记录学习率
        self.writer.add_scalar(f'{
      
      self.wrt_mode}/Learning_rate_{
      
      i}', opt_group['lr'], self.wrt_step)

    # RETURN LOSS & METRICS
    log = {
    
    'loss': self.total_loss.average,  # 返回平均损失
            **seg_metrics}  # 返回分割度量指标

    return log  # 返回日志信息

交叉验证部分，我们进行以下的定义

def _valid_epoch(self, epoch):
    if self.val_loader is None:
        self.logger.warning('Not data loader was passed for the validation step, No validation is performed !')
        return {
    
    }  # 如果没有提供验证数据加载器，发出警告并返回一个空字典
    self.logger.info('\n###### EVALUATION ######')

    self.model.eval()  # 设置模型为评估（验证）模式
    self.wrt_mode = 'val'  # 设置写入模式为'val'（用于TensorBoard记录）

    self._reset_metrics()  # 重置度量指标
    tbar = tqdm(self.val_loader, ncols=130)  # 创建一个进度条用于遍历验证数据集
    with torch.no_grad():  # 禁用梯度计算
        val_visual = []  # 用于可视化的图像列表
        for batch_idx, (data, target) in enumerate(tbar):
            #data, target = data.to(self.device), target.to(self.device)  # 将数据和目标移到指定的设备上（通常是GPU）
            # LOSS
            output = self.model(data)  # 前向传播，获取模型的输出
            loss = self.loss(output, target)  # 计算损失
            if isinstance(self.loss, torch.nn.DataParallel):  # 如果损失函数是DataParallel损失函数
                loss = loss.mean()  # 计算损失的均值
            self.total_loss.update(loss.item())  # 更新总损失

            seg_metrics = eval_metrics(output, target, self.num_classes)  # 计算分割度量指标
            self._update_seg_metrics(*seg_metrics)  # 更新分割度量指标

            # LIST OF IMAGE TO VIZ (15 images)
            if len(val_visual) < 15:  # 用于可视化的图像数量限制在15张以内
                target_np = target.data.cpu().numpy()  # 将目标从GPU移到CPU并转换为NumPy数组
                output_np = output.data.max(1)[1].cpu().numpy()  # 将模型输出的类别概率最大的类别作为预测结果
                val_visual.append([data[0].data.cpu(), target_np[0], output_np[0]])  # 添加可视化所需的图像和标签

            # PRINT INFO
            pixAcc, mIoU, _ = self._get_seg_metrics().values()  # 获取分割度量指标的值
            tbar.set_description('EVAL ({}) | Loss: {:.3f}, PixelAcc: {:.2f}, Mean IoU: {:.2f} |'.format( epoch,
                                            self.total_loss.average,
                                            pixAcc, mIoU))  # 打印验证信息

        # WRITING & VISUALIZING THE MASKS
        val_img = []  # 用于可视化的图像列表
        palette = self.train_loader.dataset.palette  # 获取调色板信息
        for d, t, o in val_visual:  # 遍历可视化图像列表
            d = self.restore_transform(d)  # 还原图像的转换（例如，去均值、缩放等）
            t, o = colorize_mask(t, palette), colorize_mask(o, palette)  # 将标签和模型输出转换为彩色掩码
            d, t, o = d.convert('RGB'), t.convert('RGB'), o.convert('RGB')  # 将图像转换为RGB格式
            [d, t, o] = [self.viz_transform(x) for x in [d, t, o]]  # 应用可视化转换
            val_img.extend([d, t, o])  # 添加可视化图像到列表中
        val_img = torch.stack(val_img, 0)  # 将可视化图像堆叠成一个张量
        val_img = make_grid(val_img.cpu(), nrow=3, padding=5)  # 使用Grid方式排列可视化图像
        self.writer.add_image(f'{
      
      self.wrt_mode}/inputs_targets_predictions', val_img, self.wrt_step)  # 将可视化图像写入TensorBoard

        # METRICS TO TENSORBOARD
        self.wrt_step = (epoch) * len(self.val_loader)  # 计算当前步数
        self.writer.add_scalar(f'{
      
      self.wrt_mode}/loss', self.total_loss.average, self.wrt_step)  # 记录平均损失到TensorBoard
        seg_metrics = self._get_seg_metrics()  # 获取分割度量指标
        for k, v in list(seg_metrics.items())[:-1]:  # 遍历分割度量指标并记录到TensorBoard
            self.writer.add_scalar(f'{
      
      self.wrt_mode}/{
      
      k}', v, self.wrt_step)

        log = {
    
    
            'val_loss': self.total_loss.average,  # 返回平均验证损失
            **seg_metrics  # 返回分割度量指标
        }

    return log  # 返回日志信息

# 以下是用于度量指标的辅助函数
def _reset_metrics(self):
    self.batch_time = AverageMeter()  # 用于记录批次处理时间的平均值
    self.data_time = AverageMeter()  # 用于记录数据加载时间的平均值
    self.total_loss = AverageMeter()  # 用于记录总损失的平均值
    self.total_inter, self.total_union = 0, 0  # 用于记录交集和并集的总和
    self.total_correct, self.total_label = 0, 0  # 用于记录正确分类和标签的总和

def _update_seg_metrics(self, correct, labeled, inter, union):
    self.total_correct += correct  # 更新正确分类的数量
    self.total_label += labeled  # 更新标签的数量
    self.total_inter += inter  # 更新交集的总和
    self.total_union += union  # 更新并集的总和

def _get_seg_metrics(self):
    pixAcc = 1.0 * self.total_correct / (np.spacing(1) + self.total_label)  # 计算像素准确率
    IoU = 1.0 * self.total_inter / (np.spacing(1) + self.total_union)  # 计算各类别的IoU
    mIoU = IoU.mean()  # 计算平均IoU
    return {
    
    
        "Pixel_Accuracy": np.round(pixAcc, 3),  # 返回像素准确率
        "Mean_IoU": np.round(mIoU, 3),  # 返回平均IoU
        "Class_IoU": dict(zip(range(self.num_classes), np.round(IoU, 3)))  # 返回各类别的IoU
    }

3. Метрики

Здесь мы используем две метрики для оценки производительности модели:

Pixel_Accuracy — это метрика, используемая для оценки производительности задач сегментации изображений. Она используется для измерения доли количества пикселей, которые модель правильно классифицировала по всему изображению, к общему количеству пикселей. Это можно выразить просто в виде следующей математической формулы:

в:

• «Количество правильно классифицированных пикселей» указывает количество пикселей, которые модель правильно классифицировала в сегментированном изображении.
• «Общее количество пикселей» представляет общее количество пикселей во всем сегментированном изображении.

Точность пикселей находится в диапазоне от 0 до 1, где 1 означает, что модель правильно классифицировала все пиксели по всему изображению, а 0 означает, что модель не классифицировала правильно ни одного пикселя.

Mean_IoU : IoU (пересечение через объединение) — это индикатор, который указывает степень перекрытия между двумя наборами и обычно используется в задачах сегментации. В задаче сегментации один набор представляет прогнозируемые моделью области сегментации, а другой набор представляет истинные области сегментации. Формула расчета IoU выглядит следующим образом:

в:

• «Область пересечения» — это область пересечения прогнозируемой модели области сегментации и реальной области сегментации.
• «Область объединения» — это площадь объединения прогнозируемой модели области сегментации и реальной области сегментации.

или

где:

• TP (True Positives): указывает количество пикселей, которые модель правильно спрогнозировала как положительные (целевая категория).
• FP (ложные срабатывания): указывает количество пикселей, в которых модель неправильно предсказывает фоновые пиксели как положительные.
• FN (ложные негативы): указывает количество пикселей, в которых модель неправильно прогнозирует положительные пиксели в качестве фона.

Среднее пересечение по объединению (Mean IoU) — это среднее значение всех категорий IoU.

Среди них N — количество категорий, IoU_i — IoU i-й категории.

4. Анализ результатов

Основные настройки параметров следующие:

"epochs": 80,
"loss": "CrossEntropyLoss2d",
"batch_size": 8,
"base_size": 400,  //图像大小调整为base_size，然后随机裁剪
"crop_size": 380,  //重新缩放后随机裁剪的大小
 "optimizer": {
    
    
        "type": "SGD",
        "differential_lr": true,
        "args":{
    
    
            "lr": 0.01,
            "weight_decay": 1e-4,
            "momentum": 0.9
        }
    },

Из-за ограниченных ресурсов графического процессора здесь запускаются только эпохи 80. Полученная информация журнала выглядит следующим образом: Информация,

записанная Tensorboard: сравнение входных данных, основной истинности и выходных данных набора перекрестной проверки проверки поезда
: некоторые детали не полностью сегментированы, но изображение можно распознать по основному телу. Между train_loss и val_loss нет большой разницы, и нет переобучения. Увеличение периода обучения может привести к лучшим результатам.

5. Тест изображения

Код тестовой части следующий:

    args = parse_arguments()  # 解析命令行参数
    config = json.load(open(args.config))  # 从JSON文件中加载配置信息

    # 根据配置信息创建数据加载器
    loader = getattr(dataloaders, config['train_loader']['type'])(**config['train_loader']['args'])
    to_tensor = transforms.ToTensor()  # 创建图像到张量的转换
    normalize = transforms.Normalize(loader.MEAN, loader.STD)  # 创建归一化转换
    num_classes = loader.dataset.num_classes  # 获取数据集中的类别数量
    palette = loader.dataset.palette  # 获取颜色映射表

    # 创建模型
    model = getattr(models, config['arch']['type'])(num_classes, **config['arch']['args'])  # 根据配置创建模型
    availble_gpus = list(range(torch.cuda.device_count()))  # 获取可用的GPU列表
    device = torch.device('cuda:0' if len(availble_gpus) > 0 else 'cpu')  # 选择运行设备（GPU或CPU）

    # 加载模型检查点
    checkpoint = torch.load(args.model, map_location=device)
    if isinstance(checkpoint, dict) and 'state_dict' in checkpoint.keys():
        checkpoint = checkpoint['state_dict']
    
    # 如果在训练期间使用了数据并行，需要处理模型
    if 'module' in list(checkpoint.keys())[0] and not isinstance(model, torch.nn.DataParallel):
        # 对于GPU推理，使用数据并行
        if "cuda" in device.type:
            model = torch.nn.DataParallel(model)
        else:
            # 对于CPU推理，移除模型的"module"前缀
            new_state_dict = OrderedDict()
            for k, v in checkpoint.items():
                name = k[7:]
                new_state_dict[name] = v
            checkpoint = new_state_dict
    
    # 加载模型权重
    model.load_state_dict(checkpoint)
    model.to(device)  # 将模型移动到所选设备
    model.eval()  # 设置模型为评估模式

    # 创建输出目录
    if not os.path.exists('outputs'):
        os.makedirs('outputs')

    # 获取图像文件列表
    image_files = sorted(glob(os.path.join(args.images, f'*.{
      
      args.extension}')))
    with torch.no_grad():
        tbar = tqdm(image_files, ncols=100)  # 创建进度条
        for img_file in tbar:
            image = Image.open(img_file).convert('RGB')  # 打开图像并将其转换为RGB格式
            input = normalize(to_tensor(image)).unsqueeze(0)  # 转换图像并添加批次维度

			#预测图像分割结果
			prediction = multi_scale_predict(model, input, scales, num_classes, device)
         
            prediction = F.softmax(torch.from_numpy(prediction), dim=0).argmax(0).cpu().numpy()  # 计算最终的预测结果
            save_images(image, prediction, args.output, img_file, palette)  # 保存预测结果的图像

Функция прогнозирования многомасштабного изображения определяется как:

def multi_scale_predict(model, image, scales, num_classes, device, flip=False):
    # 获取输入图像的尺寸
    input_size = (image.size(2), image.size(3))
    # 创建上采样层，用于将不同尺度的预测结果恢复到原始尺寸
    upsample = nn.Upsample(size=input_size, mode='bilinear', align_corners=True)
    # 初始化用于累计预测结果的数组
    total_predictions = np.zeros((num_classes, image.size(2), image.size(3)))

    # 将输入图像转换为NumPy数组，并移动到CPU上
    image = image.data.data.cpu().numpy()
    
    # 遍历不同的尺度
    for scale in scales:
        # 缩放图像
        scaled_img = ndimage.zoom(image, (1.0, 1.0, float(scale), float(scale)), order=1, prefilter=False)
        # 将缩放后的图像转换为PyTorch张量并移动到指定设备
        scaled_img = torch.from_numpy(scaled_img).to(device)
        # 使用模型进行预测并上采样到原始尺寸
        scaled_prediction = upsample(model(scaled_img).cpu())

        # 如果启用了翻转，对翻转后的图像进行预测并平均
        if flip:
            fliped_img = scaled_img.flip(-1).to(device)
            fliped_predictions = upsample(model(fliped_img).cpu())
            scaled_prediction = 0.5 * (fliped_predictions.flip(-1) + scaled_prediction)
        
        # 将当前尺度的预测结果累加到总体预测中
        total_predictions += scaled_prediction.data.cpu().numpy().squeeze(0)

    # 计算平均预测结果
    total_predictions /= len(scales)
    return total_predictions

Мы произвольно указываем входное изображение и тестируем эффект сегментации модели.Эффект

не очень хороший.Если у вас есть ресурсы, вы можете увеличить количество эпохального обучения и попробовать разные наборы данных
—————————— —————— ——————————————————————————