labelme语义分割标注工具直接生成图像的改进方法

labelme语义分割标注工具直接生成图像的改进方法


使用过语义分割模型的人应该对labelme工具不会陌生,labelme工具允许我们以正多边形拟合的方式对图像中的目标区域进行标注,标注的结果保存在json文件中,json文件和被标注的原始图像一一对应。
但是在后续的使用中,很多语义分割的训练数据要求lable结果是图像的形式,也就是需要将json的结果转换成图像进而进行训练。
在网上看了很多方法将json转成图片,但是往往需要现用labelme得到json再进行转换,也就是两步操作;之前我也是一样的操作;但是最近由于要把标注工作教给很多人,两步的转换增添了工作量,于是考虑直接改动labelme的源码,使得标注结果直接就有图片形式和json格式。

操作如下:

  1. 安装labelme,推荐使用anaconda安装
# python2
conda create --name=labelme python=2.7
source activate labelme
# conda install -c conda-forge pyside2
conda install pyqt
pip install labelme
# if you'd like to use the latest version. run below:
# pip install git+https://github.com/wkentaro/labelme.git

# python3
conda create --name=labelme python=3.6
source activate labelme
# conda install -c conda-forge pyside2
# conda install pyqt
# pip install pyqt5  # pyqt5 can be installed via pip on python3
pip install labelme
# or you can install everything by conda command
# conda install labelme -c conda-forge
  1. 安装完以后找到anaconda的安装路径,进入env中的labelme,一下是我的路径,参考。

D:\anaconda\envs\labelme\Lib\site-packages\labelme

找到其中的label_file.py文件,修改如下。

import base64
import io
import json
import os.path as osp

import PIL.Image

from labelme._version import __version__
from labelme.logger import logger
from labelme import PY2
from labelme import QT4
from labelme import utils



MODE = 'BW' # 保存的mask图片类型,二值图为BW,彩图为RGB





import numpy as np
from labelme import utils
from labelme.utils.draw import label_colormap
import PIL.Image
import cv2
import os



class LabelFileError(Exception):
    pass


