PyTorch深度学习实战(26)——多对象实例分割

news/2024/2/21 10:29:02

PyTorch深度学习实战(26)——多对象实例分割

    • 0. 前言
    • 1. 获取并准备数据
    • 2. 使用 Detectron2 训练实例分割模型
    • 3. 对新图像进行推断
    • 小结
    • 系列链接

0. 前言

我们已经学习了多种图像分割算法,在本节中,我们将学习如何使用 Detectron2 平台以及 Google 开放图像数据集执行多对象实例分割任务。Detectron2Facebook 团队打造的人工智能框架,其中包括了高性能的对象检测算法实现,包括 Mask R-CNN 模型等。Detectron2 支持一系列与目标检测和人体姿态估计等相关的任务,此外,Detectron2 还增加了对语义分割和全景分割的支持。通过利用 Detectron2,我们能够通过使用少量代码构建目标检测、分割和姿态估计模型。在本节中,我们将介绍如何从 Google 开放图像数据集中获取数据,将数据转换为 Detectron2 可接受的 COCO 格式,并训练模型执行实例分割,最后,介绍如何使用训练后的模型对新图像进行推理。

1. 获取并准备数据

Google 开放图像数据集 (Google Open Images) 是由 Google 开发和维护的一个大规模图像数据集,用于计算机视觉领域的研究和开发。该数据集包含了数百万张标记和分类的图像,涵盖了广泛的主题和场景。在本节中,我们使用 Google 开放图像数据集,为了减少训练的时间,我们仅仅使用其中一部分数据集,而不使用整个数据集。

(1) 安装所需的库:

pip install openimages

(2) 下载所需的标注文件,包括分割标注文件 train-annotations-object-segmentation.csv 和类别文件 dict.csv。

(3) 指定希望模型预测的类别 required_classes

import pandas as pdfrom glob import glob
import os
import numpy as np
import cv2required_classes = 'person,dog,bird,car,elephant,football,jug,laptop,Mushroom,Pizza,Rocket,Shirt,Traffic sign,Watermelon,Zebra'
required_classes = [c.lower() for c in required_classes.lower().split(',')]classes = pd.read_csv('dict.csv', header=None)
classes.columns = ['class','class_name']
classes = classes[classes['class_name'].map(lambda x: x in required_classes)]

(4) 获取 required_classes 对应的图像 ID 和掩码:

df = pd.read_csv('train-annotations-object-segmentation.csv')data = pd.merge(df, classes, left_on='LabelName', right_on='class')subset_data = data.groupby('class_name').agg({'ImageID': lambda x: list(x)[:2000]})
subset_data = ([y for x in subset_data.ImageID.tolist() for y in x])
subset_data = data[data['ImageID'].map(lambda x: x in subset_data)]
subset_masks = subset_data['MaskPath'].tolist()

考虑到海量数据集会增加训练成本,在 subset_data 中每个类别只使用 2000 张图像,我们也可以修改每个类别所用数据样本数量。使用以上代码,可以获取图像对应的 ImageIdMaskPath 值。接下来,我们继续从 Google 开放图像数据集中下载实际图像和蒙版。

(5) 下载掩码数据子集,Google 开放图像数据集有 16 个用于掩码 ZIP 文件,包括 train-masks-0.zip、train-masks-1.zip、train-masks-2.zip、train-masks-3.zip、train-masks-4.zip、train-masks-5.zip、train-masks-6.zip、train-masks-7.zip、train-masks-8.zip、train-masks-9.zip、train-masks-a.zip、train-masks-b.zip、train-masks-c.zip、train-masks-d.zip、train-masks-e.zip 和 train-masks-f.zip,下载以上 ZIP 文件。每个 ZIP 文件只包含 subset_masks 中的一些掩码,因此下载之后我们只读取所需掩码文件:

if not os.path.exists('masks'):os.mkdir('masks')def fname(file):return file.split('/')[-1]for c in '0123456789abcdef':tmp_mask = 'train-masks-{}/*.png'.format(c)tmp_masks = glob(tmp_mask)items = [(m, fname(m)) for m in tmp_masks]items = [(i,j) for (i,j) in items if j in subset_masks]for i,j in items:os.rename(i, f'masks/{j}')

(6) 下载 ImageId 对应的图片:

masks = glob('masks/*.png')masks = [fname(mask) for mask in masks]subset_data = subset_data[subset_data['MaskPath'].map(lambda x: x in masks)]
subset_imageIds = subset_data['ImageID'].tolist()from openimages.download import _download_images_by_idif not os.path.exists('images'):os.makedirs('images')_download_images_by_id(subset_imageIds, 'train', './images/')

(7) 压缩所有图像、掩码和标注文件并保存,用于以后模型训练:

