最近自己在做语义分割数据集,使用labelme制作,但是我的数据集中每张图片不包含所有的类别,所以在标注的时候不同的类别可能会标注成相同的颜色,这样在训练的时候会出现错误,我对代码稍微改了一下,在这里记录一下,希望对需要的人也有帮助,代码如下:
import os
import random
import shutil
import re
import argparse
import json
import os
import os.path as osp
import warnings
import PIL.Image
import yaml
from labelme import utils
import base64
def main():
warnings.warn("This script is aimed to demonstrate how to convert the\n"
"JSON file to a single image dataset, and not to handle\n"
"multiple JSON files to generate a real-use dataset.")
parser = argparse.ArgumentParser()
parser.add_argument('json_file')
parser.add_argument('-o', '--out', default=None)
args = parser.parse_args()
json_file = args.json_file
#下面类别就是自己要分的类别数目,我的共有9类,而且下面这个字典的关键字(单引号内的)是自己在用labelme时标注的类别所以我自己又加了0-9的关键字,你的改成你自己的关键字就可以了,然后冒号后面的我猜代表的是颜色的标号,这样就会出来不同的颜色。
label_name_to_value = {'_background_': 0,
'0':1,
'1':2,
'2':3,
'3':4,
'4':5,
'5':6,
'6':7,
'7':8,
'8':9}
if args.out is None:
out_dir = osp.basename(json_file).replace('.', '_')
out_dir = osp.join(osp.dirname(json_file), out_dir)
else:
out_dir = args.out
if not osp.exists(out_dir):
os.mkdir(out_dir)
count = os.listdir(json_file)
for i in range(0, len(count)):
path = os.path.join(json_file, count[i])
if os.path.isfile(path):
data = json.load(open(path))
if data['imageData']:
imageData = data['imageData']
else:
imagePath = os.path.join(os.path.dirname(path), data['imagePath'])
with open(imagePath, 'rb') as f:
imageData = f.read()
imageData = base64.b64encode(imageData).decode('utf-8')
img = utils.img_b64_to_arr(imageData)
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 must be dense
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 = utils.shapes_to_label(img.shape, data['shapes'], label_name_to_value)
captions = ['{}: {}'.format(lv, ln)
for ln, lv in label_name_to_value.items()]
lbl_viz = utils.draw_label(lbl, img, captions)
out_dir = osp.basename(count[i]).replace('.', '_')
out_dir = osp.join(osp.dirname(count[i]), out_dir)
if not osp.exists(out_dir):
os.mkdir(out_dir)
name=out_dir.split('_')
rename=name[0]+'_'+name[1]
PIL.Image.fromarray(img).save(osp.join(out_dir, rename+'_'+'img.png'))
#PIL.Image.fromarray(lbl).save(osp.join(out_dir, rename+'label.png'))
utils.lblsave(osp.join(out_dir,rename+'_'+'label.png'), lbl)
PIL.Image.fromarray(lbl_viz).save(osp.join(out_dir,rename+'_'+'label_viz.png'))
with open(osp.join(out_dir, 'label_names.txt'), 'w') as f:
for lbl_name in label_names:
f.write(lbl_name + '\n')
warnings.warn('info.yaml is being replaced by label_names.txt')
info = dict(label_names=label_names)
with open(osp.join(out_dir, 'info.yaml'), 'w') as f:
yaml.safe_dump(info, f, default_flow_style=False)
print('Saved to: %s' % out_dir)
if __name__ == '__main__':
main()
这个是我的标签为0的,颜色设置为1 的,显示为红色。
这个是标签为1,颜色设置为2的,显示为绿色。
这个是标签为2,颜色显示为3的。
运行命令中需要指定自己要处理的json放置的文件夹,我输入法人命令如下:
python3 json_to_picture.py json
我使用的是python3 ,json_to_picture.py 是上面代码的文件名,json为需要处理的json文件放置的文件夹。
有问题可以和我留言交流。