class LabelFile(object):

    suffix = '.json'

    def __init__(self, filename=None):
        self.shapes = ()
        self.imagePath = None
        self.imageData = None
        if filename is not None:
            self.load(filename)
        self.filename = filename

    @staticmethod
    def load_image_file(filename):
        try:
            image_pil = PIL.Image.open(filename)
            ##################################
            print("filename is :", filename)
        except IOError:
            logger.error('Failed opening image file: {}'.format(filename))
            return

        # apply orientation to image according to exif
        image_pil = utils.apply_exif_orientation(image_pil)

        with io.BytesIO() as f:
            ext = osp.splitext(filename)[1].lower()
            if PY2 and QT4:
                format = 'PNG'
            elif ext in ['.jpg', '.jpeg']:
                format = 'JPEG'
            else:
                format = 'PNG'
            image_pil.save(f, format=format)
            f.seek(0)
            return f.read()

    def load(self, filename):
        keys = [
            'imageData',
            'imagePath',
            'lineColor',
            'fillColor',
            'shapes',  # polygonal annotations
            'flags',   # image level flags
            'imageHeight',
            'imageWidth',
        ]
        try:
            with open(filename, 'rb' if PY2 else 'r') as f:
                data = json.load(f)
            if data['imageData'] is not None:
                imageData = base64.b64decode(data['imageData'])
                if PY2 and QT4:
                    imageData = utils.img_data_to_png_data(imageData)
            else:
                # relative path from label file to relative path from cwd
                imagePath = osp.join(osp.dirname(filename), data['imagePath'])
                imageData = self.load_image_file(imagePath)
            flags = data.get('flags') or {}
            imagePath = data['imagePath']
            self._check_image_height_and_width(
                base64.b64encode(imageData).decode('utf-8'),
                data.get('imageHeight'),
                data.get('imageWidth'),
            )
            lineColor = data['lineColor']
            fillColor = data['fillColor']
            shapes = (
                (
                    s['label'],
                    s['points'],
                    s['line_color'],
                    s['fill_color'],
                    s.get('shape_type', 'polygon'),
                    s.get('flags', {}),
                )
                for s in data['shapes']
            )
        except Exception as e:
            raise LabelFileError(e)

        otherData = {}
        for key, value in data.items():
            if key not in keys:
                otherData[key] = value


        # Only replace data after everything is loaded.
        self.flags = flags
        self.shapes = shapes
        self.imagePath = imagePath
        self.imageData = imageData
        self.lineColor = lineColor
        self.fillColor = fillColor
        self.filename = filename
        self.otherData = otherData

    @staticmethod
    def _check_image_height_and_width(imageData, imageHeight, imageWidth):
        img_arr = utils.img_b64_to_arr(imageData)
        if imageHeight is not None and img_arr.shape[0] != imageHeight:
            logger.error(
                'imageHeight does not match with imageData or imagePath, '
                'so getting imageHeight from actual image.'
            )
            imageHeight = img_arr.shape[0]
        if imageWidth is not None and img_arr.shape[1] != imageWidth:
            logger.error(
                'imageWidth does not match with imageData or imagePath, '
                'so getting imageWidth from actual image.'
            )
            imageWidth = img_arr.shape[1]
        return imageHeight, imageWidth

 '''
 保存json和使用json画图并保存,画图原理基本就是创建一个和原图尺寸相同的画布,
 然后根据json中描述的正多边形定点的键值对将正多边形内部的区域像素置换成
 指定的颜色,保存成新的图片完成标注
 '''
    def save(  
        self,
        filename,
        shapes,
        imagePath,
        imageHeight,
        imageWidth,
        imageData=None,
        lineColor=None,
        fillColor=None,
        otherData=None,
        flags=None,
    ):
        if imageData is not None:
            imageData = base64.b64encode(imageData).decode('utf-8')
            imageHeight, imageWidth = self._check_image_height_and_width(
                imageData, imageHeight, imageWidth
            )
        if otherData is None:
            otherData = {}
        if flags is None:
            flags = {}
        data = dict(
            version=__version__,
            flags=flags,
            shapes=shapes,
            lineColor=lineColor,
            fillColor=fillColor,
            imagePath=imagePath,
            imageData=imageData,
            imageHeight=imageHeight,
            imageWidth=imageWidth,
        )
        
        
        #################################################
        #imageData = base64.b64encode(imageData).decode('utf-8')
        img = utils.img_b64_to_arr(imageData)
        #print('img array is:',img)
        # 将标记的json数据转换为图片
        label_name_to_value = {'_background_': 0}
        for shape in sorted(data['shapes'], key=lambda x: x['label']):
            #print("json shape is:", shape)
            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 = utils.shapes_to_label(img.shape, data['shapes'], label_name_to_value)
        #cv2.imwrite("tower-label.jpg", lbl * 255)
        # 保存图片
        out_dir = os.path.join(os.path.abspath(os.path.join(filename,os.path.pardir)), 'mask')
        #out_dir = r"C:\Users\admin\Desktop\cv22\标注工具\mask-test"
        if not osp.exists(out_dir):
            os.mkdir(out_dir)
        #save_name = os.path.join(out_dir, imagePath.split('.')[0]+'.jpg')
        #save_name = "tower.jpg"
        save_name = os.path.join(out_dir,imagePath)
        print("out_dir is :",out_dir)
        print("imagepath is :",imagePath)
        print("save_name is :",save_name)
        if MODE == 'BW':
            lbl_tmp = cv2.resize(lbl, (500,400), interpolation=cv2.INTER_NEAREST)  #resize操作
            lbl_tmp = 255 - lbl_tmp * 255      #颠倒黑白,将数据分布在0-255范围
            lbl_tmp[lbl_tmp>100] = 255
            lbl_tmp[lbl_tmp<=100] = 0          #因为我的应用是二分类,所以全图强制二值化
            #print("img shape before resize is:",lbl_tmp.shape)
            #for i in range(224):
                #print(i in lbl_tmp)
            #print(lbl_tmp.shape)
            #lbl_tmp = lbl_tmp[255-lbl_tmp]
            #cv2.imwrite("tower1.jpg",lbl_tmp )
            #cv2.imwrite(save_name,lbl_tmp )
            lbl_pil = PIL.Image.fromarray(lbl_tmp.astype(np.uint8))
            lbl_pil.convert('1').save(save_name)
        elif MODE == 'RGB':
            lbl_pil = PIL.Image.fromarray(lbl.astype(np.uint8))
            colormap = label_colormap(255)
            lbl_pil.putpalette((colormap * 255).astype(np.uint8).flatten())
            lbl_pil.convert('RGB').save(save_name)
        else:
            print('save mode error!')
        #####################################################
        
        
        
        
        
        
        for key, value in otherData.items():
            data[key] = value
        try:
            with open(filename, 'wb' if PY2 else 'w') as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
            self.filename = filename
        except Exception as e:
            raise LabelFileError(e)

    @staticmethod
    def is_label_file(filename):
        return osp.splitext(filename)[1].lower() == LabelFile.suffix
  1. 修改完label_file.py, 重新运行labelme标注完成都既可以自动在保存的路径下创建mask文件夹,其中保存的是500x400 size的标注图像。因为我的项目的要求是resize到统一尺寸,也可以原图保存,只需要注释掉以上代码中的resize操作即可。
  2. 为了方便使用,也可以打包成exe可执行文件。
    Setup conda
    conda create --name labelme python==3.6.0
    conda activate labelme
    Build the standalone executable
    pip install .
    pip install pyinstaller
    pyinstaller labelme.spec
    dist/labelme --version

你可能感兴趣的:(labelme语义分割标注工具直接生成图像的改进方法)