注解: 下列使用场景为2类箱子识别
import json
import matplotlib.pyplot as plt
import numpy as np
import cv2
import os
from labelme import utils
from PIL import Image
json_file_path = r'D:/00-机器学习/data/json' # 保存json文件的文件夹
png_save_path = r'D:/00-机器学习/data/data_png_one/'#生成的图片保存路径
json_file_names = os.listdir(json_file_path) # json列表
for json_file_name in json_file_names:
json_file = os.path.join(json_file_path, json_file_name)
# json_file = r'D:/pycharm_projects/g3_work/data/new_data_json/1007.json'
json_name = str(json_file)
j_name_list = json_name.split('\\') # ['D:/00-机器学习/data/json', '418.json']
j_number = int(j_name_list[-1][:-5]) # 418
data = json.load(open(json_file))
''' 其他
data['imageData'] # 原图数据 str
data['shapes'] # 每个对像mask及label list
len(data['shapes']) # 返回对象个数 int
data['shapes'][0]['label'] # 返回第一个对象的标签 str
data['shapes'][0]['points'] # 返回第一个对象的边界点 list
data['shapes'][0]['points'][0] # 返回第一个对象的边界点第一个点 list
data['imagePath'] # 原图路径 str
data['fillColor'] # 填充颜色(边界内部) list
data['lineColor'] # 边界线颜色 list
'''
img = utils.img_b64_to_arr(data['imageData']) # data['imageData']:原图数据 str
lbl, lbl_names = utils.labelme_shapes_to_label(img.shape, data['shapes']) # data['shapes']:每个对像mask及label list
captions = ['%d: %s' % (l, name) for l, name in
enumerate(lbl_names)] # 转换成字符列表的形式,cations:['0: _background_', '1: box0']
'''
utils.labelme_shapes_to_label:
参数:
data['shape']:每个目标对象标注后的信息
eg.[{'label': 'box0', 'points': [[144.44444444444443, 154.1111111111111], [127.77777777777777, 365.2222222222222],
[388.88888888888886, 367.074074074074], [361.1111111111111, 157.81481481481478]],'group_id': None,
'shape_type': 'polygon', 'flags': {}}]
返回值:
lbl:数组的形式,值为0、1 (目标对象为1,背景为0)
lbl_names :字典的形式,存放的背景,对象对应的编码 eg.{'_background_': 0, 'box0': 1}
'''
'''原图中目标对象用颜色标注'''
lbl_viz = utils.draw_label(lbl, img, captions)
lbl_viz_img = Image.fromarray(lbl_viz)
# lbl_viz_img.show()
# exit()
'''掩码颜色设置'''
mask = [] # 存放
class_id = [] # 存放目标对象编号(不包括背景即0)
for i in range(1, len(lbl_names)): # 若有多个class(物体) 跳过第一个class(默认为背景)
mask.append((lbl == i).astype(np.uint8)) # 解析出像素值为1的对应,对应第一个对象 mask 为0、1组成的(0为背景,1为对象)
class_id.append(i) # mask与clas 一一对应
# mask通道数对应标签种类数,每一个通道的有目标位置像素值为1
mask = np.transpose(np.asarray(mask, np.uint8), [1, 2, 0]) # 转成[h,w,instance count]
class_id = np.asarray(class_id, np.uint8) # [instance count,]
'''背景颜色+对象颜色设置'''
background = np.zeros((512, 512, 3)) # 创建背景区域
box_dict = {'box1': 0, 'box0': 1, '_background_': 2}
# 定义掩码颜色(不同箱子同一颜色or不同类型箱子不同颜色)
color_list = [(0, 255, 0), (255, 0, 0)] # 标准箱子绿色,其余箱子红色
# color_list = [(0, 255, 0), (0, 255, 0)] # 所有箱子全为绿色
name_list = [] # 存放目标对象编号(除了背景外,该例子只有箱子故为1)
for key in lbl_names:
if key == '_background_':
continue
name_list.append(box_dict[key])
for i, cls in enumerate(name_list):
nmsl = mask[:, :, i] == 1
background[nmsl, :] = color_list[cls]
# cv2.imshow('',background)
# cv2.waitKey(0)
# cv2.destroyAllWindows()
# print(background.shape)
res_img = Image.fromarray(np.uint8(background))
res_img = res_img.convert('RGBA')
# 保存png图片,需要提前创建好文件夹
res_img.save(png_save_path+'{}.png'.format(j_number))