自定义数据集 - Dataset

这篇具有很好参考价值的文章主要介绍了自定义数据集 - Dataset。希望对大家有所帮助。如果存在错误或未考虑完全的地方,请大家不吝赐教,您也可以点击"举报违法"按钮提交疑问。

1. PASCAL VOC格式 划分训练集和验证集

import os
import random

def main():
    random.seed(0)  # 设置随机种子,保证随机结果可复现

    files_path = "./VOCdevkit/VOC2012/Annotations"  # 指定annotations目录
    assert os.path.exists(files_path), "path: '{}' does not exist.".format(files_path)

    val_rate = 0.5  # 定义划分比例

    files_name = sorted([file.split(".")[0] for file in os.listdir(files_path)])
    files_num = len(files_name)
    val_index = random.sample(range(0, files_num), k=int(files_num*val_rate))
    train_files = []
    val_files = []
    for index, file_name in enumerate(files_name):
        if index in val_index:
            val_files.append(file_name)
        else:
            train_files.append(file_name)

    try:
        train_f = open("train.txt", "x")
        eval_f = open("val.txt", "x")
        train_f.write("\n".join(train_files))
        eval_f.write("\n".join(val_files))
    except FileExistsError as e:
        print(e)
        exit(1)

if __name__ == '__main__':
    main()

2. 自定义dataset

自定义自己的数据集,需要继承torch.utils.data.Dataset,并且实现__len____getitem__方法。

  • __len__:获取数据集的大小
  • __getitem__:返回数据信息

如果使用多GPU训练,还需要实现get_height_and_width方法,获取图像的高度和宽度,如果不实现这个方法,它就会载入所有的图片,去计算图片的高和宽,这样的话就比较耗时和占内存。

自定义数据集 - Dataset,AI,pytorch,python,深度学习

from torch.utils.data import Dataset
import os
import torch
import json
from PIL import Image
from lxml import etree

