json_to_dataset DONG2020

jupyter-notbook.exe

data_to_dataset.ipynb

# labelme 制作的数据 Annotations 和 JEPGImages，重点Annotations，因为里面自带图片数据了
json_file = r".\dataset"
image_file = r".\dataset\JPEGImages"
out = r".\dataset"
label_name_to_value = {'_background_': 0, "xxname":2, "xxname":1} # 指定类名对应的标签，bk总是为第0


import os
import numpy as np
from matplotlib import pyplot as plt
import PIL.Image
import json
import glob
import base64
from image import *
from shape import *
from draw import *
from draw import label_colormap
from skimage import img_as_ubyte
from PIL import Image
import tensorflow as tf

if out is None:
    out_dir = os.path.basename(json_file).replace('.', '_')
    out_dir = os.path.join(os.path.dirname(json_file), out_dir)
else:
    out_dir = out
if not os.path.exists(out_dir):
    os.mkdir(out_dir)


def lblsave(filename, lbl):
    if os.path.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
        )
        
paths = glob.glob(json_file+"/*.json")
if paths==[]:
    print("json is NULL")
n=0
for path in paths:
    if os.path.isfile(path):
        n+=1
        data = json.load(open(path))
        out_name = os.path.split(path)[-1][:-5]
    
    # imageData
    if data['imageData']:
        imageData = data['imageData']
    else:
        file_name = os.path.split(data['imagePath'])[-1]
        imagePath = os.path.join(image_file, file_name)     
        with open(imagePath, 'rb') as f:
            imageData = f.read()
            imageData = base64.b64encode(imageData).decode('utf-8')
            
    ''' img (array type)  '''
    img = img_b64_to_arr(imageData)
    
    # label_name
    for shape in data['shapes']:
        label_name = shape['label'] # 获取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)))            

    
    ''' label  : lbl已经为0，1，2，3，的标签了'''
    lbl = shapes_to_label(img.shape, data['shapes'], label_name_to_value)      
    
    ''' lbl_viz  '''
    captions = ['{}: {}'.format(lv, ln)
    for ln, lv in label_name_to_value.items()]
    lbl_viz = draw_label(lbl, img, captions)
    
    # 保存三个类型图 raw, segmentation, visualization
    print("第{:}张图".format(n))   
    JPEGImages_outdir = out_dir + "\JPEGImages"
    if not os.path.exists(JPEGImages_outdir):
        os.mkdir(JPEGImages_outdir)    
    img_out = out_dir + "\JPEGImages\\" + out_name  + '.jpg' # + str(n).zfill(4)
    PIL.Image.fromarray(img).save(img_out)              # 原图数据输出
    
    
    Segmentation_outdir = out_dir + "\SegmentationClassPNG"
    if not os.path.exists(Segmentation_outdir):
        os.mkdir(Segmentation_outdir)
    label_out = out_dir + "\SegmentationClassPNG\\" + out_name  + '.png' # + str(n).zfill(4)
    lblsave(label_out, lbl)                             # 标签数据输出

    
    Segmentation_label_outdir = out_dir + "\SegmentationClassRaw"
    if not os.path.exists(Segmentation_label_outdir):
        os.mkdir(Segmentation_label_outdir)
    label_int8_out = out_dir + "\SegmentationClassRaw\\" + out_name  + '.png' # + str(n).zfill(4)
    with tf.io.gfile.GFile(label_int8_out, mode='w') as f:
        Image.fromarray(lbl.astype(dtype=np.uint8)).save(f, 'PNG')
    
    
    Visualization_outdir = out_dir + "\SegmentationClassVisualization"
    if not os.path.exists(Visualization_outdir):
        os.mkdir(Visualization_outdir)    
    label_viz_out = out_dir + "\SegmentationClassVisualization\\" + out_name  + '.png' # + str(n).zfill(4)
    PIL.Image.fromarray(lbl_viz).save(label_viz_out)    # 掩码数据输出    
    
    
    with open(osp.join(out_dir, 'label_names.txt'), 'w') as f:
        for lbl_name in label_names:
            f.write(lbl_name + '\n')

