YOLO数据集脚本制作

Author Avatar
根骨何尽
发表:2024-03-23 17:05:12
修改:2024-03-28 14:52:06

关于YOLO训练数据集脚本制作

任务分析

本章节以现有的项目为基础编写,从以下几个步骤实现数据集制作:

  1. 使用OS模块一键读取当前列表中的Json文件,限定读取文件类型为'.json'格式

  2. 读取当前的Json文件,将Json['shapes']列表内的所有点位都提取出,并将其转换为Float格式以供我们使用

  3. 读取点位的同时将Json文件自带的图片Base64编码提取,并将编码转换为1.二进制数据 2.UTF-8格式

我们需要根据YOLO数据集的要求,对点位进行归一化

对于高和宽和点位的归一化,正常除以图像宽高即可

对于裁剪过后的图像,计算方式为:

宽:(中点坐标减去x最小坐标,即为准确的中点坐标) / (x最大坐标-x最小坐标,即图像的宽)

高:(中点坐标减去y最小坐标,即为准确的中点坐标) / (y最大坐标-y最小坐标,即图像的宽) 此两句对裁剪后的图像恒成立

点位: 点位乘以图像改动的比例 / 图像的分辨率

  1. 对点位进行归一化后,我们需要将图片的分辨率修改为你在YOLO训练中设定的分辨率,如(640x640),最好是正方形的对称图形式分辨率

  2. 分析你训练使用的模型

yolov8n,该模型需要的格式为:class(类) 矩形左上角X 矩形左上角Y 矩形右下角X 矩形右下角Y

yolov8n-pose,该模型需要的格式为:class(类) 矩形中点x 矩形中点y 矩形宽 矩形高 点位X 点位Y 点位1X 点位1Y ... 点位nX 点位nY

  1. 开始训练,并得出结果

准备阶段

所需的库有:

  1. os库

  2. Numpy库

  3. json库

  4. cv2库

  5. base64库

所需的JS文件:

开始编写代码:

首先:我们需要定义我们所要读取的Json文件地址,以及要保存的txt文件地址和images文件地址

import os
import json
import cv2
import numpy as np
import base

# json文件目录,在字符串前加r,为反义
json_path = r'd:\project\json'
txt_path = r'd:\project\txt'
images_path = r'd:\project\images'

开始读取json_path当中的文件,我们使用os模块

path_lib = os.listdir(json_path)
# 该代码返回一个当前文件夹下的所有文件的列表,我们使用for循环遍历
for paths in path_lib:
	# 首先判断这文件是不是json格式
	if os.path.splitext(paths)[1] == '.json':
		# 使用open函数打开Json文件,r为读取模式
		with open(os.path.join(path,paths),'r') as Js:
			# 使用json库中的json.load来读取Json
			js = json.load(Js)
			# 我们通过内遍历来获取json数据中的shapes列表中的points点位
			label_points = [shape['points'][0] for shape in js['shapes'] ]
			# 创建一个空列表来存储数据
			strs = []
			# 每个points列表含有x点和y点,我们对label_points进行遍历,在strs列表中添加(x,y)元组
			for x,y in label_points:
				strs.append((x,y))

此时我们的strs列表,通过这段处理,获取了3对点数,分别是灯牌的正上方、左下角、右下角

接着我们读取json文件下的图片编码

			# 该js对应的是上段代码中打开的json文件
			image_Data = js['imageData']
			# 这行代码使用Python的base64模块对image_path进行Base64解码
			image_bytes = base64.b64decode(image_Data)
			# np.frombuffer()负责将内存中的data数据转换为np.uint8格式,即UTF-8格式
			image_np = np.frombuffer(image_bytes, np.uint8)
			# 转换格式后使用opencv的imdecode函数,函数用于从内存缓存中读取图像数据并将其转换为图像格式
			img = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
			# 最后我们使用opencv的imwrite将图片写入地址
			cv2.imwrite(os.path.join(images_path, path).replace('.json', '.jpg'),img)

切记在写入图片的时候,格式为cv2.imwrite('图片路径+图片名',图片)

os.path.join(path,name),将路径和文件名结合,切记opencv不能读取或保存带中文字符的文件

将图片提取后,我们发现图片并没有640x640的分辨率

我们需要将宽和高都转换为640,需要一个比例值:

	# 首先获取img图片的shape,一个RGB的图片通常有3个属性[高,宽,颜色位数]
	shape = img.shape[:2]
	# 改变的高的比例
	height_resize_ratio = 640 / shape[0]
	# 改变的宽的比例
	width_resize_ratio = 640 / shape[1]
	# 接下来我们直接暴力修改这张图片的尺寸
	img = cv2.resize(img, (640, 640))

此时我们拥有3对点,和一张640x640分辨率的图片,因为对图片进行了修改,点位也要进行等比例放大

	# 创建一个新的列表,存储新的放大后的点位
	points = []
	for x,y in strs:
		x = x*width_resize_ratio
		y = y*height_resize_ratio
		points.append(x,y)

接下来我们对获取的三对点位进行归一化

因为我们要创建的图片分辨率为640x640,所以归一化公式为: x = x / 640 ; y = y / 640

但在这之前,我们还需要获取矩形中点和矩形的宽和高,在此项目中我们选择的模型为yolov8n-pose

获取矩形的思路是:

  1. 根据左下角的x点和右下角的x点相减,我们得出灯牌的宽,宽/2即是横向中点

  2. 通过正上方的y点和左下或右下的y点,我们得出灯牌的高,高/2即是纵向中点

points此时的值为[[点1x,点1y],[点2x,点2y],[点3x,点3y]]

	# 我们假定xc是灯牌的宽,yc是灯牌的高
	xc = points[2][0] - points[1][0]
	yc = points[1][1] - points[0][1]
	# 横坐标中心
	x_center = points[2][0] + xc / 2
	# 纵坐标中心
	y_center = points[0][1] + yc / 2
	# 我们规定矩形的中点
	square = (x_center,y_center)
	# 我们把获取的数据,都进行归一化
	xc = xc / 640
	yc = yc / 640
	square[0] = square[0] / 640
	square[1] = square[1] / 640

我们获取了矩形的中心点,矩形的宽、高以及3对点位,我们都进行归一化后,我们接下了要将其写入txt文件当中

	# 0为标签的类,需要添加在矩形的前方
	Strs = f'0 {x_center} {y_center} {width} {height} '
	for po_x,po_y in points:
		Strs = Strs + f'{str(po_x) {str(po_y} '
	# 我们将点位以class(类) 矩形中点x 矩形中点y 矩形宽 矩形高 点位X 点位Y 点位1X 点位1Y ... 点位nX 点位nY的形式,写入文件
	with open(os.path.join(txt_path,paths).replace('.json','.txt'),'w') as ftxt:
		ftxt.writelines(Strs)

如果代码报错,说明没有此文件夹,我们只需要只用os库的创建文件夹就可以

if not os.path.exists(images_path ): os.makedirs(images_path )

if not.os.path.exists(txt_path): os.makedirs(images_path)

当然制作数据集不只这一种方法,还可以将点位和图片直接附加在一张640x640的图片之上,可以在最少的修改之内进行应用,若感兴趣可以看我的下一篇文章。

评论