YOLOV6_Seg 图片标签数据预处理代码

1.将json数据转化为yolov6训练专用的txt数据，图像数据和json数据要同时放在img_json中。

import os, cv2, json
import numpy as np
from imutils import paths
import shutil

classes = ['apple','banana','peach']  # 修改成对应的类别

base_path = './img_json'  # 指定json和图片的位置

img_outpath=r'./images'
lbe_outpath=r'./labels'

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

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

imagepaths=list(paths.list_images(base_path))

for imagepath in imagepaths:
    basename=os.path.basename(imagepath).split('.')[0]
    image=cv2.imread(imagepath,-1)
    h, w, c = image.shape
    try:
        with open(os.path.join(base_path,basename+'.json'),'r') as f:
            masks = json.load(f)['shapes']
    except:
        continue

    shutil.copy(imagepath,os.path.join(img_outpath,os.path.basename(imagepath)))
    with open(os.path.join(lbe_outpath,basename+'.txt'), 'w+') as f:
        for idx, mask_data in enumerate(masks):
            mask_label = mask_data['label']
            if '_' in mask_label:
                mask_label = mask_label.split('_')[0]
            mask = np.array([np.array(i) for i in mask_data['points']], dtype=np.float64)
            mask[:, 0] /= w
            mask[:, 1] /= h
            mask = mask.reshape((-1))
            if idx != 0:
                f.write('\n')
            f.write(f'{classes.index(mask_label)} {" ".join(list(map(lambda x: f"{x:.6f}", mask)))}')

2. 验证标签数据

import cv2
import numpy as np
import glob
import os

# 只需要给定图片文件夹和txt标签文件夹即可
pic_path = r"./images/"
txt_path = r"./labels/"

output_path=r'./suoyin_output'

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


lbes=glob.glob(txt_path + "*.txt")
color_map = np.random.randint(0, 256, (256, 3), dtype=np.uint8)

for lbe in lbes:
    basename=os.path.basename(lbe).split('.')[0]
    file_handle = open(lbe)
    img = cv2.imread(lbe.replace('labels','images').replace('txt','jpg'),-1)
    height, width, _ = img.shape
    cnt_info = file_handle.readlines()
    print("***:", cnt_info)
    new_cnt_info = [line_str.replace("\n", "").split(" ") for line_str in cnt_info]
    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_)
        color=color_map[index]
        cv2.polylines(img, [np.array(s, np.int32)], True, color=(int(color[0]),int(color[1]),int(color[2])), thickness=3)

    save_path = os.path.join(output_path,basename+'.png')
    cv2.imwrite(save_path, img)