import zipfile
files = glob('images/*.jpg') + glob('masks/*.png') + ['train-annotations-object-segmentation.csv', 'dict.csv']
with zipfile.ZipFile('data.zip','w') as zipme:for file in files:zipme.write(file, compress_type=zipfile.ZIP_DEFLATED)if not os.path.exists('train'):os.makedirs('train')
os.rename('images', 'train/myData')
os.rename('masks', 'train/annotations')

由于目标检测数据集有多种不同标注格式,为了标准化,Detectron 需要使用一种严格的训练数据格式。虽然可以使用自定义数据集格式并将其提供给 Detectron,但是将训练数据保存为 COCO 格式会使训练更加容易。这样,我们就在不改变训练数据的情况下利用其他训练算法(如 Detectron Transformers (DETR))。

(8)COCO 格式定义所需类别:

import datetimeINFO = {"description": "MyData2023","url": "None","version": "1.0","year": 2023,"contributor": "sizhky","date_created": datetime.datetime.utcnow().isoformat(' ')
}LICENSES = [{"id": 1,"name": "MIT"}
]CATEGORIES = [{'id': id+1, 'name': name.replace('/',''), 'supercategory': 'none'} for id,(_,(name, clss_name)) in enumerate(classes.iterrows())]

在以上代码中,在 CATEGORIES 的定义中,创建了一个名为 supercategory 的新键。例如: ManWoman 类别都属于 Person 这个超类别。在本节中,由于我们对超类别不感兴趣,将其指定为 none

导入相关库,并创建一个空字典,其中包含保存 COCO JSON 文件所需的键值:

from pycococreatortools import pycococreatortools
from os import listdir
from os.path import isfile, join
from PIL import Imagecoco_output = {"info": INFO,"licenses": LICENSES,"categories": CATEGORIES,"images": [],"annotations": []
}

定义包含图像位置和标注文件位置信息的变量:

ROOT_DIR = "train"
IMAGE_DIR, ANNOTATION_DIR = 'train/myData/', 'train/annotations/'
image_files = [f for f in listdir(IMAGE_DIR) if isfile(join(IMAGE_DIR, f))]
annotation_files = [f for f in listdir(ANNOTATION_DIR) if isfile(join(ANNOTATION_DIR, f))]

循环遍历每个图像文件名,并将信息填充到 coco_output 字典中的 images 键中:

image_id = 1for image_filename in image_files:image = Image.open(IMAGE_DIR + '/' + image_filename)image_info = pycococreatortools.create_image_info(image_id, os.path.basename(image_filename), image.size)coco_output["images"].append(image_info)image_id = image_id + 1

(9) 循环遍历每个分段标注,并将信息填充到 coco_output 字典中的 annotations 键中:

def stem(filename):return filename.split('.')[0]segmentation_id = 1
for annotation_filename in annotation_files:image_id = [f for f in coco_output['images'] if stem(f['file_name'])==annotation_filename.split('_')[0]][0]['id']class_id = [x['id'] for x in CATEGORIES if x['name'] in annotation_filename][0]category_info = {'id': class_id, 'is_crowd': 'crowd' in image_filename}binary_mask = np.asarray(Image.open(f'{ANNOTATION_DIR}/{annotation_filename}').convert('1')).astype(np.uint8)annotation_info = pycococreatortools.create_annotation_info(segmentation_id, image_id, category_info, binary_mask, image.size, tolerance=2)if annotation_info is not None:coco_output["annotations"].append(annotation_info)segmentation_id = segmentation_id + 1

(10)coco_output 保存在 JSON 文件中:

coco_output['categories'] = [{'id': id+1, 'name': clss_name, 'supercategory': 'none'} for id,(_,(name, clss_name)) in enumerate(classes.iterrows())]import json
with open('images.json', 'w') as output_json_file:json.dump(coco_output, output_json_file)

获取 COCO 格式的文件后,就可以很容易地使用 Detectron2 框架来训练模型。

2. 使用 Detectron2 训练实例分割模型

(1) 安装 Detectron2,在安装正确的库之前,首先应该检查 CUDAPyTorch 版本,以 PyTorch 1.9CUDA 11.3 为例,使用相应的文件:

$ pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu113/torch1.9/index.html
$ pip install pyyaml pycocotools

(2) 导入相关 detectron2 库:

from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2.engine import DefaultTrainer

重新获取所需的类:

required_classes = 'person,dog,bird,car,elephant,football,jug,laptop,Mushroom,Pizza,Rocket,Shirt,Traffic sign,Watermelon,Zebra'
required_classes = [c.lower() for c in required_classes.lower().split(',')]cclasses = pd.read_csv('dict.csv', header=None)
classes.columns = ['class','class_name']
classes = classes[classes['class_name'].map(lambda x: x in required_classes)]

