|
|
import os
|
|
|
import cv2
|
|
|
from lxml.etree import Element, SubElement, tostring
|
|
|
import xml.etree.ElementTree as ET
|
|
|
|
|
|
|
|
|
# 解析原始XML标注文件，获取人物的坐标信息
|
|
|
def parse_xml(xml_path,labels):
|
|
|
|
|
|
tree = ET.parse(xml_path)
|
|
|
root = tree.getroot()
|
|
|
|
|
|
labeks_list = []
|
|
|
for obj in root.findall('object'):
|
|
|
|
|
|
try:
|
|
|
|
|
|
name_labels = str(obj.find('name').text)
|
|
|
|
|
|
if name_labels == labels:
|
|
|
|
|
|
labels_dict = {}
|
|
|
labels_dict['name'] = name_labels
|
|
|
labels_dict['xmin'] = int(obj.find('bndbox/xmin').text)
|
|
|
labels_dict['ymin'] = int(obj.find('bndbox/ymin').text)
|
|
|
labels_dict['xmax'] = int(obj.find('bndbox/xmax').text)
|
|
|
labels_dict['ymax'] = int(obj.find('bndbox/ymax').text)
|
|
|
labeks_list.append(labels_dict)
|
|
|
except:
|
|
|
|
|
|
continue
|
|
|
#如果有需要的信息，就返回列表
|
|
|
if labeks_list:
|
|
|
return labeks_list
|
|
|
|
|
|
|
|
|
# 筛选手部和头部信息
|
|
|
def select_infors(xml_path,person):
|
|
|
|
|
|
infors_list = []
|
|
|
|
|
|
# 分别获得头部和手部信息
|
|
|
hands = parse_xml(xml_path=xml_path,labels='phone')
|
|
|
head = parse_xml(xml_path=xml_path,labels='head')
|
|
|
|
|
|
# 分别处理头部和手部信息
|
|
|
if hands:
|
|
|
|
|
|
for h,hand in enumerate(hands):
|
|
|
|
|
|
hand_result = change_bbox(infor=hand,contrast=person)
|
|
|
|
|
|
if hand_result:
|
|
|
|
|
|
infors_list.append(hand_result)
|
|
|
|
|
|
if head:
|
|
|
|
|
|
for h,head_infor in enumerate(head):
|
|
|
|
|
|
head_result = change_bbox(infor=head_infor,contrast=person)
|
|
|
|
|
|
if head_result:
|
|
|
infors_list.append(head_result)
|
|
|
else:
|
|
|
pass
|
|
|
|
|
|
return infors_list
|
|
|
|
|
|
|
|
|
# 转换坐标值
|
|
|
def change_bbox(infor,contrast):
|
|
|
|
|
|
if contrast['ymin'] <= infor['ymin'] and contrast['ymax'] >= infor['ymax'] and contrast['xmin'] <= infor['xmin'] and contrast['xmax'] >= infor['xmax']:
|
|
|
|
|
|
# 截取图片后的相对坐标
|
|
|
|
|
|
infor['ymin'] = infor['ymin'] - contrast['ymin']
|
|
|
infor['ymax'] = infor['ymax'] - contrast['ymin']
|
|
|
infor['xmin'] = infor['xmin'] - contrast['xmin']
|
|
|
infor['xmax'] = infor['xmax'] - contrast['xmin']
|
|
|
|
|
|
return infor
|
|
|
|
|
|
else:
|
|
|
pass
|
|
|
|
|
|
|
|
|
# 根据坐标值截图并保存
|
|
|
def cut_img_bbox(image_path,infors,num):
|
|
|
|
|
|
# # 读取大图
|
|
|
image = cv2.imread(image_path)
|
|
|
|
|
|
# 取大一些范围的截图尺寸
|
|
|
infors.update(ymin=infors['ymin'] - 25)
|
|
|
infors.update(ymax=infors['ymax'] + 25)
|
|
|
infors.update(xmin=infors['xmin'] - 25)
|
|
|
infors.update(xmax=infors['xmax'] + 25)
|
|
|
|
|
|
# 截取人物图片
|
|
|
person_image = image[infors['ymin']:infors['ymax'], infors['xmin']:infors['xmax']]
|
|
|
|
|
|
# 保存人物图片
|
|
|
person_filename = os.path.basename(image_path).split('.')[0] +'_'+ f"person_{num}.jpg"
|
|
|
# print(person_filename)
|
|
|
person_image_path = os.path.join(output_img_folder, person_filename)
|
|
|
cv2.imwrite(person_image_path, person_image)
|
|
|
|
|
|
return person_filename
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# 读取大图并截取人物图片
|
|
|
def extract_people(image_path, xml_path, output_img_folder, output_xml_folder):
|
|
|
|
|
|
# 解析XML标注文件、获得人的标注信息
|
|
|
people = parse_xml(xml_path=xml_path,labels='person')
|
|
|
|
|
|
# 逐个截取人物图片
|
|
|
for i, person in enumerate(people):
|
|
|
|
|
|
# 截图保存并返回图片名
|
|
|
person_filename = cut_img_bbox(image_path=image_path,infors=person,num=i)
|
|
|
|
|
|
# 写xml
|
|
|
node_root = Element('annotation')
|
|
|
node_folder = SubElement(node_root, 'folder')
|
|
|
node_folder.text = 'Images'
|
|
|
node_filename = SubElement(node_root, 'filename')
|
|
|
node_filename.text = str(person_filename)
|
|
|
node_size = SubElement(node_root, 'size')
|
|
|
node_width = SubElement(node_size, 'width')
|
|
|
node_width.text = str(person['xmax'] - person['xmin'])
|
|
|
node_height = SubElement(node_size, 'height')
|
|
|
node_height.text = str(person['ymax'] - person['ymin'])
|
|
|
node_depth = SubElement(node_size, 'depth')
|
|
|
node_depth.text = str(3)
|
|
|
|
|
|
# 获得小图上的相对坐标
|
|
|
select_infor_list = select_infors(xml_path=xml_path,person=person)
|
|
|
|
|
|
if len(select_infor_list)>=1: # 循环写入box
|
|
|
for infors in select_infor_list:
|
|
|
node_object = SubElement(node_root, 'object')
|
|
|
node_name = SubElement(node_object, 'name')
|
|
|
node_name.text = str(infors['name'])
|
|
|
node_difficult = SubElement(node_object, 'difficult')
|
|
|
node_difficult.text = '0'
|
|
|
node_bndbox = SubElement(node_object, 'bndbox')
|
|
|
node_xmin = SubElement(node_bndbox, 'xmin')
|
|
|
node_xmin.text = str(int(infors['xmin']))
|
|
|
node_ymin = SubElement(node_bndbox, 'ymin')
|
|
|
node_ymin.text = str(int(infors['ymin']))
|
|
|
node_xmax = SubElement(node_bndbox, 'xmax')
|
|
|
node_xmax.text = str(int(infors['xmax']))
|
|
|
node_ymax = SubElement(node_bndbox, 'ymax')
|
|
|
node_ymax.text = str(int(infors['ymax']))
|
|
|
else:
|
|
|
pass
|
|
|
xml = tostring(node_root, pretty_print=True) # 格式化显示，该换行的换行
|
|
|
|
|
|
# # 保存新的XML标注文件
|
|
|
new_xml_filename = person_filename.split('.')[0] + '.xml'
|
|
|
new_xml_path = os.path.join(output_xml_folder, new_xml_filename)
|
|
|
new_xml = open(new_xml_path, "wb")
|
|
|
new_xml.write(xml)
|
|
|
new_xml.close()
|
|
|
|
|
|
print(f"Saved person image {person_filename} and corresponding annotation {new_xml_filename}")
|
|
|
|
|
|
|
|
|
# 读取文件
|
|
|
def main(image_path,xml_path,output_img_folder,output_xml_folder):
|
|
|
|
|
|
# 检查保存的文件路径
|
|
|
print("frame image save path：{}".format(output_img_folder,output_xml_folder))
|
|
|
os.makedirs(output_img_folder, exist_ok=True)
|
|
|
os.makedirs(output_xml_folder, exist_ok=True)
|
|
|
|
|
|
img_file_list = os.listdir(image_path)
|
|
|
xml_file_list = os.listdir(xml_path)
|
|
|
|
|
|
for img_file in img_file_list:
|
|
|
|
|
|
img_file_full = os.path.join(image_path,img_file)
|
|
|
|
|
|
xml_file_combination_name = img_file.split('.')[0] + '.xml'
|
|
|
|
|
|
xml_file_combination = os.path.join(xml_path,xml_file_combination_name)\
|
|
|
|
|
|
if xml_file_combination_name in xml_file_list:
|
|
|
|
|
|
# 遍历文件处理数据
|
|
|
extract_people(image_path=img_file_full,
|
|
|
xml_path=xml_file_combination,
|
|
|
output_img_folder=output_img_folder,
|
|
|
output_xml_folder=output_xml_folder
|
|
|
)
|
|
|
else:
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
|
|
# 原始大图像路径
|
|
|
image_path = "E:/Images_Data/XZ/bank2/yolov8/images"
|
|
|
|
|
|
# 原始XML标注文件路径
|
|
|
xml_path = "E:/Images_Data/XZ/bank2/yolov8/XML_collect"
|
|
|
|
|
|
# 输出文件夹路径
|
|
|
output_img_folder = "E:/Images_Data/XZ/bank2/yolov8/output_files/images"
|
|
|
output_xml_folder = 'E:/Images_Data/XZ/bank2/yolov8/output_files/Annotations'
|
|
|
|
|
|
# parse_xml('E:/Images_Data/XZ/bank2/yolov8/XML_collect/Test_1_000014.xml')
|
|
|
|
|
|
main(image_path=image_path,
|
|
|
xml_path=xml_path,
|
|
|
output_img_folder=output_img_folder,
|
|
|
output_xml_folder=output_xml_folder
|
|
|
)
|
|
|
|