批量分割mask转json

代码

import os
import cv2
import sys
import PIL
import copy
import json
import yaml
import base64
import numpy as np
import skimage.io as io
from glob import glob
try:
    from labelme import __version__ as labelme_version
except:
    labelme_version = '4.2.9'
sys.path.append('..')
currentCV_version = cv2.__version__


def rm(filepath):
    p = open(filepath, 'r+')
    lines = p.readlines()
    d = ""
    for line in lines:
        c = line.replace('"group_id": "null",', '"group_id": null,')
        d += c
    p.seek(0)
    p.truncate()
    p.write(d)
    p.close()


def imgEncode(img_or_path):
    if isinstance(img_or_path, np.ndarray):
        """
        copy from labelme image.py    
        """
        img_pil = PIL.Image.fromarray(img_or_path)
        f = io.BytesIO()
        img_pil.save(f, format='PNG')
        img_bin = f.getvalue()
        if hasattr(base64, 'encodebytes'):
            img_b64 = base64.encodebytes(img_bin)
        else:
            img_b64 = base64.encodestring(img_bin)
        return img_b64
    else:
        if isinstance(img_or_path, str):
            i = open(img_or_path, 'rb')
        elif isinstance(img_or_path, io.BufferedReader):
            i = img_or_path
        else:
            raise TypeError('Input type error!')
        base64_data = base64.b64encode(i.read())
        return base64_data.decode()


def rs(st: str):
    s = st.replace('\n', '').strip()
    return s


def readYmal(filepath, labeledImg=None):
    if os.path.exists(filepath):
        if filepath.endswith('.yaml'):
            f = open(filepath)
            y = yaml.load(f, Loader=yaml.FullLoader)
            f.close()
            # print(y)
            tmp = y['label_names']
            # print(tmp["tag1"])
            objs = zip(tmp.keys(), tmp.values())
            return sorted(objs)
        elif filepath.endswith('.txt'):
            f = open(filepath, 'r', encoding='utf-8')
            classList = f.readlines()
            f.close()
            l3 = [rs(i) for i in classList]
            l = list(range(1, len(classList)+1))
            objs = zip(l3, l)
            return sorted(objs)
    elif labeledImg is not None and filepath == "":
        """
        should make sure your label is correct!!!
        """
        labeledImg = np.array(labeledImg, dtype=np.uint8)

        labeledImg[labeledImg > 0] = 255
        labeledImg[labeledImg != 255] = 0
        # print(labeledImg)
        _, labels, stats, centroids = cv2.connectedComponentsWithStats(
            labeledImg)

        labels = np.max(labels) + 1
        labels = [x for x in range(1, labels)]

        classes = []
        for i in range(0, len(labels)):
            classes.append("class{}".format(i))

        return zip(classes, labels)
    else:
        raise FileExistsError('file not found')


def get_approx(img, contour, length_p=0.005):
    """获取逼近多边形
    :param img: 处理图片
    :param contour: 连通域
    :param length_p: 逼近长度百分比
    """
    img_adp = img.copy()
    # 逼近长度计算
    epsilon = length_p * cv2.arcLength(contour, True)
    # 获取逼近多边形
    approx = cv2.approxPolyDP(contour, epsilon, True)
    return approx


def getBinary(img_or_path, minConnectedArea=1):
    if isinstance(img_or_path, str):
        i = cv2.imread(img_or_path)
    elif isinstance(img_or_path, np.ndarray):
        i = img_or_path
    else:
        raise TypeError('Input type error')

    if len(i.shape) == 3:
        img_gray = cv2.cvtColor(i, cv2.COLOR_BGR2GRAY)

    else:
        img_gray = i

    ret, img_bin = cv2.threshold(img_gray, 127, 255, cv2.THRESH_BINARY)

    _, labels, stats, centroids = cv2.connectedComponentsWithStats(img_bin, connectivity=4)
    # labels:图像上每一像素的标记,用数字1、2、3…表示(不同的数字表示不同的连通域)
    # stats:每一个标记的统计信息,是一个5列的矩阵,每一行对应每个连通区域的外接矩形的x、y、width、height和面积,示例如下: 0 0 720 720 291805
    # centroids:连通域的中心点
    # print(stats.shape)  (19,5)
    # 删除区域小的图片
    for index in range(1, stats.shape[0]):
        if stats[index][4] < minConnectedArea or stats[index][4] < 0.0001 * (
                stats[index][2] * stats[index][3]):
            labels[labels == index] = 0

    labels[labels != 0] = 1

    img_bin = np.array(img_bin * labels).astype(np.uint8)
    return i, img_bin


