我要评论"""
因为我打算先搞一下目标检测
所以先用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
如对本文有疑问， 点击进行留言回复！！
网友评论