# 继承torch.utils.data.Dataset
class VOCDataSet(Dataset):
    """读取解析PASCAL VOC2007/2012数据集"""
    """
    voc_root:数据集所在根目录
    year:数据集年份
    transforms:数据预处理方法
    text_name:读取txt文件名称 train.txt/val.txt
    """
    def __init__(self, voc_root, year="2012", transforms=None, txt_name: str = "train.txt"):
        assert year in ["2007", "2012"], "year must be in ['2007', '2012']"
        # 增加容错能力
        if "VOCdevkit" in voc_root:
            self.root = os.path.join(voc_root, f"VOC{year}")
        else:
            self.root = os.path.join(voc_root, "VOCdevkit", f"VOC{year}")
        self.img_root = os.path.join(self.root, "JPEGImages")
        self.annotations_root = os.path.join(self.root, "Annotations")

        # 读取 train.txt 或 val.txt文件
        txt_path = os.path.join(self.root, "ImageSets", "Main", txt_name)
        assert os.path.exists(txt_path), "not found {} file.".format(txt_name)

        with open(txt_path) as read:
            # strip() 方法用于移除字符串头尾指定的字符
            xml_list = [os.path.join(self.annotations_root, line.strip() + ".xml")
                        for line in read.readlines() if len(line.strip()) > 0]

        self.xml_list = []
        # 确认文件是否存在
        for xml_path in xml_list:
            if os.path.exists(xml_path) is False:
                print(f"Warning: not found '{xml_path}', skip this annotation file.")
                continue

            # 排除图像中没有目标的数据
            with open(xml_path) as fid:
                xml_str = fid.read()
            xml = etree.fromstring(xml_str)
            data = self.parse_xml_to_dict(xml)["annotation"]
            if "object" not in data:
                print(f"INFO: no objects in {xml_path}, skip this annotation file.")
                continue

            self.xml_list.append(xml_path)

        assert len(self.xml_list) > 0, "in '{}' file does not find any information.".format(txt_path)

        # 读取数据集类别信息
        json_file = './pascal_voc_classes.json'
        assert os.path.exists(json_file), "{} file not exist.".format(json_file)
        with open(json_file, 'r') as f:
            self.class_dict = json.load(f)

        self.transforms = transforms

    # 获取数据集大小
    def __len__(self):
        return len(self.xml_list)

    # 根据传入的索引值获取数据信息
    def __getitem__(self, idx):
        # 读取xml文件
        xml_path = self.xml_list[idx]
        with open(xml_path) as fid:
            xml_str = fid.read()
        xml = etree.fromstring(xml_str)
        data = self.parse_xml_to_dict(xml)["annotation"]
        img_path = os.path.join(self.img_root, data["filename"])
        image = Image.open(img_path)
        if image.format != "JPEG":
            raise ValueError("Image '{}' format not JPEG".format(img_path))

        boxes = []
        labels = []
        iscrowd = []
        assert "object" in data, "{} lack of object information.".format(xml_path)
        for obj in data["object"]:
            xmin = float(obj["bndbox"]["xmin"])
            xmax = float(obj["bndbox"]["xmax"])
            ymin = float(obj["bndbox"]["ymin"])
            ymax = float(obj["bndbox"]["ymax"])

            # 进一步检查数据,有的标注信息中可能有w或h为0的情况,这样的数据会导致计算回归loss为nan
            if xmax <= xmin or ymax <= ymin:
                print("Warning: in '{}' xml, there are some bbox w/h <=0".format(xml_path))
                continue
            
            boxes.append([xmin, ymin, xmax, ymax])
            labels.append(self.class_dict[obj["name"]])
            if "difficult" in obj:
                iscrowd.append(int(obj["difficult"]))
            else:
                iscrowd.append(0)

        # 将数据转为tensor类型
        boxes = torch.as_tensor(boxes, dtype=torch.float32)
        labels = torch.as_tensor(labels, dtype=torch.int64)
        iscrowd = torch.as_tensor(iscrowd, dtype=torch.int64)
        image_id = torch.tensor([idx])
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])

        target = {}
        target["boxes"] = boxes
        target["labels"] = labels
        target["image_id"] = image_id
        target["area"] = area
        target["iscrowd"] = iscrowd

        if self.transforms is not None:
            image, target = self.transforms(image, target)

        return image, target

    # 获取图像高度和宽度
    def get_height_and_width(self, idx):
        # read xml
        xml_path = self.xml_list[idx]
        with open(xml_path) as fid:
            xml_str = fid.read()
        xml = etree.fromstring(xml_str)
        data = self.parse_xml_to_dict(xml)["annotation"]
        data_height = int(data["size"]["height"])
        data_width = int(data["size"]["width"])
        return data_height, data_width

    """
    将xml文件解析成字典形式
    """
    def parse_xml_to_dict(self, xml):
        if len(xml) == 0:  # 遍历到底层,直接返回tag对应的信息
            return {xml.tag: xml.text}

        result = {}
        for child in xml:
            child_result = self.parse_xml_to_dict(child)  # 递归遍历标签信息
            if child.tag != 'object':
                result[child.tag] = child_result[child.tag]
            else:
                if child.tag not in result:  # 因为object可能有多个,所以需要放入列表里
                    result[child.tag] = []
                result[child.tag].append(child_result[child.tag])
        return {xml.tag: result}

    def coco_index(self, idx):
        """
        该方法是专门为pycocotools统计标签信息准备,不对图像和标签作任何处理
        由于不用去读取图片,可大幅缩减统计时间

        Args:
            idx: 输入需要获取图像的索引
        """
        # read xml
        xml_path = self.xml_list[idx]
        with open(xml_path) as fid:
            xml_str = fid.read()
        xml = etree.fromstring(xml_str)
        data = self.parse_xml_to_dict(xml)["annotation"]
        data_height = int(data["size"]["height"])
        data_width = int(data["size"]["width"])
        # img_path = os.path.join(self.img_root, data["filename"])
        # image = Image.open(img_path)
        # if image.format != "JPEG":
        #     raise ValueError("Image format not JPEG")
        boxes = []
        labels = []
        iscrowd = []
        for obj in data["object"]:
            xmin = float(obj["bndbox"]["xmin"])
            xmax = float(obj["bndbox"]["xmax"])
            ymin = float(obj["bndbox"]["ymin"])
            ymax = float(obj["bndbox"]["ymax"])
            boxes.append([xmin, ymin, xmax, ymax])
            labels.append(self.class_dict[obj["name"]])
            iscrowd.append(int(obj["difficult"]))

        # convert everything into a torch.Tensor
        boxes = torch.as_tensor(boxes, dtype=torch.float32)
        labels = torch.as_tensor(labels, dtype=torch.int64)
        iscrowd = torch.as_tensor(iscrowd, dtype=torch.int64)
        image_id = torch.tensor([idx])
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])

        target = {}
        target["boxes"] = boxes
        target["labels"] = labels
        target["image_id"] = image_id
        target["area"] = area
        target["iscrowd"] = iscrowd

        return (data_height, data_width), target

    @staticmethod
    def collate_fn(batch):
        return tuple(zip(*batch))

注意:在进行数据预处理的时候,与进行图像分类数据预处理操作有些不同,如进行图像的翻转时,bbox的坐标信息也应该随之进行改变。

自定义数据集 - Dataset,AI,pytorch,python,深度学习

所以,我们自己定义了一个transform.py文章来源地址https://www.toymoban.com/news/detail-817729.html

import random
from torchvision.transforms import functional as F

class Compose(object):
    """组合多个transform函数"""
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, image, target):
        for t in self.transforms:
            image, target = t(image, target)
        return image, target

class ToTensor(object):
    """将PIL图像转为Tensor"""
    def __call__(self, image, target):
        image = F.to_tensor(image)
        return image, target

