用于训练对象检测的参考脚本，实例分割和人员关键点检测可轻松支持 添加新的自定义数据集。数据集应继承自标准类，并实现 和 .torch.utils.data.Dataset__len__和__getitem__
__len____:数据集数量
__getitem__：返回图片和图片信息
import numpy as np
from torch.utils.data import Dataset
import os
import torch
import json
from PIL import Image
from lxml import etree
# 创建一个类 继承Dataset这个父类
class VOCDataSet(Dataset):
"""读取解析PASCAL VOC2007/2012数据集"""
# 初始化函数 voc_root：voc2012路径 transforms预处理方法 train.txt：训练集
def __init__(self, voc_root, year="2012", transforms=None, txt_name: str = "train.txt"):
voc_root = 'D:\\python\\deep learning\\VOCtrainval_11-May-2012\\VOCdevkit\\'
# 年份 voc2007 还是voc2012根目录 在预处理可以设置
assert year in ["2007", "2012"], "year must be in ['2007', '2012']"
# 增加容错能力 写入文件夹地址 voc_root+VOCdevkit+ f"VOC{year}
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")
# read train.txt or val.txt file
txt_path = os.path.join(self.root, "ImageSets", "Main", txt_name)
assert os.path.exists(txt_path), "not found {} file.".format(txt_name)
# 打开txt文件并读取
with open(txt_path) as read:
# self.xml_list 是 .xml文件 遍历.xml文件
# read.readlines() 遍历txt文件每一行
self.xml_list = [os.path.join(self.annotations_root, line.strip() + ".xml")
for line in read.readlines() if len(line.strip()) > 0]
# check file
assert len(self.xml_list) > 0, "in '{}' file does not find any information.".format(txt_path)
for xml_path in self.xml_list:
assert os.path.exists(xml_path), "not found '{}' file.".format(xml_path)
# read class_indict
# json_file：类别名称及索引信息
json_file = './pascal_voc_classes.json'
assert os.path.exists(json_file), "{} file not exist.".format(json_file)
# 打开并遍历json文件的每一行
with open(json_file, 'r') as f:
# 放入self.class_dict下
self.class_dict = json.load(f)
# 预处理
self.transforms = transforms
# 获取文件个数
def __len__(self):
return len(self.xml_list)
# idx传入索引值
def __getitem__(self, idx):
# read xml
# self.xml_list列表获得xml_path
xml_path = self.xml_list[idx]
# 打开xml文件
with open(xml_path) as fid:
xml_str = fid.read()
# from lxml import etree
# 通过etree读取xml文件
xml = etree.fromstring(xml_str)
# 将annotation传入到.parse_xml_to_dict方法中，将xml文件解析成字典形式
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 = [] # coco数据集才有
assert "object" in data, "{} lack of object information.".format(xml_path)
# 遍历object信息
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)
# convert everything into a torch.Tensor
# 列表转化成tenor格式
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])
# 建立一个字典
# 字典是一种可变的无序数据结构，
# 它由键（key）和值（value）组成的键值对（key-value pair）集合。
# 字典使用花括号 {} 来创建，每个键值对之间使用冒号 : 分隔，键和值之间使用逗号 , 分隔。
# 创建字典student = {"name": "Alice", "age": 20, "gender": "female"}
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
# 获取图像高度和宽度的函数
# 所要的信息在annotation文件里的size下
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
def parse_xml_to_dict(self, xml):
"""
将xml文件解析成字典形式，参考tensorflow的recursive_parse_xml_to_dict
Args:
xml: xml tree obtained by parsing XML file contents using lxml.etree
Returns:
Python dictionary holding XML contents.
"""
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))
# import transforms
# from draw_box_utils import draw_objs
# from PIL import Image
# import json
# import matplotlib.pyplot as plt
# import torchvision.transforms as ts
# import random
#
# # read class_indict
# category_index = {}
# try:
# json_file = open('./pascal_voc_classes.json', 'r')
# class_dict = json.load(json_file)
# category_index = {str(v): str(k) for k, v in class_dict.items()}
# except Exception as e:
# print(e)
# exit(-1)
#
# data_transform = {
# "train": transforms.Compose([transforms.ToTensor(),
# transforms.RandomHorizontalFlip(0.5)]),
# "val": transforms.Compose([transforms.ToTensor()])
# }
#
# # load train data set
# train_data_set = VOCDataSet(os.getcwd(), "2012", data_transform["train"], "train.txt")
# print(len(train_data_set))
# for index in random.sample(range(0, len(train_data_set)), k=5):
# img, target = train_data_set[index]
# img = ts.ToPILImage()(img)
# plot_img = draw_objs(img,
# target["boxes"].numpy(),
# target["labels"].numpy(),
# np.ones(target["labels"].shape[0]),
# category_index=category_index,
# box_thresh=0.5,
# line_thickness=3,
# font='arial.ttf',
# font_size=20)
# plt.imshow(plot_img)
# plt.show()
1.2Faster RCNN源码解析(pytorch)-霹雳吧啦Wz-深度学习-目标检测篇-哔哩哔哩视频 (bilibili.com)
1万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
