标注神器之auto_labelme

1、深度学习的数据标注是一件惨绝人寰的事情。

2、可利用一些参数量大的模型生成掩码，再利用opencv寻找边界，将边界点写入json文件。用labelme读入这些json文件，进行矫正，能够大幅度减少人为标定的工作量。

3、代码：

# -*- coding: utf-8 -*-

import os
import torch
from cv2 import findContours,contourArea,approxPolyDP,RETR_EXTERNAL,CHAIN_APPROX_SIMPLE
import json
import base64
import numpy as np
from PIL import Image
from scipy.ndimage import binary_fill_holes,binary_opening


class MyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        else:
            return super(MyEncoder, self).default(obj)

def data_process(x):
    x = np.array(x, dtype='float32') / 255
    x1 = np.expand_dims(x, 0)
    x = np.transpose(x1, (0, 3, 1, 2))
    x = torch.from_numpy(x)

    mean = [0.517446, 0.360147, 0.310427]
    std = [0.061526, 0.049087, 0.041330]
    mean = torch.tensor(mean, dtype=torch.float32)
    std = torch.tensor(std, dtype=torch.float32)
    return x.sub_(mean[:, None, None]).div_(std[:, None, None])

def result_save_to_json(res, img_path, img_name):
    ##########　teeth segmentation label ##########
    teeth = res == 1
    teeth = binary_fill_holes(teeth)
    teeth = binary_opening(teeth, iterations=10)
    teeth = teeth.astype('uint8')*255
    _, contours, _= findContours(teeth.copy(), RETR_EXTERNAL, CHAIN_APPROX_SIMPLE)

    shapes_arrary = []
    for contour in contours: # 遍历多个轮廓
        if contourArea(contour) < 5000: # remove small region
            continue
        contour = approxPolyDP(contour, 2.5, True) # 直线拟合轮廓，减少轮廓点的数量

        shape = {}
        shape["line_color"] = [255, 0, 0, 128]
        shape["shape_type"] = "polygon"
        shape["flags"] = {}
        shape["fill_color"] = [0, 255, 0, 128]
        shape["label"] = "tooth"
        shape["points"] = []
        for point in contour: # 遍历一个轮廓中的点
            point = point[0]
            shape["points"].append([point[0], point[1]])
        shapes_arrary.append(shape)

    ########## gum segmentation label ##########
    gum = res == 2
    gum = binary_fill_holes(gum)
    gum = binary_opening(gum, iterations=10)
    gum = gum.astype('uint8')*255
    _, contours, _= findContours(gum.copy(), RETR_EXTERNAL, CHAIN_APPROX_SIMPLE)
    for contour in contours: # 遍历多个轮廓
        if contourArea(contour) < 5000: # remove small region
            continue
        contour = approxPolyDP(contour, 2.5, True) # 直线拟合轮廓，减少轮廓点的数量

        shape = {}
        shape["line_color"] = [0, 255, 0, 128]
        shape["shape_type"] = "polygon"
        shape["flags"] = {}
        shape["fill_color"] = [0, 0, 255, 128]
        shape["label"] = "gum"
        shape["points"] = []
        for point in contour: # 遍历一个轮廓中的点
            point = point[0]
            shape["points"].append([point[0], point[1]])
        shapes_arrary.append(shape)

    with open(img_path + img_name, "rb") as f:
        base64_data= base64.b64encode(f.read())
        base64_str = base64_data.decode()

    info_dict = {
        "imagePath" : "./"+img_name,
        "imageData" : base64_str,
        "shapes": shapes_arrary,
        "version": "3.16.1",
        "flags":{},
        "fillColor":[255, 0, 0, 128],
        "lineColor":[0, 255, 0, 128],
        "imageWidth": 640,
        "imageHeight": 480
    }

    json_name = img_path + img_name.split('.')[0]+'.json'
    json_str = json.dumps(info_dict, cls=MyEncoder)
    with open(json_name, 'w') as json_file:
        json_file.write(json_str)

def network_process(model_path, image_path):

    os.system("rm -rf " + image_path + "*.json")

    model = torch.load(model_path)
    model.eval()

    img_folder = os.listdir(image_path)

    for iter, image_name in enumerate(img_folder):
        if iter % 100 == 0:
            print("### proc %d / %d"%(iter, len(img_folder)))

        if image_name.endswith(('.jpg','.JPG','.png','.PNG','.bmp')):
            img_file = os.path.join(image_path,image_name)

            try:
                image = Image.open(img_file)
            except (OSError, NameError):
                print("OSError %s can't open"%(img_file))
                continue

            x = data_process(image)
            y = model(x.cuda().float())
            y = torch.nn.Softmax2d()(y)
            numpy_y = y.cpu().detach().numpy()[0, ...]
            y1 = np.argmax(numpy_y, axis=0)

            result_save_to_json(y1, image_path, image_name)


if __name__ =="__main__":
    model_path=r"/home/anchao/桌面/super_labelme_mini/model/0.8976_model.path"
    img_path=r"/home/anchao/桌面/super_labelme_mini/images/"
    network_process(model_path, img_path)