关键点数据增强

1.关键点数据增强

# 关键点数据增强
from PIL import Image, ImageDraw
import random
import json
from pathlib import Path# 创建一个黑色背景图像
width, height = 5000, 5000  # 图像宽度和高度
background_color = (0, 0, 0)  # 黑色填充# 随机分布图像
num_images = 1  # 要随机分布的图像数量
folder_path = Path("E:/2") # 测试图像目录
output_path = Path("E:/5") # 输出图像目录
for file in folder_path.rglob("*.jpg"):# eg: file = "目录名/123.jpg"，file_name = "123.jpg"file_name = file.nameimage_origin = Image.open(file)width_origin,height_origin = image_origin.sizefor _ in range(num_images):#随机选择图像的位置x = random.randint(0, width - width_origin)y = random.randint(0, height - height_origin)print(x,y)canvas = Image.new("RGB", (width,height), background_color) #新建一个mask，全黑填充canvas.paste(image_origin, (x,y)) #将原图从(x,y)处粘贴到mask上Path.mkdir(output_path, exist_ok=True)img_name = 'a' + '_' + file_name #改变增强后图片的名字canvas.save(output_path / img_name)jsonFile = file.with_suffix(".json")print(jsonFile)if Path.exists(jsonFile): #判断图片是否有对应的json文件print(f"找到{file}的Json文件")with open(jsonFile, "r", encoding="utf-8") as f:objectDict = json.load(f)objectDict["imageData"] =  None # 清空json文件里加密的imgdataobjectDict["imageHeight"] = heightobjectDict["imageWidth"] = widthjson_name = 'a' + '_' + jsonFile.name #改变增强后json文件的名字for i in range(len(objectDict["shapes"])):if objectDict["shapes"][i]["shape_type"] in ["rectangle","line"]: #矩形框、线段objectDict["shapes"][i]['points'][0][0]+=xobjectDict["shapes"][i]['points'][0][1]+=yobjectDict["shapes"][i]['points'][1][0]+=xobjectDict["shapes"][i]['points'][1][1]+=yif objectDict["shapes"][i]["shape_type"] in ["polygon"]: #多段线for polygonMat in objectDict["shapes"][i]['points']:polygonMat[0]+=xpolygonMat[1]+=yif objectDict["shapes"][i]["shape_type"] in ["point"]: #关键点objectDict["shapes"][i]['points'][0][0]+=xobjectDict["shapes"][i]['points'][0][1]+=ywith open(output_path / json_name, 'w',encoding='utf-8') as f:json.dump(objectDict, f)else:print("没有Json文件")

2.关键点可视化