image.py

import base64
import io

import numpy as np
import PIL.ExifTags
import PIL.Image
import PIL.ImageOps


def img_b64_to_arr(img_b64):
    f = io.BytesIO()
    f.write(base64.b64decode(img_b64))
    img_arr = np.array(PIL.Image.open(f))
    return img_arr


def img_arr_to_b64(img_arr):
    img_pil = PIL.Image.fromarray(img_arr)
    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


def img_data_to_png_data(img_data):
    with io.BytesIO() as f:
        f.write(img_data)
        img = PIL.Image.open(f)

        with io.BytesIO() as f:
            img.save(f, 'PNG')
            f.seek(0)
            return f.read()


def apply_exif_orientation(image):
    try:
        exif = image._getexif()
    except AttributeError:
        exif = None

    if exif is None:
        return image

    exif = {
        PIL.ExifTags.TAGS[k]: v
        for k, v in exif.items()
        if k in PIL.ExifTags.TAGS
    }

    orientation = exif.get('Orientation', None)

    if orientation == 1:
        # do nothing
        return image
    elif orientation == 2:
        # left-to-right mirror
        return PIL.ImageOps.mirror(image)
    elif orientation == 3:
        # rotate 180
        return image.transpose(PIL.Image.ROTATE_180)
    elif orientation == 4:
        # top-to-bottom mirror
        return PIL.ImageOps.flip(image)
    elif orientation == 5:
        # top-to-left mirror
        return PIL.ImageOps.mirror(image.transpose(PIL.Image.ROTATE_270))
    elif orientation == 6:
        # rotate 270
        return image.transpose(PIL.Image.ROTATE_270)
    elif orientation == 7:
        # top-to-right mirror
        return PIL.ImageOps.mirror(image.transpose(PIL.Image.ROTATE_90))
    elif orientation == 8:
        # rotate 90
        return image.transpose(PIL.Image.ROTATE_90)
    else:
        return image

shape.py

import math

import numpy as np
import PIL.Image
import PIL.ImageDraw

import warnings

def polygons_to_mask(img_shape, polygons, shape_type=None):
    warnings.warn(
        "The 'polygons_to_mask' function is deprecated, "
        "use 'shape_to_mask' instead."
    )
    return shape_to_mask(img_shape, points=polygons, shape_type=shape_type)


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 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 labelme_shapes_to_label(img_shape, shapes):
    warnings.warn('labelme_shapes_to_label is deprecated, so please use '
                'shapes_to_label.')

    label_name_to_value = {'_background_': 0}
    for shape in 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

    lbl = shapes_to_label(img_shape, shapes, label_name_to_value)
    return lbl, label_name_to_value


def masks_to_bboxes(masks):
    if masks.ndim != 3:
        raise ValueError(
            'masks.ndim must be 3, but it is {}'
            .format(masks.ndim)
        )
    if masks.dtype != bool:
        raise ValueError(
            'masks.dtype must be bool type, but it is {}'
            .format(masks.dtype)
        )
    bboxes = []
    for mask in masks:
        where = np.argwhere(mask)
        (y1, x1), (y2, x2) = where.min(0), where.max(0) + 1
        bboxes.append((y1, x1, y2, x2))
    bboxes = np.asarray(bboxes, dtype=np.float32)
    return bboxes

draw.py

import io
import os.path as osp

import numpy as np
import PIL.Image
import PIL.ImageDraw
import PIL.ImageFont


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 _validate_colormap(colormap, n_labels):
    if colormap is None:
        colormap = label_colormap(n_labels)
    else:
        assert colormap.shape == (colormap.shape[0], 3), \
            'colormap must be sequence of RGB values'
        assert 0 <= colormap.min() and colormap.max() <= 1, \
            'colormap must ranges 0 to 1'
    return colormap


