语义分割mask掩码转化为labelme格式(json文件)
- 前言
- 代码
- 完整代码
- 基于自己的任务运行代码
- 应用
前言
当我们数据集的语义标签为mask掩码格式时,而又想转换成labelme数据格式(json文件),达到如下图所示的结果,该如何实现呢?
代码
完整代码
骚话少说,直接上完整代码mask2json.py。
# 导入包
import os
import io
import json
import numpy as np
from pycococreatortools import pycococreatortools
from PIL import Image
import base64
def img_tobyte(img_pil):
'''
该函数用于将图像转化为base64字符类型
:param img_pil: Image类型
:return base64_string: 字符串
'''
ENCODING = 'utf-8'
img_byte = io.BytesIO()
img_pil.save(img_byte, format='PNG')
binary_str2 = img_byte.getvalue()
imageData = base64.b64encode(binary_str2)
base64_string = imageData.decode(ENCODING)
return base64_string
# 定义路径
ROOT_DIR = '' # 请输入你文件的根目录
Image_DIR = os.path.join(ROOT_DIR, "Image") # 目录底下包含图片
Label_DIR = os.path.join(ROOT_DIR, "GT") # 目录底下包含label文件
# 读取路径下的掩码
Label_files = os.listdir(Label_DIR)
# 指定png中index中对应的label
class_names = ['_background_', 'basketball', 'person'] # 分别表示label标注图中1对应basketball,2对应person。
for Label_filename in Label_files:
# 创建一个json文件
Json_output = {
"version": "3.16.7",
"flags": {},
"fillColor": [255, 0, 0, 128],
"lineColor": [0, 255, 0, 128],
"imagePath": {},
"shapes": [],
"imageData": {}}
print(Label_filename)
name = Label_filename.split('.', 3)[0]
name1 = name + '.jpg'
Json_output["imagePath"] = name1
# 打开原图并将其转化为labelme json格式
image = Image.open(Image_DIR + '/' + name1)
imageData = img_tobyte(image)
Json_output["imageData"] = imageData
# 获得注释的掩码
binary_mask = np.asarray(np.array(Image.open(Label_DIR + '/' + Label_filename))
).astype(np.uint8)
# 分别对掩码中的label结果绘制边界点
for i in np.unique(binary_mask):
if i != 0:
temp_mask = np.where(binary_mask == i, 1, 0)
segmentation = pycococreatortools.binary_mask_to_polygon(temp_mask, tolerance=2) # tolerancec参数控制无误差
for item in segmentation:
if (len(item) > 10):
list1 = []
for j in range(0, len(item), 2):
list1.append([item[j], item[j + 1]])
label = class_names[i] #
seg_info = {'points': list1, "fill_color": None, "line_color": None, "label": label,
"shape_type": "polygon", "flags": {}}
Json_output["shapes"].append(seg_info)
Json_output["imageHeight"] = binary_mask.shape[0]
Json_output["imageWidth"] = binary_mask.shape[1]
# 保存在根目录下的json文件中
full_path = '{}/json/' + name + '.json'
with open(full_path.format(ROOT_DIR), 'w') as output_json_file:
json.dump(Json_output, output_json_file)
基于自己的任务运行代码
大家根据自己的需求运行上述代码时还需要下载对应的工具包并修改一下以下参数。
1.需要下载pycococreatortools包
下载地址:https://github.com/waspinator/pycococreator
把红框文件夹放在与mask2json.py放在同一个目录里,如下图所示:
2.改文件路径
ROOT_DIR = '' # 请输入你文件的根目录
根目录下的文件夹按照下图方式的设置。
当然也可以按照自己喜欢的方式设置啦,这就需要友友们自行魔改代码。
3.改类别名。
class_names = ['_background_', 'basketball', 'person']
在本代码中设置为上面代码形式,那是因为语义分割label标签图中basketball区域的标签为1,person对应得区域为2。但为什么在本例中1为红色,2为绿色呢?这在我的另一篇blog中将会介绍。因此大家也需要根据自己的项目进行更改。
4. 改tolerance参数(可修改)
segmentation = pycococreatortools.binary_mask_to_polygon(temp_mask, tolerance=2)
tolerance参数的作用通俗一点说就是轮廓点构成的区域与mask标签区域的误差,tolerance越小表示两者的误差越小,但是轮廓点越多。
应用
如果大家也跟我一样为标注语义分割数据集而苦恼,我们可以利用上述代码减轻一些工作量。
1.用少量的数据集训练语义分割模型,然后用语义分割模型去预测,然后用labelme软件打开json文件手动进行微调。
2.也可以基于最近很火的大模型SAM,让大模型帮助获得分割掩码,然后基于labelme软件进行微调编辑。