# 可视化关键点位置
import cv2
import numpy as np
import json
import matplotlib.pyplot as plt# 载入图像
img_path = 'meter_6_25.jpg'
img_bgr = cv2.imread(img_path)# 载入labelme格式的json标注文件
labelme_path = 'meter_6_25.json'with open(labelme_path, 'r', encoding='utf-8') as f:labelme = json.load(f)# 查看标注信息  rectangle:矩形  point:点  polygon:多边形
# print(labelme.keys())
# dict_keys(['version', 'flags', 'shapes', 'imagePath', 'imageData', 'imageHeight', 'imageWidth'])
# print(labelme['shapes'])# <<<<<<<<<<<<<<<<<<可视化框（rectangle）标注>>>>>>>>>>>>>>>>>>>>>
# 框可视化配置
bbox_color = (255, 129, 0)           # 框的颜色
bbox_thickness = 5                   # 框的线宽# 框类别文字
bbox_labelstr = {'font_size':6,         # 字体大小'font_thickness':14,   # 字体粗细'offset_x':0,          # X 方向，文字偏移距离，向右为正'offset_y':-80,       # Y 方向，文字偏移距离，向下为正
}
# 画框
for each_ann in labelme['shapes']:  # 遍历每一个标注if each_ann['shape_type'] == 'rectangle':  # 筛选出框标注# 框的类别bbox_label = each_ann['label']# 框的两点坐标bbox_keypoints = each_ann['points']bbox_keypoint_A_xy = bbox_keypoints[0]bbox_keypoint_B_xy = bbox_keypoints[1]# 左上角坐标bbox_top_left_x = int(min(bbox_keypoint_A_xy[0], bbox_keypoint_B_xy[0]))bbox_top_left_y = int(min(bbox_keypoint_A_xy[1], bbox_keypoint_B_xy[1]))# 右下角坐标bbox_bottom_right_x = int(max(bbox_keypoint_A_xy[0], bbox_keypoint_B_xy[0]))bbox_bottom_right_y = int(max(bbox_keypoint_A_xy[1], bbox_keypoint_B_xy[1]))# 画矩形：画框img_bgr = cv2.rectangle(img_bgr, (bbox_top_left_x, bbox_top_left_y), (bbox_bottom_right_x, bbox_bottom_right_y),bbox_color, bbox_thickness)# 写框类别文字：图片，文字字符串，文字左上角坐标，字体，字体大小，颜色，字体粗细img_bgr = cv2.putText(img_bgr, bbox_label, (bbox_top_left_x + bbox_labelstr['offset_x'],bbox_top_left_y + bbox_labelstr['offset_y']),cv2.FONT_HERSHEY_SIMPLEX, bbox_labelstr['font_size'], bbox_color,bbox_labelstr['font_thickness'])# <<<<<<<<<<<<<<<<<<可视化关键点（keypoint）标注>>>>>>>>>>>>>>>>>>>>>
# 关键点的可视化配置
# 关键点配色
kpt_color_map = {'0': {'name': '0', 'color': [0, 0, 255], 'radius': 25, 'thickness':-1},'1': {'name': '1', 'color': [255, 0, 0], 'radius': 25, 'thickness':-1},'2': {'name': '2', 'color': [255, 0, 0], 'radius': 25, 'thickness':-1},'3': {'name': '3', 'color': [0, 255, 0], 'radius': 25, 'thickness':-1},'4': {'name': '4', 'color': [0, 255, 0], 'radius': 25, 'thickness':-1},'5': {'name': '5', 'color': [193, 182, 255], 'radius': 25, 'thickness':-1},'6': {'name': '6', 'color': [193, 182, 255], 'radius': 25, 'thickness':-1},# '7': {'name': '7', 'color': [16, 144, 247], 'radius': 25},# '8': {'name': '8', 'color': [16, 144, 247], 'radius': 25},
}# 点类别文字
kpt_labelstr = {'font_size':4,             # 字体大小'font_thickness':12,       # 字体粗细'offset_x':30,             # X 方向，文字偏移距离，向右为正'offset_y':100,            # Y 方向，文字偏移距离，向下为正
}# 画点
for each_ann in labelme['shapes']:  # 遍历每一个标注if each_ann['shape_type'] == 'point':  # 筛选出关键点标注kpt_label = each_ann['label']  # 该点的类别# 该点的 XY 坐标kpt_xy = each_ann['points'][0]kpt_x, kpt_y = int(kpt_xy[0]), int(kpt_xy[1])# 该点的可视化配置kpt_color = kpt_color_map[kpt_label]['color']  # 颜色kpt_radius = kpt_color_map[kpt_label]['radius']  # 半径kpt_thickness = kpt_color_map[kpt_label]['thickness']  # 线宽（-1代表填充）# 画圆：画该关键点img_bgr = cv2.circle(img_bgr, (kpt_x, kpt_y), kpt_radius, kpt_color, kpt_thickness)# 写该点类别文字：图片，文字字符串，文字左上角坐标，字体，字体大小，颜色，字体粗细img_bgr = cv2.putText(img_bgr, kpt_label, (kpt_x + kpt_labelstr['offset_x'], kpt_y + kpt_labelstr['offset_y']),cv2.FONT_HERSHEY_SIMPLEX, kpt_labelstr['font_size'], kpt_color,kpt_labelstr['font_thickness'])# # <<<<<<<<<<<<<<<<<<可视化多段线（polygon）标注>>>>>>>>>>>>>>>>>>>>>
# # 多段线的可视化配置
# poly_color = (151, 57, 224)
# poly_thickness = 3
#
# poly_labelstr = {
#     'font_size':4,          # 字体大小
#     'font_thickness':12,    # 字体粗细
#     'offset_x':-200,        # X 方向，文字偏移距离，向右为正
#     'offset_y':0,           # Y 方向，文字偏移距离，向下为正
# }
#
# # 画多段线
# img_mask = np.ones(img_bgr.shape, np.uint8) #创建一个和img_bgr一样大小的黑色mask
#
# for each_ann in labelme['shapes']:  # 遍历每一个标注
#
#     if each_ann['shape_type'] == 'polygon':  # 筛选出多段线（polygon）标注
#
#         poly_label = each_ann['label']  # 该多段线的类别
#
#         poly_points = [np.array(each_ann['points'], np.int32).reshape((-1, 1, 2))]  #reshape后增加一个维度
#
#         # 该多段线平均 XY 坐标，用于放置多段线类别文字
#         x_mean = int(np.mean(poly_points[0][:, 0, :][:, 0])) #取出所有点的x坐标并求平均值
#         y_mean = int(np.mean(poly_points[0][:, 0, :][:, 1])) #取出所有点的y坐标并求平均值
#
#         # 画该多段线轮廓
#         img_bgr = cv2.polylines(img_bgr, poly_points, isClosed=True, color=poly_color, thickness=poly_thickness)
#
#         # 画该多段线内部填充
#         img_mask = cv2.fillPoly(img_mask, poly_points, color=poly_color) #填充的颜色为color=poly_color
#
#         # 写该多段线类别文字：图片，文字字符串，文字左上角坐标，字体，字体大小，颜色，字体粗细
#         img_bgr = cv2.putText(img_bgr, poly_label,
#                               (x_mean + poly_labelstr['offset_x'], y_mean + poly_labelstr['offset_y']),
#                               cv2.FONT_HERSHEY_SIMPLEX, poly_labelstr['font_size'], poly_color,
#                               poly_labelstr['font_thickness'])# opacity = 0.8 # 透明度，越大越接近原图
# img_bgr = cv2.addWeighted(img_bgr, opacity, img_mask, 1-opacity, 0)# 可视化
plt.imshow(img_bgr[:,:,::-1]) # 将bgr通道转换成rgb通道
plt.show()# 当前目录下保存可视化结果
cv2.imwrite('meter.jpg', img_bgr)

