这篇文章是一篇记录，记录了我在制作数据集的时候用到的代码。
这些代码都是博主在网上找到到的，某些地方也做出了更适合博主的更改，所以这些代码并不一定适合每一个人，切记切记。
1、首先博主使用的是andconda3，创建的环境使用的是python 3.10
可能会有人问为什么不用最新的python，这是因为老版本更让人放心，与其它会用到的包兼容性更高，至于我使用python 3.10完全是因为比较认可“十全十美”这个词。
2、博主用到的标注软件是labelme
这个标注软件还是挺好用的，我个人是比较中意它使用“A”和“D”就可以进行翻页。它也有ai自动分割的功能，不过对于使用目标识别的人来说用处不大。
它标注完之后生成的是'.json'文件，这个文件是不能拿来直接给yolo使用的，这也是下面一堆代码存在的意义。
3、进行json文件--xml文件--txt文件的转化
首先我们要强调一点，在这里博主使用的照片是‘.bmp’格式的，如果想使用之后的代码，首先你要将你的图片变为‘.bmp’格式，下面是将‘.jpg’变为‘.bmp’格式的代码：
from PIL import Image
import os
def convert_jpg_to_bmp(input_folder, output_folder):
# 检查输出文件夹是否存在，如果不存在则创建
if not os.path.exists(output_folder):
os.makedirs(output_folder)
# 遍历输入文件夹中的所有文件
for filename in os.listdir(input_folder):
if filename.endswith(".jpg"):
# 构建完整的输入文件路径
input_path = os.path.join(input_folder, filename)
# 构建完整的输出文件路径
output_path = os.path.join(output_folder, os.path.splitext(filename)[0] + ".bmp")
# 打开JPG文件并转换为BMP
with Image.open(input_path) as img:
img.save(output_path, "BMP")
print(f"Converted {input_path} to {output_path}")
# 使用示例
input_folder = r"你照片所在的绝对路径"
output_folder = r"你希望改变格式之后的照片存在的绝对路径"
convert_jpg_to_bmp(input_folder, output_folder)
当然这段代码也可以改变别的照片格式，只需要将代码中的‘.jpg’更改一下就可以。
之后我们就要进行json文件转xml文件的操作，这段代码是我从网上找到的，是我感觉转的最细致的代码，缺点就是有点太多。
首先创建read_json_anno.py这个文件，并将下面的代码粘贴：
import numpy as np
import json
class ReadAnno:
def __init__(self, json_path, process_mode="rectangle"):
self.json_data = json.load(open(json_path))
self.filename = self.json_data['imagePath']
self.width = self.json_data['imageWidth']
self.height = self.json_data['imageHeight']
self.coordis = []
assert process_mode in ["rectangle", "polygon"]
if process_mode == "rectangle":
self.process_polygon_shapes()
elif process_mode == "polygon":
self.process_polygon_shapes()
def process_rectangle_shapes(self):
for single_shape in self.json_data['shapes']:
bbox_class = single_shape['label']
xmin = single_shape['points'][0][0]
ymin = single_shape['points'][0][1]
xmax = single_shape['points'][1][0]
ymax = single_shape['points'][1][1]
self.coordis.append([xmin, ymin, xmax, ymax, bbox_class])
def process_polygon_shapes(self):
for single_shape in self.json_data['shapes']:
bbox_class = single_shape['label']
temp_points = []
for couple_point in single_shape['points']:
x = float(couple_point[0])
y = float(couple_point[1])
temp_points.append([x, y])
temp_points = np.array(temp_points)
xmin, ymin = temp_points.min(axis=0)
xmax, ymax = temp_points.max(axis=0)
self.coordis.append([xmin, ymin, xmax, ymax, bbox_class])
def get_width_height(self):
return self.width, self.height
def get_filename(self):
return self.filename
def get_coordis(self):
return self.coordis
然后创建create_xml_anno.py文件：
from xml.dom.minidom import Document
class CreateAnno:
def __init__(self, ):
self.doc = Document() # 创建DOM文档对象
self.anno = self.doc.createElement('annotation') # 创建根元素
self.doc.appendChild(self.anno)
self.add_folder()
self.add_path()
self.add_source()
self.add_segmented()
# self.add_filename()
# self.add_pic_size(width_text_str=str(width), height_text_str=str(height), depth_text_str=str(depth))
def add_folder(self, floder_text_str='JPEGImages'):
floder = self.doc.createElement('floder') ##建立自己的开头
floder_text = self.doc.createTextNode(floder_text_str) ##建立自己的文本信息
floder.appendChild(floder_text) ##自己的内容
self.anno.appendChild(floder)
def add_filename(self, filename_text_str='00000.jpg'):
filename = self.doc.createElement('filename')
filename_text = self.doc.createTextNode(filename_text_str)
filename.appendChild(filename_text)
self.anno.appendChild(filename)
def add_path(self, path_text_str="None"):
path = self.doc.createElement('path')
path_text = self.doc.createTextNode(path_text_str)
path.appendChild(path_text)
self.anno.appendChild(path)
def add_source(self, database_text_str="Unknow"):
source = self.doc.createElement('source')
database = self.doc.createElement('database')
database_text = self.doc.createTextNode(database_text_str) # 元素内容写入
database.appendChild(database_text)
source.appendChild(database)
self.anno.appendChild(source)
def add_pic_size(self, width_text_str="0", height_text_str="0", depth_text_str="3"):
size = self.doc.createElement('size')
width = self.doc.createElement('width')
width_text = self.doc.createTextNode(width_text_str) # 元素内容写入
width.appendChild(width_text)
size.appendChild(width)
height = self.doc.createElement('height')
height_text = self.doc.createTextNode(height_text_str)
height.appendChild(height_text)
size.appendChild(height)
depth = self.doc.createElement('depth')
depth_text = self.doc.createTextNode(depth_text_str)
depth.appendChild(depth_text)
size.appendChild(depth)
self.anno.appendChild(size)
def add_segmented(self, segmented_text_str="0"):
segmented = self.doc.createElement('segmented')
segmented_text = self.doc.createTextNode(segmented_text_str)
segmented.appendChild(segmented_text)
self.anno.appendChild(segmented)
def add_object(self,
name_text_str="None",
xmin_text_str="0",
ymin_text_str="0",
xmax_text_str="0",
ymax_text_str="0",
pose_text_str="Unspecified",
truncated_text_str="0",
difficult_text_str="0"):
object = self.doc.createElement('object')
name = self.doc.createElement('name')
name_text = self.doc.createTextNode(name_text_str)
name.appendChild(name_text)
object.appendChild(name)
pose = self.doc.createElement('pose')
pose_text = self.doc.createTextNode(pose_text_str)
pose.appendChild(pose_text)
object.appendChild(pose)
truncated = self.doc.createElement('truncated')
truncated_text = self.doc.createTextNode(truncated_text_str)
truncated.appendChild(truncated_text)
object.appendChild(truncated)
difficult = self.doc.createElement('difficult')
difficult_text = self.doc.createTextNode(difficult_text_str)
difficult.appendChild(difficult_text)
object.appendChild(difficult)
bndbox = self.doc.createElement('bndbox')
xmin = self.doc.createElement('xmin')
xmin_text = self.doc.createTextNode(xmin_text_str)
xmin.appendChild(xmin_text)
bndbox.appendChild(xmin)
ymin = self.doc.createElement('ymin')
ymin_text = self.doc.createTextNode(ymin_text_str)
ymin.appendChild(ymin_text)
bndbox.appendChild(ymin)
xmax = self.doc.createElement('xmax')
xmax_text = self.doc.createTextNode(xmax_text_str)
xmax.appendChild(xmax_text)
bndbox.appendChild(xmax)
ymax = self.doc.createElement('ymax')
ymax_text = self.doc.createTextNode(ymax_text_str)
ymax.appendChild(ymax_text)
bndbox.appendChild(ymax)
object.appendChild(bndbox)
self.anno.appendChild(object)
def get_anno(self):
return self.anno
def get_doc(self):
return self.doc
def save_doc(self, save_path):
with open(save_path, "w") as f:
self.doc.writexml(f, indent='\t', newl='\n', addindent='\t', encoding='utf-8')
最后创建main.py文件，当然这个你也可以改名字，至于上面两个.py文件我个人是建议新手就不要操作了：
import os
from tqdm import tqdm
from read_json_anno import ReadAnno
from create_xml_anno import CreateAnno
def json_transform_xml(json_path, xml_path, process_mode="rectangle"):
json_path = json_path
json_anno = ReadAnno(json_path, process_mode=process_mode)
width, height = json_anno.get_width_height()
filename = json_anno.get_filename()
coordis = json_anno.get_coordis()
xml_anno = CreateAnno()
xml_anno.add_filename(filename)
xml_anno.add_pic_size(width_text_str=str(width), height_text_str=str(height), depth_text_str=str(3))
for xmin, ymin, xmax, ymax, label in coordis:
xml_anno.add_object(name_text_str=str(label),
xmin_text_str=str(int(xmin)),
ymin_text_str=str(int(ymin)),
xmax_text_str=str(int(xmax)),
ymax_text_str=str(int(ymax)))
xml_anno.save_doc(xml_path)
if __name__ == "__main__":
root_json_dir = r"你的图片和json的路径，不要将图片与json分到两个文件夹"
root_save_xml_dir = r"你的xml想存到什么位置"
for json_filename in tqdm(os.listdir(root_json_dir)):
if json_filename.split('.')[-1]=='json':
json_path = os.path.join(root_json_dir, json_filename)
save_xml_path = os.path.join(root_save_xml_dir, json_filename.replace(".json", ".xml"))
json_transform_xml(json_path, save_xml_path, process_mode="rectangle") # labelme原数据的标注方式(矩形rectangle和多边形polygon)
当你运行完main.py这个文件你就可以得到xml文件。
下面这步是数据增强，这段代码的增强也是效果极好的，不过只能使用xml文件进行增强，这也是博主为什么不辞辛苦的将json转到xml而不是直接变成txt，这段代码你要改的地方在if __name__ == '__main__':这个位置：
# -*- coding=utf-8 -*-
# 裁剪(需改变bbox)
# 2. 平移(需改变bbox)
# 3. 改变亮度
# 4. 加噪声
# 5. 旋转角度(需要改变bbox)
# 6. 镜像(需要改变bbox)
# 7. cutout
# 注意:
# random.seed(),相同的seed,产生的随机数是一样
import time
import random
import copy
import cv2
import os
import math
import numpy as np
from skimage.util import random_noise
from lxml import etree, objectify
import xml.etree.ElementTree as ET
import argparse
# 显示图片
def show_pic(img, bboxes=None):
'''
输入:
img:图像array
bboxes:图像的所有boudning box list, 格式为[[x_min, y_min, x_max, y_max]....]
names:每个box对应的名称
'''
for i in range(len(bboxes)):
bbox = bboxes[i]
x_min = bbox[0]
y_min = bbox[1]
x_max = bbox[2]
y_max = bbox[3]
cv2.rectangle(img, (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 3)
cv2.namedWindow('pic', 0) # 1表示原图
cv2.moveWindow('pic', 0, 0)
cv2.resizeWindow('pic', 1200, 800) # 可视化的图片大小
cv2.imshow('pic', img)
cv2.waitKey(0)
cv2.destroyAllWindows()
# 图像均为cv2读取
class DataAugmentForObjectDetection():
def __init__(self, rotation_rate=0.5, max_rotation_angle=5,
crop_rate=0.5, shift_rate=0.5, change_light_rate=0.5,
add_noise_rate=0.5, flip_rate=0.5,
cutout_rate=0.5, cut_out_length=50, cut_out_holes=1, cut_out_threshold=0.5,
is_addNoise=True, is_changeLight=True, is_cutout=False, is_rotate_img_bbox=True,
is_crop_img_bboxes=True, is_shift_pic_bboxes=True, is_filp_pic_bboxes=True):
# 配置各个操作的属性
self.rotation_rate = rotation_rate
self.max_rotation_angle = max_rotation_angle
self.crop_rate = crop_rate
self.shift_rate = shift_rate
self.change_light_rate = change_light_rate
self.add_noise_rate = add_noise_rate
self.flip_rate = flip_rate
self.cutout_rate = cutout_rate
self.cut_out_length = cut_out_length
self.cut_out_holes = cut_out_holes
self.cut_out_threshold = cut_out_threshold
# 是否使用某种增强方式
self.is_addNoise = is_addNoise
self.is_changeLight = is_changeLight
self.is_cutout = is_cutout
self.is_rotate_img_bbox = is_rotate_img_bbox
self.is_crop_img_bboxes = is_crop_img_bboxes
self.is_shift_pic_bboxes = is_shift_pic_bboxes
self.is_filp_pic_bboxes = is_filp_pic_bboxes
# 加噪声
def _addNoise(self, img):
'''
输入:
img:图像array
输出:
加噪声后的图像array,由于输出的像素是在[0,1]之间,所以得乘以255
'''
# return cv2.GaussianBlur(img, (11, 11), 0)
np.random.seed(int(time.time()))
return random_noise(img, mode='gaussian', clip=True) * 255
# 调整亮度
def _changeLight(self, img):
alpha = random.uniform(0.35, 1)
blank = np.zeros(img.shape, img.dtype)
return cv2.addWeighted(img, alpha, blank, 1 - alpha, 0)
# cutout
def _cutout(self, img, bboxes, length=100, n_holes=1, threshold=0.5):
'''
原版本：https://github.com/uoguelph-mlrg/Cutout/blob/master/util/cutout.py
Randomly mask out one or more patches from an image.
Args:
img : a 3D numpy array,(h,w,c)
bboxes : 框的坐标
n_holes (int): Number of patches to cut out of each image.
length (int): The length (in pixels) of each square patch.
'''
def cal_iou(boxA, boxB):
'''
boxA, boxB为两个框，返回iou
boxB为bouding box
'''
# determine the (x, y)-coordinates of the intersection rectangle
xA = max(boxA[0], boxB[0])
yA = max(boxA[1], boxB[1])
xB = min(boxA[2], boxB[2])
yB = min(boxA[3], boxB[3])
if xB <= xA or yB <= yA:
return 0.0
# compute the area of intersection rectangle
interArea = (xB - xA + 1) * (yB - yA + 1)
# compute the area of both the prediction and ground-truth
# rectangles
boxAArea = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1)
boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)
iou = interArea / float(boxBArea)
return iou
# 得到h和w
if img.ndim == 3:
h, w, c = img.shape
else:
_, h, w, c = img.shape
mask = np.ones((h, w, c), np.float32)
for n in range(n_holes):
chongdie = True # 看切割的区域是否与box重叠太多
while chongdie:
y = np.random.randint(h)
x = np.random.randint(w)
y1 = np.clip(y - length // 2, 0,
h) # numpy.clip(a, a_min, a_max, out=None), clip这个函数将将数组中的元素限制在a_min, a_max之间，大于a_max的就使得它等于 a_max，小于a_min,的就使得它等于a_min
y2 = np.clip(y + length // 2, 0, h)
x1 = np.clip(x - length // 2, 0, w)
x2 = np.clip(x + length // 2, 0, w)
chongdie = False
for box in bboxes:
if cal_iou([x1, y1, x2, y2], box) > threshold:
chongdie = True
break
mask[y1: y2, x1: x2, :] = 0.
img = img * mask
return img
# 旋转
def _rotate_img_bbox(self, img, bboxes, angle=5, scale=1.):
'''
参考:https://blog.csdn.net/u014540717/article/details/53301195crop_rate
输入:
img:图像array,(h,w,c)
bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
angle:旋转角度
scale:默认1
输出:
rot_img:旋转后的图像array
rot_bboxes:旋转后的boundingbox坐标list
'''
# ---------------------- 旋转图像 ----------------------
w = img.shape[1]
h = img.shape[0]
# 角度变弧度
rangle = np.deg2rad(angle) # angle in radians
# now calculate new image width and height
nw = (abs(np.sin(rangle) * h) + abs(np.cos(rangle) * w)) * scale
nh = (abs(np.cos(rangle) * h) + abs(np.sin(rangle) * w)) * scale
# ask OpenCV for the rotation matrix
rot_mat = cv2.getRotationMatrix2D((nw * 0.5, nh * 0.5), angle, scale)
# calculate the move from the old center to the new center combined
# with the rotation
rot_move = np.dot(rot_mat, np.array([(nw - w) * 0.5, (nh - h) * 0.5, 0]))
# the move only affects the translation, so update the translation
rot_mat[0, 2] += rot_move[0]
rot_mat[1, 2] += rot_move[1]
# 仿射变换
rot_img = cv2.warpAffine(img, rot_mat, (int(math.ceil(nw)), int(math.ceil(nh))), flags=cv2.INTER_LANCZOS4)
# ---------------------- 矫正bbox坐标 ----------------------
# rot_mat是最终的旋转矩阵
# 获取原始bbox的四个中点，然后将这四个点转换到旋转后的坐标系下
rot_bboxes = list()
for bbox in bboxes:
xmin = bbox[0]
ymin = bbox[1]
xmax = bbox[2]
ymax = bbox[3]
point1 = np.dot(rot_mat, np.array([(xmin + xmax) / 2, ymin, 1]))
point2 = np.dot(rot_mat, np.array([xmax, (ymin + ymax) / 2, 1]))
point3 = np.dot(rot_mat, np.array([(xmin + xmax) / 2, ymax, 1]))
point4 = np.dot(rot_mat, np.array([xmin, (ymin + ymax) / 2, 1]))
# 合并np.array
concat = np.vstack((point1, point2, point3, point4))
# 改变array类型
concat = concat.astype(np.int32)
# 得到旋转后的坐标
rx, ry, rw, rh = cv2.boundingRect(concat)
rx_min = rx
ry_min = ry
rx_max = rx + rw
ry_max = ry + rh
# 加入list中
rot_bboxes.append([rx_min, ry_min, rx_max, ry_max])
return rot_img, rot_bboxes
# 裁剪
def _crop_img_bboxes(self, img, bboxes):
'''
裁剪后的图片要包含所有的框
输入:
img:图像array
bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
输出:
crop_img:裁剪后的图像array
crop_bboxes:裁剪后的bounding box的坐标list
'''
# ---------------------- 裁剪图像 ----------------------
w = img.shape[1]
h = img.shape[0]
x_min = w # 裁剪后的包含所有目标框的最小的框
x_max = 0
y_min = h
y_max = 0
for bbox in bboxes:
x_min = min(x_min, bbox[0])
y_min = min(y_min, bbox[1])
x_max = max(x_max, bbox[2])
y_max = max(y_max, bbox[3])
d_to_left = x_min # 包含所有目标框的最小框到左边的距离
d_to_right = w - x_max # 包含所有目标框的最小框到右边的距离
d_to_top = y_min # 包含所有目标框的最小框到顶端的距离
d_to_bottom = h - y_max # 包含所有目标框的最小框到底部的距离
# 随机扩展这个最小框
crop_x_min = int(x_min - random.uniform(0, d_to_left))
crop_y_min = int(y_min - random.uniform(0, d_to_top))
crop_x_max = int(x_max + random.uniform(0, d_to_right))
crop_y_max = int(y_max + random.uniform(0, d_to_bottom))
# 随机扩展这个最小框 , 防止别裁的太小
# crop_x_min = int(x_min - random.uniform(d_to_left//2, d_to_left))
# crop_y_min = int(y_min - random.uniform(d_to_top//2, d_to_top))
# crop_x_max = int(x_max + random.uniform(d_to_right//2, d_to_right))
# crop_y_max = int(y_max + random.uniform(d_to_bottom//2, d_to_bottom))
# 确保不要越界
crop_x_min = max(0, crop_x_min)
crop_y_min = max(0, crop_y_min)
crop_x_max = min(w, crop_x_max)
crop_y_max = min(h, crop_y_max)
crop_img = img[crop_y_min:crop_y_max, crop_x_min:crop_x_max]
# ---------------------- 裁剪boundingbox ----------------------
# 裁剪后的boundingbox坐标计算
crop_bboxes = list()
for bbox in bboxes:
crop_bboxes.append([bbox[0] - crop_x_min, bbox[1] - crop_y_min, bbox[2] - crop_x_min, bbox[3] - crop_y_min])
return crop_img, crop_bboxes
# 平移
def _shift_pic_bboxes(self, img, bboxes):
'''
参考:https://blog.csdn.net/sty945/article/details/79387054
平移后的图片要包含所有的框
输入:
img:图像array
bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
输出:
shift_img:平移后的图像array
shift_bboxes:平移后的bounding box的坐标list
'''
# ---------------------- 平移图像 ----------------------
w = img.shape[1]
h = img.shape[0]
x_min = w # 裁剪后的包含所有目标框的最小的框
x_max = 0
y_min = h
y_max = 0
for bbox in bboxes:
x_min = min(x_min, bbox[0])
y_min = min(y_min, bbox[1])
x_max = max(x_max, bbox[2])
y_max = max(y_max, bbox[3])
d_to_left = x_min # 包含所有目标框的最大左移动距离
d_to_right = w - x_max # 包含所有目标框的最大右移动距离
d_to_top = y_min # 包含所有目标框的最大上移动距离
d_to_bottom = h - y_max # 包含所有目标框的最大下移动距离
x = random.uniform(-(d_to_left - 1) / 3, (d_to_right - 1) / 3)
y = random.uniform(-(d_to_top - 1) / 3, (d_to_bottom - 1) / 3)
M = np.float32([[1, 0, x], [0, 1, y]]) # x为向左或右移动的像素值,正为向右负为向左; y为向上或者向下移动的像素值,正为向下负为向上
shift_img = cv2.warpAffine(img, M, (img.shape[1], img.shape[0]))
# ---------------------- 平移boundingbox ----------------------
shift_bboxes = list()
for bbox in bboxes:
shift_bboxes.append([bbox[0] + x, bbox[1] + y, bbox[2] + x, bbox[3] + y])
return shift_img, shift_bboxes
# 镜像
def _filp_pic_bboxes(self, img, bboxes):
'''
参考:https://blog.csdn.net/jningwei/article/details/78753607
平移后的图片要包含所有的框
输入:
img:图像array
bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
输出:
flip_img:平移后的图像array
flip_bboxes:平移后的bounding box的坐标list
'''
# ---------------------- 翻转图像 ----------------------
flip_img = copy.deepcopy(img)
h, w, _ = img.shape
sed = random.random()
if 0 < sed < 0.33: # 0.33的概率水平翻转，0.33的概率垂直翻转,0.33是对角反转
flip_img = cv2.flip(flip_img, 0) # _flip_x
inver = 0
elif 0.33 < sed < 0.66:
flip_img = cv2.flip(flip_img, 1) # _flip_y
inver = 1
else:
flip_img = cv2.flip(flip_img, -1) # flip_x_y
inver = -1
# ---------------------- 调整boundingbox ----------------------
flip_bboxes = list()
for box in bboxes:
x_min = box[0]
y_min = box[1]
x_max = box[2]
y_max = box[3]
if inver == 0:
#0：垂直翻转
flip_bboxes.append([x_min, h - y_max, x_max, h - y_min])
elif inver == 1:
# 1：水平翻转
flip_bboxes.append([w - x_max, y_min, w - x_min, y_max])
elif inver == -1:
# -1：水平垂直翻转
flip_bboxes.append([w - x_max, h - y_max, w - x_min, h - y_min])
return flip_img, flip_bboxes
# 图像增强方法
def dataAugment(self, img, bboxes):
'''
图像增强
输入:
img:图像array
bboxes:该图像的所有框坐标
输出:
img:增强后的图像
bboxes:增强后图片对应的box
'''
change_num = 0 # 改变的次数
# print('------')
while change_num < 1: # 默认至少有一种数据增强生效
if self.is_rotate_img_bbox:
if random.random() > self.rotation_rate: # 旋转
change_num += 1
angle = random.uniform(-self.max_rotation_angle, self.max_rotation_angle)
scale = random.uniform(0.7, 0.8)
img, bboxes = self._rotate_img_bbox(img, bboxes, angle, scale)
if self.is_shift_pic_bboxes:
if random.random() < self.shift_rate: # 平移
change_num += 1
img, bboxes = self._shift_pic_bboxes(img, bboxes)
if self.is_changeLight:
if random.random() > self.change_light_rate: # 改变亮度
change_num += 1
img = self._changeLight(img)
if self.is_addNoise:
if random.random() < self.add_noise_rate: # 加噪声
change_num += 1
img = self._addNoise(img)
if self.is_cutout:
if random.random() < self.cutout_rate: # cutout
change_num += 1
img = self._cutout(img, bboxes, length=self.cut_out_length, n_holes=self.cut_out_holes,
threshold=self.cut_out_threshold)
if self.is_filp_pic_bboxes:
if random.random() < self.flip_rate: # 翻转
change_num += 1
img, bboxes = self._filp_pic_bboxes(img, bboxes)
return img, bboxes
# xml解析工具
class ToolHelper():
# 从xml文件中提取bounding box信息, 格式为[[x_min, y_min, x_max, y_max, name]]
def parse_xml(self, path):
'''
输入：
xml_path: xml的文件路径
输出：
从xml文件中提取bounding box信息, 格式为[[x_min, y_min, x_max, y_max, name]]
'''
tree = ET.parse(path)
root = tree.getroot()
objs = root.findall('object')
coords = list()
for ix, obj in enumerate(objs):
name = obj.find('name').text
box = obj.find('bndbox')
x_min = int(box[0].text)
y_min = int(box[1].text)
x_max = int(box[2].text)
y_max = int(box[3].text)
coords.append([x_min, y_min, x_max, y_max, name])
return coords
# 保存图片结果
def save_img(self, file_name, save_folder, img):
cv2.imwrite(os.path.join(save_folder, file_name), img)
# 保持xml结果
def save_xml(self, file_name, save_folder, img_info, height, width, channel, bboxs_info):
'''
:param file_name:文件名
:param save_folder:#保存的xml文件的结果
:param height:图片的信息
:param width:图片的宽度
:param channel:通道
:return:
'''
folder_name, img_name = img_info # 得到图片的信息
E = objectify.ElementMaker(annotate=False)
anno_tree = E.annotation(
E.folder(folder_name),
E.filename(img_name),
E.path(os.path.join(folder_name, img_name)),
E.source(
E.database('Unknown'),
),
E.size(
E.width(width),
E.height(height),
E.depth(channel)
),
E.segmented(0),
)
labels, bboxs = bboxs_info # 得到边框和标签信息
for label, box in zip(labels, bboxs):
anno_tree.append(
E.object(
E.name(label),
E.pose('Unspecified'),
E.truncated('0'),
E.difficult('0'),
E.bndbox(
E.xmin(box[0]),
E.ymin(box[1]),
E.xmax(box[2]),
E.ymax(box[3])
)
))
etree.ElementTree(anno_tree).write(os.path.join(save_folder, file_name), pretty_print=True)
if __name__ == '__main__':
need_aug_num = 5 # 每张图片需要增强的次数
is_endwidth_dot = True # 文件是否以.jpg或者png结尾
dataAug = DataAugmentForObjectDetection() # 数据增强工具类
toolhelper = ToolHelper() # 工具
# 获取相关参数
parser = argparse.ArgumentParser()
parser.add_argument('--source_img_path', type=str, default=r'你增强之前照片的绝对位置')
parser.add_argument('--source_xml_path', type=str, default=r'你增强之前xml的绝对位置')
parser.add_argument('--save_img_path', type=str, default=r'你增强之后照片的绝对位置')
parser.add_argument('--save_xml_path', type=str, default=r'你增强之后xml的绝对位置')
args = parser.parse_args()
source_img_path = args.source_img_path # 图片原始位置
source_xml_path = args.source_xml_path # xml的原始位置
save_img_path = args.save_img_path # 图片增强结果保存文件
save_xml_path = args.save_xml_path # xml增强结果保存文件
# 如果保存文件夹不存在就创建
if not os.path.exists(save_img_path):
os.mkdir(save_img_path)
if not os.path.exists(save_xml_path):
os.mkdir(save_xml_path)
for parent, _, files in os.walk(source_img_path):
files.sort()
for file in files:
cnt = 0
pic_path = os.path.join(parent, file)
xml_path = os.path.join(source_xml_path, file[:-4] + '.xml')
values = toolhelper.parse_xml(xml_path) # 解析得到box信息，格式为[[x_min,y_min,x_max,y_max,name]]
coords = [v[:4] for v in values] # 得到框
labels = [v[-1] for v in values] # 对象的标签
# 如果图片是有后缀的
if is_endwidth_dot:
# 找到文件的最后名字
dot_index = file.rfind('.')
_file_prefix = file[:dot_index] # 文件名的前缀
_file_suffix = file[dot_index:] # 文件名的后缀
img = cv2.imread(pic_path)
# show_pic(img, coords) # 显示原图
while cnt < need_aug_num: # 继续增强
auged_img, auged_bboxes = dataAug.dataAugment(img, coords)
auged_bboxes_int = np.array(auged_bboxes).astype(np.int32)
height, width, channel = auged_img.shape # 得到图片的属性
img_name = '{}_{}{}'.format(_file_prefix, cnt + 1, _file_suffix) # 图片保存的信息
toolhelper.save_img(img_name, save_img_path,
auged_img) # 保存增强图片
toolhelper.save_xml('{}_{}.xml'.format(_file_prefix, cnt + 1),
save_xml_path, (save_img_path, img_name), height, width, channel,
(labels, auged_bboxes_int)) # 保存xml文件
# show_pic(auged_img, auged_bboxes) # 强化后的图
print(img_name)
cnt += 1 # 继续增强下一张
下面我们要开始将xml转换为txt文件：
import xml.etree.ElementTree as ET
import os
def convert(size, box):
x_center = (box[0] + box[1]) / 2.0
y_center = (box[2] + box[3]) / 2.0
# 分别计算纵坐标和横坐标的中心点
x = x_center / size[0]
y = y_center / size[1]
w = (box[1] - box[0]) / size[0]
h = (box[3] - box[2]) / size[1]
# print(x, y, w, h)
return (x, y, w, h)
def convert_annotation(xml_files_path, save_txt_files_path, classes):
xml_files = os.listdir(xml_files_path)
print(xml_files)
for xml_name in xml_files:
print(xml_name)
xml_file = os.path.join(xml_files_path, xml_name)
out_txt_path = os.path.join(save_txt_files_path, xml_name.split('.')[0] + '.txt')
out_txt_f = open(out_txt_path, 'w')
tree = ET.parse(xml_file)
root = tree.getroot()
size = root.find('size')
w = int(size.find('width').text)
h = int(size.find('height').text)
for obj in root.iter('object'):
difficult = obj.find('difficult').text
cls = obj.find('name').text
if cls not in classes or int(difficult) == 1:
continue
cls_id = classes.index(cls)
xmlbox = obj.find('bndbox')
b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
float(xmlbox.find('ymax').text))
# b=(xmin, xmax, ymin, ymax)
print(w, h, b)
bb = convert((w, h), b)
out_txt_f.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
if __name__ == "__main__":
# 1、指定yolo类别
classes1 = ["A", "B"， ...] #这个地方填入你lable的名字，
# 2、voc格式的xml标签文件路径
xml_files1 = r'你的绝对路径'
# 3、转化为yolo格式的txt标签文件存储路径
save_txt_files1 = r'你的绝对路径'
convert_annotation(xml_files1, save_txt_files1, classes1)
with open(save_txt_files1 + '/classes.txt', 'w') as file:
for class_name in classes1:
file.write(class_name + '\n')
4、当你完成这些我们就完成一大半的任务了，之后我们要开始将划分数据集了，这里我只划分了train和val，占比是8：2，之后也可以在代码中更改：
import os
import shutil
import random
def split_dataset(image_dir, label_dir, output_dir, train_ratio=0.8):
# 创建输出目录结构
os.makedirs(os.path.join(output_dir, 'train', 'images'), exist_ok=True)
os.makedirs(os.path.join(output_dir, 'train', 'labels'), exist_ok=True)
os.makedirs(os.path.join(output_dir, 'val', 'images'), exist_ok=True)
os.makedirs(os.path.join(output_dir, 'val', 'labels'), exist_ok=True)
# 获取所有图像文件
images = [f for f in os.listdir(image_dir) if f.endswith('.bmp')]
random.shuffle(images)
# 计算训练集和验证集的大小
train_size = int(len(images) * train_ratio)
# 分配图像到训练集和验证集
train_images = images[:train_size]
val_images = images[train_size:]
# 复制图像和标注文件到相应目录
for image_set, split in zip([train_images, val_images], ['train', 'val']):
for image in image_set:
# 复制图像文件
shutil.copy(os.path.join(image_dir, image), os.path.join(output_dir, split, 'images', image))
# 复制对应的标注文件
label_file = image.replace('.bmp', '.txt')
shutil.copy(os.path.join(label_dir, label_file), os.path.join(output_dir, split, 'labels', label_file))
# 示例用法
image_dir = r'你图片的绝对路径'
label_dir = r'你txt文件的绝对路径'
output_dir = '你要生成dataset文件夹的路径'
split_dataset(image_dir, label_dir, output_dir)
下面就是创建你的yaml文件，这个就没有生成的代码了，是要自己打的下面我会给你模板：
# 数据集配置文件
train: 你的训练集的路径
val: 你的测试集的路径
# 类别数
nc: 2
# 类别名称
names:
0: A
1: B
当你完成上面这些你就可以开始训练了具体的可以参考：
YOLOv8教程系列：一、使用自定义数据集训练YOLOv8模型（详细版教程，你只看一篇-＞调参攻略），包含环境搭建/数据准备/模型训练/预测/验证/导出等_yolov8训练自己的数据集-CSDN博客
3824
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
