【实例分割】YOLO自定义数据集准备：COCO转YOLO（可视化标注）

YOLOv5实例分割数据集准备：

1. 将COCO格式JSON文件划分为多个LabelMe格式文件

【目标检测】JSON转换：COCO格式转LabelMe格式（大JSON文件划分）_ericdiii的博客-优快云博客将大JSON文件划分为每张图片对应的标注文件https://blog.youkuaiyun.com/ericdiii/article/details/128281336?spm=1001.2014.3001.5502

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]