YOLOv5实例分割数据集准备:
1. 将COCO格式JSON文件划分为多个LabelMe格式文件
2. 批量转换:LabelMe to YOLO
import glob
import numpy as np
import json
import os
import cv2
json_path = r"coco2labelme/train"
json_files = glob.glob(json_path + "/*.json")
for json_file in json_files:
print(json_file)
f = open(json_file)
json_info = json.load(f)
# print(json_info.keys())
img = cv2.imread(os.path.join(json_path, json_info["imagePath"]))
height, width, _ = img.shape
np_w_h = np.array([[width, height]], np.int32)
txt_file = json_file.replace(".json", ".txt")
f = open(txt_file, "a")
for point_json in json_info["shapes"]:
txt_content = ""
np_points = np.array(point_json["points"], np.int32)
label = point_json["label"]
if label == 'class 1':
label_index = 0
elif label == 'class 2':
label_index = 1
elif label == 'class 3':
label_index = 2
np_points = np.array(point_json["points"], np.int32)
norm_points = np_points / np_w_h
norm_points_list = norm_points.tolist()
txt_content += f"{label_index} " + " ".join([" ".join([str(cell[0]), str(cell[1])]) for cell in norm_points_list]) + "\n"
f.write(txt_content)
3. 可视化以检查数据集标注
import cv2
import numpy as np
pic_path = r"coco2labelme/val/1.jpg"
txt_path = r"coco2labelme/val/1.txt"
img = cv2.imread(pic_path)
height, width, _ = img.shape
file_handle = open(txt_path)
cnt_info = file_handle.readlines()
new_cnt_info = [line_str.replace("\n", "").split(" ") for line_str in cnt_info]
color_map = [(0, 255, 255), (255, 0, 255), (255, 255, 0)]
for new_info in new_cnt_info:
s = []
for i in range(1, len(new_info), 2):
b = [float(tmp) for tmp in new_info[i:i + 2]]
s.append([int(b[0] * width), int(b[1] * height)])
class_ = new_info[0]
index = int(class_)
cv2.polylines(img, [np.array(s, np.int32)], True, color_map[index], thickness = 3)
cv2.imshow('img2', img)
cv2.waitKey()
对子文件夹内所有数据的标注进行可视化并保存:
import cv2
import numpy as np
import glob
# path
pic_path = r"images/train/"
txt_path = r"labels/train/"
pic = glob.glob(pic_path + "*.jpg")
for pic_file in pic:
img = cv2.imread(pic_file)
substrings = pic_file.split('\\')
substrings = substrings[1].split('.')
num=substrings[0]
height, width, _ = img.shape
txt_file = txt_path + num + ".txt"
file_handle = open(txt_file)
cnt_info = file_handle.readlines()
new_cnt_info = [line_str.replace("\n", "").split(" ") for line_str in cnt_info]
color_map = [(0, 255, 255), (255, 0, 255), (255, 255, 0)]
for new_info in new_cnt_info:
s = []
for i in range(1, len(new_info), 2):
b = [float(tmp) for tmp in new_info[i:i + 2]]
s.append([int(b[0] * width), int(b[1] * height)])
class_ = new_info[0]
index = int(class_)
cv2.polylines(img, [np.array(s, np.int32)], True, color_map[index], thickness = 3)
save_path = 'visual/train/' + num + '.jpg'
cv2.imwrite(save_path, img)
参考:[1]