228 lines
7.7 KiB
Python
228 lines
7.7 KiB
Python
|
# -*- coding: utf-8 -*-
|
|||
|
|
# -*- coding: utf-8 -*-
|
|||
|
|
# @Time : 20210610
|
|||
|
|
# @Author : dejahu
|
|||
|
|
# @File : gen_yolo_data.py
|
|||
|
|
# @Software: PyCharm
|
|||
|
|
# @Brief : 生成测试、验证、训练的图片和标签
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import shutil
|
|||
|
|
from pathlib import Path
|
|||
|
|
from shutil import copyfile
|
|||
|
|
|
|||
|
|
import cv2
|
|||
|
|
from PIL import Image, ImageDraw
|
|||
|
|
from xml.dom.minidom import parse
|
|||
|
|
import numpy as np
|
|||
|
|
import os.path as osp
|
|||
|
|
import random
|
|||
|
|
|
|||
|
|
# main，首先在当前目录生成数据的划分
|
|||
|
|
|
|||
|
|
# 开始转化数据集
|
|||
|
|
|
|||
|
|
# 更换代码中images替换的逻辑，和文件对应起来
|
|||
|
|
|
|||
|
|
# todo 修改为你的数据的根目录
|
|||
|
|
FILE_ROOT = "/scm/data/xianyu/Mask/"
|
|||
|
|
IMAGE_SET_ROOT = FILE_ROOT + "VOC2021_Mask/ImageSets/Main" # 图片区分文件的路径
|
|||
|
|
IMAGE_PATH = FILE_ROOT + "VOC2021_Mask/JPEGImages" # 图片的位置
|
|||
|
|
ANNOTATIONS_PATH = FILE_ROOT + "VOC2021_Mask/Annotations" # 数据集标签文件的位置
|
|||
|
|
LABELS_ROOT = FILE_ROOT + "VOC2021_Mask/Labels" # 进行归一化之后的标签位置
|
|||
|
|
DEST_PPP = FILE_ROOT + "mask_yolo_format"
|
|||
|
|
DEST_IMAGES_PATH = "mask_yolo_format/images" # 区分训练集、测试集、验证集的图片目标路径
|
|||
|
|
DEST_LABELS_PATH = "mask_yolo_format/labels" # 区分训练集、测试集、验证集的标签文件目标路径
|
|||
|
|
|
|||
|
|
if osp.isdir(LABELS_ROOT):
|
|||
|
|
shutil.rmtree(LABELS_ROOT)
|
|||
|
|
print("Labels存在，已删除")
|
|||
|
|
|
|||
|
|
if osp.isdir(DEST_PPP):
|
|||
|
|
shutil.rmtree(DEST_PPP)
|
|||
|
|
print("Dest目录存在，已删除")
|
|||
|
|
# todo 修改为你数据集的标签名称
|
|||
|
|
label_names = ['face', 'face_mask']
|
|||
|
|
|
|||
|
|
|
|||
|
|
def cord_converter(size, box):
|
|||
|
|
"""
|
|||
|
|
将标注的 xml 文件标注转换为 darknet 形的坐标
|
|||
|
|
:param size: 图片的尺寸： [w,h]
|
|||
|
|
:param box: anchor box 的坐标 [左上角x,左上角y,右下角x,右下角y,]
|
|||
|
|
:return: 转换后的 [x,y,w,h]
|
|||
|
|
"""
|
|||
|
|
x1 = int(box[0])
|
|||
|
|
y1 = int(box[1])
|
|||
|
|
x2 = int(box[2])
|
|||
|
|
y2 = int(box[3])
|
|||
|
|
|
|||
|
|
dw = np.float32(1. / int(size[0]))
|
|||
|
|
dh = np.float32(1. / int(size[1]))
|
|||
|
|
|
|||
|
|
w = x2 - x1
|
|||
|
|
h = y2 - y1
|
|||
|
|
x = x1 + (w / 2)
|
|||
|
|
y = y1 + (h / 2)
|
|||
|
|
|
|||
|
|
x = x * dw
|
|||
|
|
w = w * dw
|
|||
|
|
y = y * dh
|
|||
|
|
h = h * dh
|
|||
|
|
return [x, y, w, h]
|
|||
|
|
|
|||
|
|
|
|||
|
|
def save_file(img_jpg_file_name, size, img_box):
|
|||
|
|
# print("保存图片")
|
|||
|
|
save_file_name = LABELS_ROOT + '/' + img_jpg_file_name + '.txt'
|
|||
|
|
# print(save_file_name)
|
|||
|
|
file_path = open(save_file_name, "a+")
|
|||
|
|
# 默认给定的是id为0，防止错误数据的出现
|
|||
|
|
for box in img_box:
|
|||
|
|
box_name = box[0]
|
|||
|
|
cls_num = 0
|
|||
|
|
if box_name in label_names:
|
|||
|
|
cls_num = label_names.index(box_name)
|
|||
|
|
new_box = cord_converter(size, box[1:])
|
|||
|
|
file_path.write(f"{cls_num} {new_box[0]} {new_box[1]} {new_box[2]} {new_box[3]}\n")
|
|||
|
|
file_path.flush()
|
|||
|
|
file_path.close()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_dataset_box_feature(file_name, point_array):
|
|||
|
|
"""
|
|||
|
|
使用样本数据测试数据集的建议框
|
|||
|
|
:param image_name: 图片文件名
|
|||
|
|
:param point_array: 全部的点 [建议框sx1,sy1,sx2,sy2]
|
|||
|
|
:return: None
|
|||
|
|
"""
|
|||
|
|
im = Image.open(rf"{IMAGE_PATH}\{file_name}")
|
|||
|
|
imDraw = ImageDraw.Draw(im)
|
|||
|
|
for box in point_array:
|
|||
|
|
x1 = box[1]
|
|||
|
|
y1 = box[2]
|
|||
|
|
x2 = box[3]
|
|||
|
|
y2 = box[4]
|
|||
|
|
imDraw.rectangle((x1, y1, x2, y2), outline='red')
|
|||
|
|
im.show()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_xml_data(file_path, img_xml_file):
|
|||
|
|
img_path = file_path + '/' + img_xml_file + '.xml'
|
|||
|
|
# print(img_path)
|
|||
|
|
dom = parse(img_path)
|
|||
|
|
root = dom.documentElement
|
|||
|
|
img_name = root.getElementsByTagName("filename")[0].childNodes[0].data
|
|||
|
|
img_jpg_file_name = img_xml_file + '.jpg'
|
|||
|
|
# print(img_jpg_file_name)
|
|||
|
|
cv2.imread(img_jpg_file_name)
|
|||
|
|
img_size = root.getElementsByTagName("size")[0]
|
|||
|
|
if len(img_size) == 0:
|
|||
|
|
img_h, img_w, c = cv2.imread(img_jpg_file_name).shape
|
|||
|
|
else:
|
|||
|
|
img_w = img_size.getElementsByTagName("width")[0].childNodes[0].data
|
|||
|
|
img_h = img_size.getElementsByTagName("height")[0].childNodes[0].data
|
|||
|
|
img_c = img_size.getElementsByTagName("depth")[0].childNodes[0].data
|
|||
|
|
objects = root.getElementsByTagName("object")
|
|||
|
|
|
|||
|
|
img_box = []
|
|||
|
|
for box in objects:
|
|||
|
|
cls_name = box.getElementsByTagName("name")[0].childNodes[0].data
|
|||
|
|
# todo 更换坐标点转换的逻辑
|
|||
|
|
x1 = int(float(box.getElementsByTagName("xmin")[0].childNodes[0].data))
|
|||
|
|
y1 = int(float(box.getElementsByTagName("ymin")[0].childNodes[0].data))
|
|||
|
|
x2 = int(float(box.getElementsByTagName("xmax")[0].childNodes[0].data))
|
|||
|
|
y2 = int(float(box.getElementsByTagName("ymax")[0].childNodes[0].data))
|
|||
|
|
# print("box:(c,xmin,ymin,xmax,ymax)", cls_name, x1, y1, x2, y2)
|
|||
|
|
|
|||
|
|
img_box.append([cls_name, x1, y1, x2, y2])
|
|||
|
|
# test_dataset_box_feature(img_jpg_file_name, img_box)
|
|||
|
|
save_file(img_xml_file, [img_w, img_h], img_box)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def copy_data(img_set_source, img_labels_root, imgs_source, type):
|
|||
|
|
file_name = img_set_source + '/' + type + ".txt"
|
|||
|
|
file = open(file_name)
|
|||
|
|
|
|||
|
|
# 判断文件夹是否存在，不存在则创建
|
|||
|
|
root_file = Path(FILE_ROOT + DEST_IMAGES_PATH + '/' + type)
|
|||
|
|
if not root_file.exists():
|
|||
|
|
print(f"Path {root_file} is not exit")
|
|||
|
|
os.makedirs(root_file)
|
|||
|
|
|
|||
|
|
root_file = Path(FILE_ROOT + DEST_LABELS_PATH + '/' + type)
|
|||
|
|
if not root_file.exists():
|
|||
|
|
print(f"Path {root_file} is not exit")
|
|||
|
|
os.makedirs(root_file)
|
|||
|
|
|
|||
|
|
# 遍历文件夹
|
|||
|
|
for line in file.readlines():
|
|||
|
|
# print(line)
|
|||
|
|
img_name = line.strip('\n')
|
|||
|
|
img_sor_file = imgs_source + '/' + img_name + '.jpg'
|
|||
|
|
label_sor_file = img_labels_root + '/' + img_name + '.txt'
|
|||
|
|
|
|||
|
|
# 复制图片
|
|||
|
|
DICT_DIR = FILE_ROOT + DEST_IMAGES_PATH + '/' + type
|
|||
|
|
img_dict_file = DICT_DIR + '/' + img_name + '.jpg'
|
|||
|
|
copyfile(img_sor_file, img_dict_file)
|
|||
|
|
|
|||
|
|
# 复制 label
|
|||
|
|
DICT_DIR = FILE_ROOT + DEST_LABELS_PATH + '/' + type
|
|||
|
|
img_dict_file = DICT_DIR + '/' + img_name + '.txt'
|
|||
|
|
copyfile(label_sor_file, img_dict_file)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
# 将文件进行 train 和 val 的区分
|
|||
|
|
# 用于生成测试集使用，目前还在测试阶段，暂时不发布
|
|||
|
|
img_set_root = IMAGE_SET_ROOT
|
|||
|
|
imgs_root = IMAGE_PATH
|
|||
|
|
img_labels_root = LABELS_ROOT
|
|||
|
|
if osp.isdir(img_labels_root) == False:
|
|||
|
|
os.makedirs(img_labels_root)
|
|||
|
|
os.makedirs(DEST_PPP)
|
|||
|
|
names = os.listdir(ANNOTATIONS_PATH)
|
|||
|
|
# real_names = []
|
|||
|
|
# for name in names:
|
|||
|
|
# if name.split(".")[-1] == "xml":
|
|||
|
|
# real_names.append(name.split(".")[0])
|
|||
|
|
# # print(real_names)
|
|||
|
|
# # print(real_names)
|
|||
|
|
# random.shuffle(real_names)
|
|||
|
|
# # print(real_names)
|
|||
|
|
# length = len(real_names)
|
|||
|
|
# split_point = int(length * 0.2)
|
|||
|
|
#
|
|||
|
|
# val_names = real_names[:split_point]
|
|||
|
|
# train_names = real_names[split_point:]
|
|||
|
|
#
|
|||
|
|
# # 开始生成文件
|
|||
|
|
# np.savetxt('data/val.txt', np.array(val_names), fmt="%s", delimiter="\n")
|
|||
|
|
# np.savetxt('data/test.txt', np.array(val_names), fmt="%s", delimiter="\n")
|
|||
|
|
# np.savetxt('data/train.txt', np.array(train_names), fmt="%s", delimiter="\n")
|
|||
|
|
# print("txt文件生成完毕，请放在VOC2012的ImageSets/Main的目录下")
|
|||
|
|
# 生成标签
|
|||
|
|
root = ANNOTATIONS_PATH
|
|||
|
|
files_all = os.listdir(root)
|
|||
|
|
files = []
|
|||
|
|
for file in files_all:
|
|||
|
|
if file.split(".")[-1] == "xml":
|
|||
|
|
files.append(file.split(".")[0])
|
|||
|
|
# print(len(files))
|
|||
|
|
for file in files:
|
|||
|
|
# print("file name: ", file)
|
|||
|
|
file_xml = file.split(".")
|
|||
|
|
# try:
|
|||
|
|
get_xml_data(root, file_xml[0])
|
|||
|
|
# except:
|
|||
|
|
# print(file_xml[0])
|
|||
|
|
# break
|
|||
|
|
# copy_data(img_set_root, img_labels_root, imgs_root, "train")
|
|||
|
|
# copy_data(img_set_root, img_labels_root, imgs_root, "val")
|
|||
|
|
# copy_data(img_set_root, img_labels_root, imgs_root, "test")
|
|||
|
|
# data = list(np.loadtxt("tt100k.txt", dtype=str))
|
|||
|
|
# print(data)
|
|||
|
|
# print(len(data))
|
|||
|
|
print("数据转换已完成！")
|