第一步解析XML文件
'''
def get_xml_label_point(path):
tree = ET.ElementTree(file = path)
root = tree.getroot()
labels = []
points = []
point_json = []
basename = ""
for child1 in root:
if child1.tag == "object":
for child2 in child1:
if child2.tag == "name":
labels.append(child2.text)
elif child2.tag == "bndbox":
point = []
point.append(child2[0].text)
point.append(child2[1].text)
point.append(child2[2].text)
point.append(child2[3].text)
points.append(point)
if child1.tag == "filename":
basename = child1.text
return basename, labels, points
'''
使用get_xml_label_point获得xml文件的文件名，类别和四个坐标点。
'''
def transform_labels_points(labels, points):
person_number = 1
head_weared_number = 1
head_not_weared_number = 1
head_unsure_number = 1
points_json = []
for i in range(len(labels)):
if labels[i] == "person":
labels[i] = labels[i] + '_' + str(person_number)
person_number = person_number + 1
if labels[i] == "head_weared":
labels[i] = labels[i] + '_' + str(head_weared_number)
head_weared_number = head_weared_number + 1
if labels[i] == "head_not_weared":
labels[i] = labels[i] + '_' + str(head_not_weared_number)
head_not_weared_number = head_not_weared_number + 1
if labels[i] == "head_unsure_number":
labels[i] = labels[i] + '_' + str(head_unsure_number)
head_unsure_number = head_unsure_number + 1
for i in range(len(points)):
i0 = points[i][0]
i1 = points[i][1]
i2 = points[i][2]
i3 = points[i][3]
point0 = [i0, i1]
point1 = [i2, i1]
point2 = [i2, i3]
point3 = [i0, i3]
point = []
point.append(point0)
point.append(point1)
point.append(point2)
point.append(point3)
points_json.append(point)
return labels, points_json
def generate_json(picture_basename, labels, points):
str_json = {}
shapes = []
for i in range(len(labels)):
intPoint = []
shape = {}
for each in points[i]:
each_line=list(map(lambda x: int(float(x)), each))
intPoint.append(each_line)
shape["label"] = labels[i]
shape["points"] = intPoint
shapes.append(shape)
str_json["version"] = "3.14.1"
str_json["shapes"] = shapes
str_json["lineColor"] = [0, 255, 0, 128]
str_json["fillColor"] = [255, 0, 0, 128]
str_json["imagePath"] = picture_basename
return str_json
'''
将上面获得的文件名，类别和四个坐标点转换成符合mask rcnn所需要的json的文件格式。
'''
def base64encode_img(image_path):
src_image = Image.open(image_path)
output_buffer = io.BytesIO()
src_image.save(output_buffer, format='JPEG')
byte_data = output_buffer.getvalue()
base64_str = base64.b64encode(byte_data).decode('utf-8')
return base64_str
'''
将原图编码。
全部代码如下：
'''
import xml.etree.ElementTree as ET
import json
import base64
from PIL import Image
import io
import demjson
import re
import os
def get_xml_label_point(path):
tree = ET.ElementTree(file = path)
root = tree.getroot()
labels = []
points = []
point_json = []
basename = ""
for child1 in root:
if child1.tag == "object":
for child2 in child1:
if child2.tag == "name":
labels.append(child2.text)
elif child2.tag == "bndbox":
point = []
point.append(child2[0].text)
point.append(child2[1].text)
point.append(child2[2].text)
point.append(child2[3].text)
points.append(point)
if child1.tag == "filename":
basename = child1.text
return basename, labels, points
def transform_labels_points(labels, points):
person_number = 1
head_weared_number = 1
head_not_weared_number = 1
head_unsure_number = 1
points_json = []
for i in range(len(labels)):
if labels[i] == "person":
labels[i] = labels[i] + '_' + str(person_number)
person_number = person_number + 1
if labels[i] == "head_weared":
labels[i] = labels[i] + '_' + str(head_weared_number)
head_weared_number = head_weared_number + 1
if labels[i] == "head_not_weared":
labels[i] = labels[i] + '_' + str(head_not_weared_number)
head_not_weared_number = head_not_weared_number + 1
if labels[i] == "head_unsure_number":
labels[i] = labels[i] + '_' + str(head_unsure_number)
head_unsure_number = head_unsure_number + 1
for i in range(len(points)):
i0 = points[i][0]
i1 = points[i][1]
i2 = points[i][2]
i3 = points[i][3]
point0 = [i0, i1]
point1 = [i2, i1]
point2 = [i2, i3]
point3 = [i0, i3]
point = []
point.append(point0)
point.append(point1)
point.append(point2)
point.append(point3)
points_json.append(point)
return labels, points_json
def generate_json(picture_basename, labels, points):
str_json = {}
shapes = []
for i in range(len(labels)):
intPoint = []
shape = {}
for each in points[i]:
each_line=list(map(lambda x: int(float(x)), each))
intPoint.append(each_line)
shape["label"] = labels[i]
shape["points"] = intPoint
shapes.append(shape)
str_json["version"] = "3.14.1"
str_json["shapes"] = shapes
str_json["lineColor"] = [0, 255, 0, 128]
str_json["fillColor"] = [255, 0, 0, 128]
str_json["imagePath"] = picture_basename
return str_json
def base64encode_img(image_path):
src_image = Image.open(image_path)
output_buffer = io.BytesIO()
src_image.save(output_buffer, format='JPEG')
byte_data = output_buffer.getvalue()
base64_str = base64.b64encode(byte_data).decode('utf-8')
return base64_str
for file in os.listdir("."):
file_first = os.path.splitext(file)[0]
file_second = os.path.splitext(file)[1]
if file_second == ".xml":
file_name = file_first
print(file_name)
picture_name = file_name + ".jpg"
picture_basename, labels, points = get_xml_label_point("./" + file)
labels, points = transform_labels_points(labels, points)
str_json = generate_json(picture_basename, labels, points)
src_image_str = base64encode_img("./" + picture_name)
str_json["imageData"] = src_image_str
json_data = json.dumps(str_json)
f = open("./json/" + file_name + ".json", 'w')
f.write(json_data)
f.close()
'''