# similar function as skimage.color.label2rgb
def label2rgb(
    lbl, img=None, n_labels=None, alpha=0.5, thresh_suppress=0, colormap=None,
):
    if n_labels is None:
        n_labels = len(np.unique(lbl))

    colormap = _validate_colormap(colormap, n_labels)
    colormap = (colormap * 255).astype(np.uint8)

    lbl_viz = colormap[lbl]
    lbl_viz[lbl == -1] = (0, 0, 0)  # unlabeled

    if img is not None:
        img_gray = PIL.Image.fromarray(img).convert('LA')
        img_gray = np.asarray(img_gray.convert('RGB'))
        # img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        # img_gray = cv2.cvtColor(img_gray, cv2.COLOR_GRAY2RGB)
        lbl_viz = alpha * lbl_viz + (1 - alpha) * img_gray
        lbl_viz = lbl_viz.astype(np.uint8)

    return lbl_viz


def draw_label(label, img=None, label_names=None, colormap=None, **kwargs):
    """Draw pixel-wise label with colorization and label names.

    label: ndarray, (H, W)
        Pixel-wise labels to colorize.
    img: ndarray, (H, W, 3), optional
        Image on which the colorized label will be drawn.
    label_names: iterable
        List of label names.
    """
    import matplotlib.pyplot as plt

    backend_org = plt.rcParams['backend']
    plt.switch_backend('agg')

    plt.subplots_adjust(left=0, right=1, top=1, bottom=0,
                        wspace=0, hspace=0)
    plt.margins(0, 0)
    plt.gca().xaxis.set_major_locator(plt.NullLocator())
    plt.gca().yaxis.set_major_locator(plt.NullLocator())

    if label_names is None:
        label_names = [str(l) for l in range(label.max() + 1)]

    colormap = _validate_colormap(colormap, len(label_names))

    label_viz = label2rgb(
        label, img, n_labels=len(label_names), colormap=colormap, **kwargs
    )
    plt.imshow(label_viz)
    plt.axis('off')

    plt_handlers = []
    plt_titles = []
    for label_value, label_name in enumerate(label_names):
        if label_value not in label:
            continue
        fc = colormap[label_value]
        p = plt.Rectangle((0, 0), 1, 1, fc=fc)
        plt_handlers.append(p)
        plt_titles.append('{value}: {name}'
                          .format(value=label_value, name=label_name))
    plt.legend(plt_handlers, plt_titles, loc='lower right', framealpha=.5)

    f = io.BytesIO()
    plt.savefig(f, bbox_inches='tight', pad_inches=0)
    plt.cla()
    plt.close()

    plt.switch_backend(backend_org)

    out_size = (label_viz.shape[1], label_viz.shape[0])
    out = PIL.Image.open(f).resize(out_size, PIL.Image.BILINEAR).convert('RGB')
    out = np.asarray(out)
    return out


def draw_instances(
    image=None,
    bboxes=None,
    labels=None,
    masks=None,
    captions=None,
):
    import matplotlib

    # TODO(wkentaro)
    assert image is not None
    assert bboxes is not None
    assert labels is not None
    assert masks is None
    assert captions is not None

    viz = PIL.Image.fromarray(image)
    draw = PIL.ImageDraw.ImageDraw(viz)

    font_path = osp.join(
        osp.dirname(matplotlib.__file__),
        'mpl-data/fonts/ttf/DejaVuSans.ttf'
    )
    font = PIL.ImageFont.truetype(font_path)

    colormap = label_colormap(255)
    for bbox, label, caption in zip(bboxes, labels, captions):
        color = colormap[label]
        color = tuple((color * 255).astype(np.uint8).tolist())

        xmin, ymin, xmax, ymax = bbox
        draw.rectangle((xmin, ymin, xmax, ymax), outline=color)
        draw.text((xmin, ymin), caption, font=font)

    return np.asarray(viz)