(3) 使用 register_coco_instances 注册创建的数据集:

from detectron2.data.datasets import register_coco_instances
register_coco_instances("dataset_train", {}, "images.json", "train/myData")

register_coco_instances 函数接受四个参数,第一个参数是数据集的名称,用于唯一标识该数据集;第二个参数是一个空字典,通常用于存储一些额外的元数据信息;第三个参数是标注文件的路径;最后一个参数是图像文件夹的路径。可以通过类似方式注册任意数量的数据集。

(4) 定义 Detectron2 配置文件 cfg 中的参数。

配置文件 (cfg) 是一个特殊的 Detectron 对象,它包含用于训练模型的所有相关信息:

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("dataset_train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml") # pretrained weights
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025 # pick a good LR
cfg.SOLVER.MAX_ITER = 8000 # instead of epochs, we train on 5000 batches
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512
cfg.MODEL.ROI_HEADS.NUM_CLASSES = len(classes) 

使用配置文件可以设置训练模型所需的所有主要超参数,merge_from_file 方法用于从配置文件导入核心参数,mask_rcnn_R_50_FPN_3x.yaml 配置文件使用以 FPN 作为主干网络进行预训练的 mask_rcnn,此配置文件还包含有关预训练模型的其他信息,例如优化器和损失函数。

(5) 训练模型:

os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg) 
trainer.resume_or_load(resume=True)
trainer.train()

训练完成后,可以使用该模型预测类别、边界框以及自定义数据集中定义的类别对象的分割掩码。

模型训练
将模型保存在指定文件夹中:

import shutil
import random
from matplotlib import pyplot as plt
shutil.copy('output/model_final.pth', 'output/trained_model.pth')

模型训练完成后,在下一小节中,我们将对新图像进行推断。

3. 对新图像进行推断

为了对新图像进行推断,我们首先加载模型,设置概率阈值,并将其传递给 DefaultPredictor 方法。

(1) 加载训练好的模型权重,使用相同的 cfg 配置文件并加载模型权重:

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "trained_model.pth")

(2) 设置对象属于某一类别的概率阈值:

cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.25

(3) 定义 predictor 方法:

predictor = DefaultPredictor(cfg)

(4) 对新图像进行分割并将其可视化。

加载图像文件,随机绘制 40 张图像分割结果:

from detectron2.utils.visualizer import ColorMode
files = glob('train/myData/*.jpg')
for _ in range(40):im = cv2.imread(random.choice(files))outputs = predictor(im)v = Visualizer(im[:, :, ::-1],scale=0.5,metadata=MetadataCatalog.get("dataset_train"),instance_mode=ColorMode.IMAGE_BW )out = v.draw_instance_predictions(outputs["instances"].to("cpu"))plt.imshow(out.get_image())plt.show()

VisualizerDetectron2 绘制对象实例的方式,模型预测输出是张量字典,Visualizer 将它们转换为像素信息并将绘制在图像上。Visualizer 接受以下参数作为输入:

  • im:要可视化的图像
  • scale:绘制图像的大小,在以上代码中,将图像缩小 50%
  • metadata:数据集的类别信息,主要是索引到类别的映射,用于将原始张量输入解码为人类的可读实际类别
  • instance_mode:指定模型突出显示的分割像素

创建了 Visualizer 类对象后,就可以用其绘制模型的实例预测并显示图像,代码输出结果如下所示:

模型预测结果

从以上输出可以看出,模型能够准确地识别出人物对应的像素。

小结

Detectron2 是用于计算机视觉任务的深度学习库,构建与 PyTorch 框架之上,并提供了一套灵活且高性能的工具,用于对象检测、实例分割和人体关键点检测等任务。其使用了先进的技术和优化策略,如特征金字塔网络、RoI Align 操作等,以提高模型的感受野和对小目标的检测能力。本节中,介绍了如何安装和使用 Detectron2 库,并通过使用 COCO 数据格式提高模型训练效率。

系列链接

