欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页  >  网络运营

PyTorch VOC数据加载

程序员文章站 2023-11-21 11:51:34
这是一个读取VOC数据集的例子,供大家参考但是我想这更适合作为一个云备份,哈哈"""因为我打算先搞一下目标检测所以先用VOC这个轻量级的来做"""import torchimport torch.utils.data as dataimport numpy as npimport cv2import xml # 标注是xml格式try: import xml.etree.cElementTree as ET # 解析xml的c语言版的模块except Import....
  • 这是一个读取VOC数据集的例子,供大家参考
  • 但是我想这更适合作为一个云备份,哈哈
"""
因为我打算先搞一下目标检测
所以先用VOC这个轻量级的来做
"""

import torch
import torch.utils.data as data
import numpy as np
import cv2
import xml  # 标注是xml格式

try:
    import xml.etree.cElementTree as ET  # 解析xml的c语言版的模块
except ImportError:
    import xml.etree.ElementTree as ET


VOC_CLASSES = {
    'aeroplane', 'bicycle', 'bird', 'boat',
    'bottle', 'bus', 'car', 'cat', 'chair',
    'cow', 'diningtable', 'dog', 'horse',
    'motorbike', 'person', 'pottedplant',
    'sheep', 'sofa', 'train', 'tvmonitor'
}

# 把str映射为int
dict_classes = dict(zip(VOC_CLASSES, range(len(VOC_CLASSES))))
# print(dict_classes['aeroplane'])


class ReadVOC(data.Dataset):
    def __init__(self, root):
        print("reading voc...")
        self.root = root

        self.img_idx = []
        self.ano_idx = []

        self.bbox = []
        self.obj_name = []  # 类别

        train_txt_path = self.root + "/ImageSets/Main/train_val.txt"  # train这个文件夹里面数量太少 换掉

        self.img_path = self.root + "/JPEGImages/"
        self.ano_path = self.root + "/Annotations/"

        # 首先读取txt文件进行训练集图片索引
        train_txt = open(train_txt_path)

        lines = train_txt.readlines()

        for line in lines:
            name = line.strip().split()[0]
            # print(name)  # name is in str type

            self.img_idx.append(self.img_path + name + '.jpg')
            self.ano_idx.append(self.ano_path + name + '.xml')  # 最好是在这直接解析出bbox

    def __getitem__(self, item):
        # print("getitem...")
        # print(self.img_idx[item])
        img = cv2.imread(self.img_idx[item])
        height, width, channels = img.shape

        targrts = ET.parse(self.ano_idx[item])  # .getroot()   # 运行时解析 逻辑更加清晰
        res = []  # 标注输出

        # find all obj in xml
        for obj in targrts.iter("object"):  # 便利物体
            name = obj.find('name').text.lower().strip()
            class_idx = dict_classes[name]

            bbox = obj.find('bndbox')

            pts = ['xmin', 'ymin', 'xmax', 'ymax']
            obj_bbox = []
            for i, pt in enumerate(pts):
                cur_pt = int(bbox.find(pt).text)
                cur_pt = cur_pt / width if i % 2 == 0 else cur_pt / height  # scale height or width
                obj_bbox.append(cur_pt)

            res.append(obj_bbox)  # 当前obj的所有bboxdf
            res.append(class_idx)

        img, res = self.data_trans(img, res)

        return img, res

    def __len__(self):
        data_lenth = len(self.img_idx)
        # print('data lenth is ', data_lenth)
        return data_lenth

    # 标注输入使用w h归一化的相对坐标
    def data_trans(self, img_input, bbox_input):
        # print("trans...")

        goal_size = (400, 400)

        # 在这时候,图像尺寸可以变化,只要目标不发生平移等等
        img = cv2.resize(img_input, goal_size)

        # pre-process input img
        img = torch.from_numpy(img).permute(2, 0, 1).float()

        # 把bbox转换成绝对坐标
        # bbox = [bbox_input[0] * goal_size[0], bbox_input[1] * goal_size[1], bbox_input[2] * goal_size[0], bbox_input[3] * goal_size[1]]
        # bbox = list(map(int, bbox))
        bbox = torch.tensor(bbox_input[0])

        return img, bbox


if __name__ == "__main__":
    ReadVOC(root='/home/.../data/VOCdevkit/VOC2012')

本文地址:https://blog.csdn.net/m0_38139098/article/details/107426408