模型推理结果生成XML文件
很多时候标注数据是一个令人头疼的问题,不管是时间成本还是金钱成本,所以我们可以用小样本数据训练之后的模型权重去推理得到一批数据,在这批推理后的数据上进行微调作为新的样本再加入到网络中,也可以通过推理后的图片筛选出当前模型对哪些场景的识别效果差的数据,着重标注这批数据以提高模型的泛化性(也就是困难样本),在目标检测模型中大多数用到的是xml格式的label文件,所以我们可以通过生成xml文件进行预标注,所以记录一下生成xml文件的方法
上代码
import xml.etree.ElementTree as ET
‘shape-->推理得到的bbox的坐标
filename_-->被推理图片的名字
w_f-->图片所在的文件夹名字
class_-->储存类别的列表(如[0,1,2,3,...])’
def to_xml(shape,filename_,w_f, class_):
tree = ET.parse("/data1/fxs/normal_72954.xml")#将xml解析为树:这里的xml是一个无标签的模板文件
root = tree.getroot()
root.find('folder').text=w_f#更改现有节点的信息
root.find('filename').text='{}'.format(str(filename_))
root.find('path').text='/home/dell/fxs_hand/yolov5/my_tools/test/{}'.format(str(filename_))
# print(len(shape))
for j in range(len(shape)):
root_a=ET.Element("object")#新增节点
root_a=root.append(root_a)
students = root.findall("object")
x=ET.Element("name")#在object节点下新增name节点
root_c=students[-1].append(x)
if class_[j] == 0:
students[-1].find('name').text='people'对应到每个类别
else:
students[-1].find('name').text='dog'
x = ET.Element("pose")
root_c = students[-1].append(x)
students[-1].find('pose').text = 'Unspecified'
x = ET.Element("truncated")
root_c = students[-1].append(x)
students[-1].find('truncated').text = '0'
x = ET.Element("difficult")
root_c = students[-1].append(x)
students[-1].find('difficult').text = '0'
x = ET.Element("bndbox")
root_c = students[-1].append(x)
z_sku=students[-1].find('bndbox')
x1=ET.Element("xmin")
z_sku.append(x1)
z_sku.find('xmin').text=str(shape[j][0])
y1 = ET.Element("ymin")
z_sku.append(y1)
z_sku.find('ymin').text = str(shape[j][1])
x2 = ET.Element("xmax")
z_sku.append(x2)
z_sku.find('xmax').text = str(shape[j][2])
y2 = ET.Element("ymax")
z_sku.append(y2)
z_sku.find('ymax').text = str(shape[j][3])
tree.write('/home/dell/fxs_hand/yolov5/my_tools/test/{}.xml'.format(filename_.split('.')[0]))