class RandomHorizontalFlip(object):
    """随机水平翻转图像以及bboxes"""
    def __init__(self, prob=0.5):
        self.prob = prob

    def __call__(self, image, target):
        if random.random() < self.prob:
            height, width = image.shape[-2:]
            image = image.flip(-1)  # 水平翻转图片
            bbox = target["boxes"]
            # bbox: xmin, ymin, xmax, ymax
            bbox[:, [0, 2]] = width - bbox[:, [2, 0]]  # 翻转对应bbox坐标信息
            target["boxes"] = bbox
        return image, target

到了这里,关于自定义数据集 - Dataset的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处: 如若内容造成侵权/违法违规/事实不符,请点击违法举报进行投诉反馈,一经查实,立即删除!

领支付宝红包 赞助服务器费用

相关文章

  • 【深度学习】PyTorch的dataloader制作自定义数据集

    PyTorch的dataloader是用于读取训练数据的工具,它可以自动将数据分割成小batch,并在训练过程中进行数据预处理。以下是制作PyTorch的dataloader的简单步骤: 导入必要的库 定义数据集类 需要自定义一个继承自 torch.utils.data.Dataset 的类,在该类中实现 __len__ 和 __getitem__ 方法。 创建

    2024年02月10日
    浏览(51)
  • 如何使用pytorch的Dataset, 来定义自己的Dataset

    Dataset与DataLoader的关系 Dataset: 构建一个数据集,其中含有所有的数据样本 DataLoader:将构建好的Dataset,通过shuffle、划分batch、多线程num_workers运行的方式,加载到可训练的迭代容器。 实战1:CSV数据集(结构化数据集) 实战2:图片数据集 ├── flower_data —├── flower_photo

    2024年01月22日
    浏览(48)
  • 跟着李沐学AI(动手学深度学习 PyTorch版)学习笔记——03安装(环境配置d2l、pytorch)(python3.7版本+Windows+各种问题解决措施)

    1.下载Miniconda下载地址 2.在安装过程中需要勾选“Add Anaconda to the system PATH environment variable”选项 3.检验win+R,输入cmd,在文本框输入conda --version 1.点击该链接+点击jupyter记事本下载压缩包 2.解压该压缩包 3.在解压后的文件夹地址栏输入cmd回车进入命令模式。 1.conda和pip默认使⽤

    2024年02月12日
    浏览(54)
  • 深度学习--PyTorch定义Tensor以及索引和切片

    ​这些方法只是开辟了空间,所附的初始值(非常大,非常小,0),后面还需要我们进行数据的存入。 torch.empty():返回一个没有初始化的Tensor,默认是FloatTensor类型。 torch.FloatTensor():返回没有初始化的FloatTensor。 torch.IntTensor():返回没有初始化的IntTensor。 随机均匀分布:

    2023年04月20日
    浏览(46)
  • 学习pytorch 2 导入查看dataset

    B站小土堆视频 https://download.pytorch.org/tutorial/hymenoptera_data.zip

    2024年02月13日
    浏览(32)
  • OpenCV与AI深度学习 | 实战 | YOLOv8自定义数据集训练实现手势识别 (标注+训练+预测 保姆级教程)

    本文来源公众号 “OpenCV与AI深度学习” ,仅用于学术分享,侵权删,干货满满。 原文链接:实战 | YOLOv8自定义数据集训练实现手势识别 (标注+训练+预测 保姆级教程)     本文将手把手教你用YoloV8训练自己的数据集并实现手势识别。 【1】安装torch, torchvision对应版本,这里先

    2024年04月23日
    浏览(79)
  • Pytorch Dataset类的使用(个人学习笔记)

    训练模型一般都是先处理 数据的输入问题 和 预处理问题 。 Pytorch提供了几个有用的工具: torch.utils.data.Dataset类 和 torch.utils.data.DataLoader类。 流程是先把 原始数据 转变成 torch.utils.data.Dataset类 , 随后再把得到 torch.utils.data.Dataset类 当作一个参数传递给 torch.utils.data.DataLoader类

    2024年02月05日
    浏览(37)
  • 【AI】《动手学-深度学习-PyTorch版》笔记(三):PyTorch常用函数

    返回一维张量(一维数组),官网说明,常见的三种用法如下 tensor.shape:查看张量的形状 tensor.reshape:返回改变形状后的张量,原张量不变

    2024年02月15日
    浏览(46)
  • 【AI】《动手学-深度学习-PyTorch版》笔记(八):线性回归

    线性函数如下: y ^ = w 1 x 1 + . . . + w d x d

    2024年02月14日
    浏览(48)
  • 探索深度学习世界:掌握PyTorch,成为AI领域的行家

    🏘️🏘️个人简介:以山河作礼。 🎖️🎖️: Python领域新星创作者,CSDN实力新星认证 PyTorch是一个开源的机器学习框架,由Facebook AI研究院开发和维护。它基于Torch,是一个动态图计算框架,可以支持动态构建计算图,使得它更加灵活和易于使用。 它是由Torch7团队开发,是

    2024年02月04日
    浏览(44)

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

博客赞助

微信扫一扫打赏

请作者喝杯咖啡吧~博客赞助

支付宝扫一扫领取红包,优惠每天领

二维码1

领取红包

二维码2

领红包