def getMultiRegion(img, img_bin):
    """
    for multiple objs in same class
    """
    if float(currentCV_version[0:3]) < 3.5:
        img_bin, contours, hierarchy = cv2.findContours(
            img_bin, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
    else:
        contours, hierarchy = cv2.findContours(img_bin, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

    regions = []
    if len(contours) >= 1:
        for i in range(0, len(contours)):
            if i:
                # print(len(contours[i]))
                region = get_approx(img, contours[i], 0.0001)
                # print(region)
                if region.shape[0] > 3:
                    regions.append(region)

        return regions
    else:
        return []


def process(oriImg):
    img, img_bin = getBinary(oriImg)
    return getMultiRegion(img, img_bin)


def getMultiShapes(oriImgPath, labelPath, savePath='', labelYamlPath='', flag=False):
    """
    oriImgPath : for change img to base64  \n
    labelPath : after fcn/unet or other machine learning objects outlining , the generated label img
                or labelme labeled imgs(after json files converted to mask files)  \n
    savePath : json file save path  \n
    labelYamlPath : after json files converted to mask files. if doesn't have this file,should have a labeled img.
                    but the classes should change by yourself(labelme 4.2.9 has a bug,when change the label there will be an error.
                    )   \n
    """
    if isinstance(labelPath, str):
        if os.path.exists(labelPath):
            label_img = io.imread(labelPath)
        else:
            raise FileNotFoundError('mask/labeled image not found')
    else:
        label_img = labelPath

    # print(np.max(label_img))

    if np.max(label_img) > 127:
        # print('too many classes! \n maybe binary?')
        label_img[label_img > 127] = 255
        label_img[label_img != 255] = 0
        label_img = label_img / 255

    labelShape = label_img.shape

    labels = readYmal(labelYamlPath, label_img)
    # print(list(labels))
    shapes = []
    obj = dict()
    obj['version'] = labelme_version
    obj['flags'] = {}
    for la in list(labels):

        if la[1] > 0:
            # print(la[0])
            img = copy.deepcopy(label_img)  # img = label_img.copy()
            img = img.astype(np.uint8)
            img[img == la[1]] = 255
            img[img != 255] = 0

            region = process(img.astype(np.uint8))

            if isinstance(region, np.ndarray):

                points = []
                for i in range(0, region.shape[0]):
                    print(len(region[i][0]))
                    points.append(region[i][0].tolist())
                shape = dict()
                shape['label'] = la[0]
                shape['points'] = points
                shape['group_id'] = 'null'
                shape['shape_type'] = 'polygon'
                shape['flags'] = {}
                shapes.append(shape)

            elif isinstance(region, list):
                # print(len(region))
                for subregion in region:
                    points = []
                    for i in range(0, subregion.shape[0]):
                        points.append(subregion[i][0].tolist())
                    shape = dict()
                    shape['label'] = la[0]
                    shape['points'] = points
                    shape['group_id'] = 'null'
                    shape['shape_type'] = 'polygon'
                    shape['flags'] = {}
                    shapes.append(shape)

    # print(len(shapes))
    obj['shapes'] = shapes
    # print(shapes)
    (_, imgname) = os.path.split(oriImgPath)
    obj['imagePath'] = imgname
    # print(obj['imagePath'])
    obj['imageData'] = str(imgEncode(oriImgPath))

    obj['imageHeight'] = labelShape[0]
    obj['imageWidth'] = labelShape[1]

    j = json.dumps(obj, sort_keys=True, indent=4)
    # print(j)

    if not flag:
        saveJsonPath = savePath + os.sep + obj['imagePath'][:-4] + '.json'
        # print(saveJsonPath)
        with open(saveJsonPath, 'w') as f:
            f.write(j)

        rm(saveJsonPath)
    else:
        return j


if __name__ == "__main__":
    path = ''
    init_path = '%s/image' % path
    mask_path = '%s/mask' % path

    yaml_file = '%s/label_names.yaml' % path
    save_json = '%s/json' % path

    mask_images_list = glob(os.path.join(mask_path, "*.png"))
    init_images_list = glob(os.path.join(init_path, "*.png"))

    if not os.path.exists(save_json):
        os.mkdir(save_json)

    for mask_image, init_image in zip(mask_images_list, init_images_list):
        print(mask_image)
        getMultiShapes(init_image, mask_image, save_json, yaml_file)

label_name.yaml格式

label_names:
    Tag1: 1
    Tag2: 2
    Tag3: 3
    类别: 掩码像素值
    ....

参考

https://github.com/guchengxi1994/mask2json

  • 3
    点赞
  • 22
    收藏
    觉得还不错? 一键收藏
  • 17
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 17
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值