这个代码的主网络修改为resnet50。
write_txt.py
from curses import endwin
import xml.etree.ElementTree as ET
import os
import random
"""
这个文件的主要功能是读取文件夹内所有的xml的文件以及信息，将这些信息(name,bbox,class)写入一个txt文件中，并且按照7:3划分训练集和测试集
"""
VOC_CLASSES = ( # 定义所有的类名
'aeroplane', 'bicycle', 'bird', 'boat',
'bottle', 'bus', 'car', 'cat', 'chair',
'cow', 'diningtable', 'dog', 'horse',
'motorbike', 'person', 'pottedplant',
'sheep', 'sofa', 'train', 'tvmonitor') # 使用其他训练集需要更改
# VOC_CLASSES=[]# 自己的数据集
# 切换成当前路径-需要修改
# os.chdir('/root/workspace/YOLOV1-pytorch/')
# 定义一些参数
train_set = open('voctrain.txt', 'w')
test_set = open('voctest.txt', 'w')
# Annotations = 'VOCdevkit//VOC2007//Annotations//'
Annotations='VOCdevkit/VOC2007/Annotations/'
xml_files = os.listdir(Annotations)
random.shuffle(xml_files) # 打乱数据集
train_num = int(len(xml_files) * 0.7) # 训练集数量
train_lists = xml_files[:train_num] # 训练列表
test_lists = xml_files[train_num:] # 测测试列表
# 输出一些信息
print("train_lists:",len(train_lists))
print("test_lists:",len(test_lists))
def parse_rec(filename): # 输入xml文件名
"""
读取xml文件信息，在"object"目录下查看"difficult"值是否为1，若不为1则在名为"obj_struct"的字典中存入"bbox"和"name"的信息，
再将这个字典作为名为"objects"的列表的元素，最终输出这个列表。所以这个名为"objects"的列表中的每一个元素都是一个字典。
"""
tree = ET.parse(filename)# 生成一个总目录名为tree
objects = []
for obj in tree.findall('object'):
obj_struct = {
}
difficult = int(obj.find('difficult').text)
if difficult == 1: # 若为1则跳过本次循环
continue
obj_struct['name'] = obj.find('name').text
bbox = obj.find('bndbox')
obj_struct['bbox'] = [int(float(bbox.find('xmin').text)),
int(float(bbox.find('ymin').text)),
int(float(bbox.find('xmax').text)),
int(float(bbox.find('ymax').text))]
objects.append(obj_struct)
return objects
def write_txt():
count = 0
for train_list in train_lists: # 生成训练集txt
count += 1
image_name = train_list.split('.')[0] + '.jpg' # 图片文件名
results = parse_rec(Annotations + train_list)
if len(results) == 0:
print(train_list)
continue
train_set.write(image_name)
for result in results:
class_name = result['name']
# # 添加类别名字
# if class_name not in VOC_CLASSES:
# VOC_CLASSES.append(class_name)
bbox = result['bbox']
class_name = VOC_CLASSES.index(class_name)
train_set.write(' ' + str(bbox[0]) +
' ' + str(bbox[1]) +
' ' + str(bbox[2]) +
' ' + str(bbox[3]) +
' ' + str(class_name))
train_set.write('\n')
train_set.close()
for test_list in test_lists: # 生成测试集txt
count += 1
image_name = test_list.split('.')[0] + '.jpg' # 图片文件名
results = parse_rec(Annotations + test_list)
if len(results) == 0:
print(test_list)
continue
test_set.write(image_name)
for result in results:
class_name = result['name']
# # 添加类别名字
# if class_name not in VOC_CLASSES:
# VOC_CLASSES.append(class_name)
bbox = result['bbox']
class_name = VOC_CLASSES.index(class_name)
test_set.write(' ' + str(bbox[0]) +
' ' + str(bbox[1]) +
' ' + str(bbox[2]) +
' ' + str(bbox[3]) +
' ' + str(class_name))
test_set.write('\n')
test_set.close()
"""
if __name__ == "__main__": 的作用
在Python中，每个Python文件（模块）都可以作为脚本直接运行，也可以被其他文件导入。__name__ 是一个特殊变量，
当文件被直接运行时，__name__ 的值被设置为 "__main__"。如果文件是被导入的，则 __name__ 的值会被设置为该模块的名字。
if __name__ == "__main__": 这行代码的作用是判断该文件是否作为主程序运行。如果是，则执行该条件语句块下的代码。
这种方式通常用于提供一个可执行的入口点给该文件，同时也允许该文件中的函数和类被其他文件导入而不会自动执行这些代码。
"""
if __name__ == '__main__':
write_txt()
print(VOC_CLASSES)# 类别名称
yoloData.py
import torch
import cv2
import os
import os.path
import random
import numpy as np
from torch.utils.data import DataLoader, Dataset
from torchvision.transforms import ToTensor
from PIL import Image
# from write_txt import VOC_CLASSES # 这个使用要谨慎，因为对应文件里面定义的两个txt是全局变量，导入的时候里面的全局变量会重新赋值
CLASS_NUM = 20 # 使用其他训练集需要更改
# CLASS_NUM=len(VOC_CLASSES) # 类别的数量
# os.chdir('/root/workspace/YOLOV1-pytorch/')
class yoloDataset(Dataset):
image_size = 448 # 输入图片大小
def __init__(self, img_root, list_file, train, transform): # list_file为txt文件 img_root为图片路径
"""
逐行读取生成的文本文件的内容，然后对其进行分类，将信息保存在fnames，boxes，labels三个列表中
"""
self.root = img_root
self.train = train
self.transform = transform
# 后续要提取txt文件信息，分类后装入以下三个列表
self.fnames = []
self.boxes = []
self.labels = []
self.S = 7 # YOLOV1
self.B = 2 # 相关
self.C = CLASS_NUM # 参数
self.mean = (123, 117, 104) # RGB
file_txt = open(list_file,'r')
lines = file_txt.readlines() # 读取txt文件每一行
for line in lines: # 逐行开始操作
# strip() # 移除首位的换行符号；split() # 以空格为分界线，将所有元素组成一个列表
splited = line.strip().split() # 移除首位的换行符号再生成一张列表
self.fnames.append(splited[0]) # 存储图片的名字
num_boxes = (len(splited) - 1) // 5 # 每一幅图片里面有多少个bbox
box = []
label = []
for i in range(num_boxes): # bbox四个角的坐标
x = float(splited[1 + 5 * i])
y = float(splited[2 + 5 * i])
x2 = float(splited[3 + 5 * i])
y2 = float(splited[4 + 5 * i])
c = splited[5 + 5 * i] # 代表物体的类别，即是20种物体里面的哪一种 值域 0-19
box.append([x, y, x2, y2])
label.append(int(c))
self.boxes.append(torch.Tensor(box))
self.labels.append(torch.LongTensor(label))
self.num_samples = len(self.boxes)
# 访问坐标的时候就会直接执行这个函数
def __getitem__(self, idx):
fname = self.fnames[idx]
img = cv2.imread(os.path.join(self.root + fname))
boxes = self.boxes[idx].clone()
labels = self.labels[idx].clone()
if self.train: # 数据增强里面的各种变换用torch自带的transform是做不到的，因为对图片进行旋转、随即裁剪等会造成bbox的坐标也会发生变化，所以需要自己来定义数据增强
img, boxes = self.random_flip(img, boxes) # 随机翻转
img, boxes = self.randomScale(img, boxes) # 随机伸缩变换
img = self.randomBlur(img)# 随机模糊处理
img = self.RandomBrightness(img)# 随机调整亮度
# img = self.RandomHue(img)
# img = self.RandomSaturation(img)
img, boxes, labels = self.randomShift(img, boxes, labels)# 平移转换
# img, boxes, labels = self.randomCrop(img, boxes, labels)
h, w, _ = img.shape
boxes /= torch.Tensor([w, h, w, h]).expand_as(boxes) # 坐标归一化处理，为了方便训练，这个表示的bbox的宽高占整个图像的比例
img = self.BGR2RGB(img) # because pytorch pretrained model use RGB
img = self.subMean(img, self.mean) # 减去均值
"""
这里对图像resize后不需要对boxes变化，原因一是这里不是图像增强，只是方便图片输入网络；
原因二是YOLO原文写的是，对bbox的宽高做归一化，这个归一化是相当于整个原来图像的宽高进行归一化的（上面已经归一化了），而对bbox的中心坐标的归一化
是相当于bbox所在的grid cell的左上角坐标进行归一化的，也就是下面的encoder操作，所以这一步是正确的。
而且在后面使用到这个bbox的xywh的时候，是会做相应的操作的，详情可以看yoloLoss
"""
# YOLO V1输入图像大小设置为448*448* 3
img = cv2.resize(img, (self.image_size, self.image_size)) # 将所有图片都resize到指定大小，这里不是图像增强，而是为了方便网络的输入
target = self.encoder(boxes, labels) # 将图片标签编码到7x7*30的向量
for t in self.transform:
img = t(img)
# 返回的img是经过图像增强的img
return img, target
def __len__(self):
return self.num_samples
# def letterbox_image(self, image, size):
# # 对图片进行resize，使图片不失真。在空缺的地方进行padding
# iw, ih = image.size
# scale = min(size / iw, size / ih)
# nw = int(iw * scale)
# nh = int(ih * scale)
#
# image = image.resize((nw, nh), Image.BICUBIC)
# new_image = Image.new('RGB', size, (128, 128, 128))
# new_image.paste(image, ((size - nw) // 2, (size - nh) // 2))
# return new_image
def encoder(self, boxes, labels): # 输入的box为归一化形式(X1,Y1,X2,Y2) , 输出ground truth (7*7*30)
grid_num = 7
target = torch.zeros((grid_num, grid_num, int(CLASS_NUM + 10))) # 7*7*30
# cell_size 是图像宽度和高度被划分成的等分数，用于将归一化的坐标转换为网格索引。
cell_size = 1. / grid_num # 1/7
# 这个是bbox的归一化后的宽高
wh = boxes[:, 2:] - boxes[:, :2] # wh = [w, h] 1*1
# 物体中心坐标集合
cxcy = (boxes[:, 2:] + boxes[:, :2]) / 2 # 归一化含小数的中心坐标
for i in range(cxcy.size()[0]):
cxcy_sample = cxcy[i] # 中心坐标 1*1
"""
ij 并不是直接表示“左上角坐标（7*7)为整数，而是表示边界框中心点所在的网格的索引。ceil()表示向上取整；-1是因为python的索引是从0开始的
ij 是一个包含两个元素的tensor，分别表示边界框中心点所在的网格的x和y索引
下面的公式的解释可以这样理解：假设有一个图像大小为w*h，现在把图像分为7分，问坐标为（x，y）的点位于哪一个网格中，这就是小学乘法问题，
很明显，求一下x和y占比w和h的占比，分别乘以7，最后向上取整，所以答案就是（7x/w,7y/h）.ceil(),这里再看cxcy_sample本身就是已经归一化（也就是已经除以w）
了，所以直接乘7，也就是 / cell_size 就可以得到结果。-1是为了让索引从0开始。
"""
ij = (cxcy_sample / cell_size).ceil() - 1 # 左上角坐标 （7*7)为整数
# 这里先1后0是因为坐标提取就是先行后列
# 第一个框的置信度，4表示第一个标注框的置信度存储在下标为4的位置，下面9同理，并且这里的意义是，只有有标注框的置信度置位为1
target[int(ij[1]), int(ij[0]), 4] = 1
# 第二个框的置信度
target[int(ij[1]), int(ij[0]), 9] = 1
target[int(ij[1]), int(ij[0]), int(labels[i]) + 10] = 1 # 20个类别对应处的概率设置为1
xy = ij * cell_size # 归一化左上坐标 （1*1）
# 在YOLOV1原文中，其bbox的五个参数中的x，y就是中心坐标相对于其grid cell左上角的坐标的相对值
delta_xy = (cxcy_sample - xy) / cell_size # 中心与左上坐标差值 （7*7）
# 坐标w,h代表了预测的bounding box的width、height相对于整幅图像width,height的比例
target[int(ij[1]), int(ij[0]), 2:4] = wh[i] # w1,h1
target[int(ij[1]), int(ij[0]), :2] = delta_xy # x1,y1
# 每一个网格有两个边框，在真实数据中，两个边框的值是一样的
target[int(ij[1]), int(ij[0]), 7:9] = wh[i] # w2,h2
# 由此可得其实返回的中心坐标其实是相对左上角顶点的偏移，因此在进行预测的时候还需要进行解码
target[int(ij[1]), int
最低0.47元/天 解锁文章
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
