语义标签转实例标签

最近尝试使用YOLO v5做实例分割，手里只有语义分割的标签，而且制造的标签比较不同寻常（用PS做的标签），所以需要转一下标签的格式。借鉴了这位大佬的代码

如何将语义分割数据集转换为实例分割数据集_语义分割标签如何转为实力分割标签_点PY的博客-CSDN博客y​​​​​​

如果哪里可以改进，欢迎提建议，这里只要把路径稍微修改一下就好了

代码

import os
from skimage.measure import label, regionprops
import numpy as np
from skimage.morphology import erosion, square
from PIL import Image
import cv2

img_path = 
mask_path = 
label_path = 
img_list = {}
mask_list = {}
for img in os.listdir(img_path):
    img_list[img] = os.listdir(os.path.join(img_path, img))
for mask in os.listdir(mask_path):
    mask_list[mask] = os.listdir(os.path.join(mask_path, mask))

for i, j in zip(img_list.keys(),mask_list.keys()):
    img_list[i].sort()
    mask_list[j].sort()
    for img_name, mask_name in zip(img_list[i],mask_list[j]):
        assert img_name.split('.')[0] == mask_name.split('.')[0] #确保图像和mask相同
        imgpath = os.path.join(img_path, i, 'raw', img_name)
        maskpath = os.path.join(mask_path, j, mask_name)

        img_0 = Image.open(imgpath)
        img_0 = np.array(img_0)

        mask = Image.open(maskpath).convert('L')
        mask = np.array(mask)
        mask = np.where(mask == 0, 1, 0) #将mask=0的元素变为0，非0变为0
        w, h = mask.shape

        mask = erosion(mask, square(2)).astype(np.uint8) #腐蚀操作，防止粘连无法将实例分开
        label_0 = label(mask) #sklearn的一个函数 → 分离图像中的连通区域 → 获取含有实例标签的数组
        props = regionprops(label_0) #获取图像中各个区域的属性 → 面积、中心坐标、周长
        img_1 = img_0.copy()
        mask2 = mask.copy()

        for id, prop in enumerate(props):
            mask_ = np.where(label_0 == id + 1, 1, 0).astype(np.uint8)
            contours, hierarchy = cv2.findContours(mask_, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
            cv2.drawContours(img_1, contours, -1, (0, 0, 255), 3)

            contours1 = contours[0]
            segmentation = np.array(contours1.reshape(1,-1), dtype=np.float)
            segmentation = segmentation.reshape((-1, 2))
            segmentation[:, 0] = segmentation[:, 0] / w #标准化 除以图像的宽
            segmentation[:, 1] = segmentation[:, 1] / h	#除以图像的长
            segmentation = segmentation.reshape(-1) #重构图像
            segmentation = list(segmentation)
            labelPath = os.path.join(label_path, i)
            if not os.path.exists(labelPath):
                os.makedirs(labelPath)
            with open(os.path.join(labelPath, img_name.split('.')[0] + '.txt'), 'a+') as f:
                f.write(f'{0} {" ".join(list(map(lambda x:f"{x:.6f}", segmentation)))}') #写入txt文件
                f.write('\n')