PyTorch深度学习实战(1)——神经网络与模型训练过程详解
PyTorch深度学习实战(2)——PyTorch基础
PyTorch深度学习实战(3)——使用PyTorch构建神经网络
PyTorch深度学习实战(4)——常用激活函数和损失函数详解
PyTorch深度学习实战(5)——计算机视觉基础
PyTorch深度学习实战(6)——神经网络性能优化技术
PyTorch深度学习实战(7)——批大小对神经网络训练的影响
PyTorch深度学习实战(8)——批归一化
PyTorch深度学习实战(9)——学习率优化
PyTorch深度学习实战(10)——过拟合及其解决方法
PyTorch深度学习实战(11)——卷积神经网络
PyTorch深度学习实战(12)——数据增强
PyTorch深度学习实战(13)——可视化神经网络中间层输出
PyTorch深度学习实战(14)——类激活图
PyTorch深度学习实战(15)——迁移学习
PyTorch深度学习实战(16)——面部关键点检测
PyTorch深度学习实战(17)——多任务学习
PyTorch深度学习实战(18)——目标检测基础
PyTorch深度学习实战(19)——从零开始实现R-CNN目标检测
PyTorch深度学习实战(20)——从零开始实现Fast R-CNN目标检测
PyTorch深度学习实战(21)——从零开始实现Faster R-CNN目标检测
PyTorch深度学习实战(22)——从零开始实现YOLO目标检测
PyTorch深度学习实战(23)——使用U-Net架构进行图像分割
PyTorch深度学习实战(24)——从零开始实现Mask R-CNN实例分割
PyTorch深度学习实战(25)——自编码器(Autoencoder)
PyTorch深度学习实战(26)——卷积自编码器(Convolutional Autoencoder)
PyTorch深度学习实战(27)——变分自编码器(Variational Autoencoder, VAE)
PyTorch深度学习实战(28)——对抗攻击(Adversarial Attack)
PyTorch深度学习实战(29)——神经风格迁移
PyTorch深度学习实战(30)——Deepfakes
PyTorch深度学习实战(31)——生成对抗网络(Generative Adversarial Network, GAN)
PyTorch深度学习实战(32)——DCGAN详解与实现
PyTorch深度学习实战(33)——条件生成对抗网络(Conditional Generative Adversarial Network, CGAN)
PyTorch深度学习实战(34)——Pix2Pix详解与实现


https://www.xjx100.cn/news/3271247.html

相关文章

代码随想录day20--二叉树的应用8

LeetCode669.修剪二叉搜索树 题目描述: 给你二叉搜索树的根节点 root ,同时给定最小边界low 和最大边界 high。通过修剪二叉搜索树,使得所有节点的值在[low, high]中。修剪树 不应该 改变保留在树中的元素的相对结构 (即,如果没…

【Java程序设计】【C00257】基于Springboot的校园二手书交易平台(有论文)

基于Springboot的校园二手书交易平台(有论文) 项目简介项目获取开发环境项目技术运行截图 项目简介 这是一个基于Springboot的乐校园二手书交易管理系统 本系统分为系统功能模块、管理员功能模块、卖家用户功能模块以及用户功能模块。 系统功能模块&…

应急响应实战笔记02日志分析篇(2)

第2篇:Linux日志分析 0x00 前言 Linux系统拥有非常灵活和强大的日志功能,可以保存几乎所有的操作记录,并可以从中检索出我们需要的信息。 本文简介一下Linux系统日志及日志分析技巧。 0x01 日志简介 日志默认存放位置:/var/log/ 查看日志…

算法学习——LeetCode力扣二叉树篇5

算法学习——LeetCode力扣二叉树篇5 513. 找树左下角的值 513. 找树左下角的值 - 力扣(LeetCode) 描述 给定一个二叉树的 根节点 root,请找出该二叉树的 最底层 最左边 节点的值。 假设二叉树中至少有一个节点。 示例 示例 1: 输入: r…

C# CAD交互界面-自定义面板集-查找定位(六)

运行环境 vs2022 c# cad2016 调试成功 一、代码说明 1. 类成员变量声明&#xff1a; List<ObjectId> objectIds new List<ObjectId>(); // 用于存储AutoCAD实体对象的ObjectId列表 private static Autodesk.AutoCAD.Windows.PaletteSet _ps2; // 自定义浮动面板…

2024年教师资格证认定报名完整流程

&#x1f49a;网上报名流程概览 一、进入教资认定网报入口&#xff1b; 二、进行实名核验&#xff1b; 三、申请网报时间查询&#xff1b; 四、个人信息维护&#xff1b; 五、认定申请报名&#xff1b; &#x1f49a;认定所需材料 1、 身份证&#xff1b; 2 、户口本/居住证…

nba2k24 韩旭面补

nba2k23-24 韩旭面补 nba2k23-nba2k24通用 韩旭面补 下载地址&#xff1a; https://www.changyouzuhao.cn/9605.html

使用Docker快速部署MySQL

部署MySQL 使用Docker安装&#xff0c;仅仅需要一步即可&#xff0c;在命令行输入下面的命令 docker run -d \--name mysql \-p 3306:3306 \-e TZAsia/Shanghai \-e MYSQL_ROOT_PASSWORD123456 \mysql MySQL安装完毕&#xff01;通过任意客户端工具即可连接到MySQL. 当我们执…