3.json2txt（用YOLOV8进行关键点训练）

#将坐标框、关键点、线段的json标注转换为txt
import os
import json
import shutil
import numpy as np
from tqdm import tqdm# 框的类别
bbox_class = {'meter3':0
}
# 关键点的类别，注意按顺序写
keypoint_class = ['0','1','2','3','4','5','6','7','8']path = 'E:/6' #json文件存放路径
save_folder='E:/7' #转换后的txt文件存放路径# 定义单个json文件的转换
def process_single_json(labelme_path, save_folder):with open(labelme_path, 'r', encoding='utf-8') as f:labelme = json.load(f)img_width = labelme['imageWidth']  # 图像宽度img_height = labelme['imageHeight']  # 图像高度# 生成 YOLO 格式的 txt 文件suffix = labelme_path.split('.')[-2]# print(suffix)yolo_txt_path = suffix + '.txt'# print(yolo_txt_path)with open(yolo_txt_path, 'w', encoding='utf-8') as f:for each_ann in labelme['shapes']:  # 遍历每个标注if each_ann['shape_type'] == 'rectangle':  # 每个框，在 txt 里写一行yolo_str = ''# 框的信息# 框的类别 IDbbox_class_id = bbox_class[each_ann['label']]yolo_str += '{} '.format(bbox_class_id)# 左上角和右下角的 XY 像素坐标bbox_top_left_x = int(min(each_ann['points'][0][0], each_ann['points'][1][0]))bbox_bottom_right_x = int(max(each_ann['points'][0][0], each_ann['points'][1][0]))bbox_top_left_y = int(min(each_ann['points'][0][1], each_ann['points'][1][1]))bbox_bottom_right_y = int(max(each_ann['points'][0][1], each_ann['points'][1][1]))# 框中心点的 XY 像素坐标bbox_center_x = int((bbox_top_left_x + bbox_bottom_right_x) / 2)bbox_center_y = int((bbox_top_left_y + bbox_bottom_right_y) / 2)# 框宽度bbox_width = bbox_bottom_right_x - bbox_top_left_x# 框高度bbox_height = bbox_bottom_right_y - bbox_top_left_y# 框中心点归一化坐标bbox_center_x_norm = bbox_center_x / img_widthbbox_center_y_norm = bbox_center_y / img_height# 框归一化宽度bbox_width_norm = bbox_width / img_width# 框归一化高度bbox_height_norm = bbox_height / img_heightyolo_str += '{:.5f} {:.5f} {:.5f} {:.5f} '.format(bbox_center_x_norm, bbox_center_y_norm,bbox_width_norm, bbox_height_norm)## 找到该框中所有关键点，存在字典 bbox_keypoints_dict 中bbox_keypoints_dict = {}for each_ann in labelme['shapes']:  # 遍历所有标注if each_ann['shape_type'] == 'point':  # 筛选出关键点标注# 关键点XY坐标、类别x = int(each_ann['points'][0][0])y = int(each_ann['points'][0][1])label = each_ann['label']if (x > bbox_top_left_x) & (x < bbox_bottom_right_x) & (y < bbox_bottom_right_y) & \(y > bbox_top_left_y):  # 筛选出在该个体框中的关键点bbox_keypoints_dict[label] = [x, y]if each_ann['shape_type'] == 'line':  # 筛选出线段标注# 起点XY坐标、类别x0 = int(each_ann['points'][0][0])y0 = int(each_ann['points'][0][1])label = each_ann['label']bbox_keypoints_dict[label] = [x0, y0]# 终点XY坐标、类别x1 = int(each_ann['points'][1][0])y1 = int(each_ann['points'][1][1])label = int(each_ann['label']) + 1 #将字符串转为整形，并+1，代表最后一个点label = str(label) #将整型转为字符串bbox_keypoints_dict[label] = [x1, y1]# print(bbox_keypoints_dict)# if (x > bbox_top_left_x) & (x < bbox_bottom_right_x) & (y < bbox_bottom_right_y) & \#         (y > bbox_top_left_y):  # 筛选出在该个体框中的关键点#     bbox_keypoints_dict[label] = [x, y]## 把关键点按顺序排好for each_class in keypoint_class:  # 遍历每一类关键点if each_class in bbox_keypoints_dict:keypoint_x_norm = bbox_keypoints_dict[each_class][0] / img_widthkeypoint_y_norm = bbox_keypoints_dict[each_class][1] / img_heightyolo_str += '{:.5f} {:.5f} {} '.format(keypoint_x_norm, keypoint_y_norm, 2)  # 2可见不遮挡 1遮挡 0没有点else:  # 不存在的点，一律为0# yolo_str += '0 0 0 '.format(keypoint_x_norm, keypoint_y_norm, 0)yolo_str += '0 0 0 '# yolo_str += ' '# 写入 txt 文件中f.write(yolo_str + '\n')shutil.move(yolo_txt_path, save_folder) #从yolo_txt_path文件夹中移动到save_folder文件夹中# print('{} --> {} 转换完成'.format(labelme_path, yolo_txt_path))# json2txt
for labelme_path0 in os.listdir(path):labelme_path = path + '/' + labelme_path0print(labelme_path)process_single_json(labelme_path, save_folder)
print('YOLO格式的txt标注文件已保存至 ', save_folder)