YOLO数据集脚本制作
关于YOLO训练数据集脚本制作
任务分析
本章节以现有的项目为基础编写,从以下几个步骤实现数据集制作:
使用OS模块一键读取当前列表中的Json文件,限定读取文件类型为'.json'格式
读取当前的Json文件,将Json['shapes']列表内的所有点位都提取出,并将其转换为Float格式以供我们使用
读取点位的同时将Json文件自带的图片Base64编码提取,并将编码转换为1.二进制数据 2.UTF-8格式
我们需要根据YOLO数据集的要求,对点位进行归一化
对于高和宽和点位的归一化,正常除以图像宽高即可
对于裁剪过后的图像,计算方式为:
宽:(中点坐标减去x最小坐标,即为准确的中点坐标) / (x最大坐标-x最小坐标,即图像的宽)
高:(中点坐标减去y最小坐标,即为准确的中点坐标) / (y最大坐标-y最小坐标,即图像的宽) 此两句对裁剪后的图像恒成立
点位: 点位乘以图像改动的比例 / 图像的分辨率
对点位进行归一化后,我们需要将图片的分辨率修改为你在YOLO训练中设定的分辨率,如(640x640),最好是正方形的对称图形式分辨率
分析你训练使用的模型
yolov8n,该模型需要的格式为:class(类) 矩形左上角X 矩形左上角Y 矩形右下角X 矩形右下角Y
yolov8n-pose,该模型需要的格式为:class(类) 矩形中点x 矩形中点y 矩形宽 矩形高 点位X 点位Y 点位1X 点位1Y ... 点位nX 点位nY
开始训练,并得出结果
准备阶段
所需的库有:
os库
Numpy库
json库
cv2库
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
获取矩形的思路是:
根据左下角的x点和右下角的x点相减,我们得出灯牌的宽,宽/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的图片之上,可以在最少的修改之内进行应用,若感兴趣可以看我的下一篇文章。