labelme之批量生成掩码图（复制代码直接可用）

前言

当你看到这篇文章的时候，说明你在面临着标数据，这个巨烦的工作啦，我表示我懂，很难受。

然后labelme又不支持批量转换的，看网上的教程好多说要找到labelme的某个文件呀，然后在复制粘贴上去呀，在用命令行生成json文件夹的，老麻烦了，虽然我以前就是这么干的，但是干着干着就想偷懒了，所以自己翻看了labelme的代码，抽取出来用，自己用程序一步到位。

批量转换代码

代码是从labelme中抽取出来的，直接看main函数就好了，它生成的单通道彩色图，在PIL中对应的是P模式，有空我在放上opencv转成对应格式的

import json
import os
import sys

import PIL.Image
import cv2
import  numpy as np
import math
import PIL.ImageDraw
import os.path as osp
from tqdm import tqdm
def shapes_to_label(img_shape, shapes, label_name_to_value, type='class'):
    assert type in ['class', 'instance']

    cls = np.zeros(img_shape[:2], dtype=np.int32)
    if type == 'instance':
        ins = np.zeros(img_shape[:2], dtype=np.int32)
        instance_names = ['_background_']
    for shape in shapes:
        points = shape['points']
        label = shape['label']
        shape_type = shape.get('shape_type', None)
        if type == 'class':
            cls_name = label
        elif type == 'instance':
            cls_name = label.split('-')[0]
            if label not in instance_names:
                instance_names.append(label)
            ins_id = instance_names.index(label)
        cls_id = label_name_to_value[cls_name]
        mask = shape_to_mask(img_shape[:2], points, shape_type)
        cls[mask] = cls_id
        if type == 'instance':
            ins[mask] = ins_id

    if type == 'instance':
        return cls, ins
    return cls

def shape_to_mask(img_shape, points, shape_type=None,
                  line_width=10, point_size=5):
    mask = np.zeros(img_shape[:2], dtype=np.uint8)
    mask = PIL.Image.fromarray(mask)
    draw = PIL.ImageDraw.Draw(mask)
    xy = [tuple(point) for point in points]
    if shape_type == 'circle':
        assert len(xy) == 2, 'Shape of shape_type=circle must have 2 points'
        (cx, cy), (px, py) = xy
        d = math.sqrt((cx - px) ** 2 + (cy - py) ** 2)
        draw.ellipse([cx - d, cy - d, cx + d, cy + d], outline=1, fill=1)
    elif shape_type == 'rectangle':
        assert len(xy) == 2, 'Shape of shape_type=rectangle must have 2 points'
        draw.rectangle(xy, outline=1, fill=1)
    elif shape_type == 'line':
        assert len(xy) == 2, 'Shape of shape_type=line must have 2 points'
        draw.line(xy=xy, fill=1, width=line_width)
    elif shape_type == 'linestrip':
        draw.line(xy=xy, fill=1, width=line_width)
    elif shape_type == 'point':
        assert len(xy) == 1, 'Shape of shape_type=point must have 1 points'
        cx, cy = xy[0]
        r = point_size
        draw.ellipse([cx - r, cy - r, cx + r, cy + r], outline=1, fill=1)
    else:
        assert len(xy) > 2, 'Polygon must have points more than 2'
        draw.polygon(xy=xy, outline=1, fill=1)
    mask = np.array(mask, dtype=bool)
    return mask

def label_colormap(N=256):

    def bitget(byteval, idx):
        return ((byteval & (1 << idx)) != 0)

    cmap = np.zeros((N, 3))
    for i in range(0, N):
        id = i
        r, g, b = 0, 0, 0
        for j in range(0, 8):
            r = np.bitwise_or(r, (bitget(id, 0) << 7 - j))
            g = np.bitwise_or(g, (bitget(id, 1) << 7 - j))
            b = np.bitwise_or(b, (bitget(id, 2) << 7 - j))
            id = (id >> 3)
        cmap[i, 0] = r
        cmap[i, 1] = g
        cmap[i, 2] = b
    cmap = cmap.astype(np.float32) / 255
    return cmap

def lblsave(filename, lbl):
    if osp.splitext(filename)[1] != '.png':
        filename += '.png'
    # Assume label ranses [-1, 254] for int32,
    # and [0, 255] for uint8 as VOC.
    if lbl.min() >= -1 and lbl.max() < 255:
        lbl_pil = PIL.Image.fromarray(lbl.astype(np.uint8), mode='P')
        colormap = label_colormap(255)
        lbl_pil.putpalette((colormap * 255).astype(np.uint8).flatten())
        lbl_pil.save(filename)
    else:
        raise ValueError(
            '[%s] Cannot save the pixel-wise class label as PNG. '
            'Please consider using the .npy format.' % filename
        )

def main():
    img_path = 'C:/Users/Administrator/Desktop/temp/data/img'  #改这
    json_path = 'C:/Users/Administrator/Desktop/temp/data/output'#改这
    mask_save_path = 'C:/Users/Administrator/Desktop/temp/data/label'#改这
    os.makedirs(mask_save_path,exist_ok=True)

    # 标签类型  初始化为'_background_' 改这里就ok
    label_name_to_value ={'_background_': 0,'phone':1}

    json_name_list = os.listdir(json_path)
    json_name_list = [file for file in json_name_list if file.endswith('json')]

    error_list = []

    pbar =tqdm(total=len(json_name_list))
    for json_name in json_name_list:

        try:
            data = json.load(open(f'{json_path}/{json_name}'))

            name = json_name.split('.')[0]

            img = cv2.imread(f'{img_path}/{name}.jpg')
            if img is None:
                print("空图片")
                #sys.exit()   #报错时检查图片是否读取到

            for shape in data['shapes']:
                label_name = shape['label']
                if label_name in label_name_to_value:
                    label_value = label_name_to_value[label_name]
                else:
                    label_value = len(label_name_to_value)
                    label_name_to_value[label_name] = label_value

                label_values, label_names = [], []
                for ln, lv in sorted(label_name_to_value.items(), key=lambda x: x[1]):
                    label_values.append(lv)
                    label_names.append(ln)
                assert label_values == list(range(len(label_values)))

                lbl = shapes_to_label(img.shape, data['shapes'], label_name_to_value)

                output_name = json_name.split('.')[0]
                lblsave(osp.join(mask_save_path, f'{output_name}.png'), lbl)
        except Exception as e:
            error_list.append({f'{json_name}:{e}'})
        pbar.update(1)
    print(f"error_list len ={len(error_list)}")
    print(error_list)

if __name__ =="__main__":
    main()