import os
from tqdm import tqdm
import numpy as np
import cv2
import json
from PIL import Image
import io
import base64
def change_label(json_path, old_class, new_class):
json_files = os.listdir(json_path)
for json_file in tqdm(json_files):
json_file_path = os.path.join(json_path, json_file)
with open(json_file_path, 'r') as f:
data = f.read()
if data == '':
print(json_file_path, 'is error')
continue
data = json.loads(data)
objects = data['shapes']
for i, obj in enumerate(objects):
label = obj['label']
if label == old_class:
label = new_class
else:
continue
objects[i]['label'] = label
data['shapes'] = objects
with open(json_file_path, 'w') as f:
json.dump(data, f, cls=NumpyArrayEncoder)
def set_image_data(json_path, img_path):
json_files = os.listdir(json_path)
img_suffix = __get_data_suffix(img_path)
for json_file in tqdm(json_files):
json_file_path = os.path.join(json_path, json_file)
with open(json_file_path, 'r') as f:
data = f.read()
if data == '':
print(json_file_path, 'is error')
continue
data = json.loads(data)
if data['imageData'] is not None:
continue
img_file_path = os.path.join(img_path, json_file.split('.')[0] + '.' + img_suffix)
img = Image.open(img_file_path)
imgdata = __img_tobyte(img)
data['imageData'] = imgdata
with open(json_file_path, 'w') as f:
json.dump(data, f, cls=NumpyArrayEncoder)
def __get_data_suffix(data_path):
data_files = os.listdir(data_path)
suffix = data_files[0].split('.')[-1]
return suffix
class NumpyArrayEncoder(json.JSONEncoder):
def default(self, obj):
if isinstance(obj, np.ndarray):
return obj.tolist()
return json.JSONEncoder.default(self, obj)
def __img_tobyte(img_pil):
# 类型转换 重要代码
# img_pil = Image.fromarray(roi)
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
def mask_to_json(img_path, json_path, mask_path, class_dict):
os.makedirs(json_path, exist_ok=True)
def func(img_file_path, mask, class_dict, ) -> dict:
gray = mask.astype(np.uint8)
img = Image.open(img_file_path)
imgData = __img_tobyte(img)
dic = {"version": "5.1.1", "flags": {}, "shapes": list(), "imagePath": img_file_path, "imageData": imgData,
"imageHeight": gray.shape[0], "imageWidth": gray.shape[1]}
for k, v in class_dict.items():
if v == 0:
continue
binary = gray.copy()
binary[binary != v] = 0
binary[binary == v] = 255
# 只检测外轮廓并存储所有的轮廓点
contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
for contour in contours:
temp = list()
if len(contour) < 10:
continue
for point in contour:
temp.append([float(point[0][0]), float(point[0][1])])
dic["shapes"].append({"label": k, "points": temp, "group_id": None,
"shape_type": "polygon", "flags": {}})
return dic
mask_suffix = __get_data_suffix(mask_path)
img_files = os.listdir(img_path)
for img_file in img_files:
save_file = img_file.split('.')[0] + '.json'
save_file_path = os.path.join(json_path, save_file)
mask_file_path = os.path.join(mask_path, img_file.split('.')[0] + '.' + mask_suffix)
mask = cv2.imread(mask_file_path)
mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
with open(save_file_path, mode='w', encoding='utf-8') as f:
img_file_path = os.path.join(img_path, img_file)
json.dump(func(img_file_path, mask, class_dict), f)
def json_to_mask(json_path, img_path, mask_path, class_dict):
def create_background_mask(img_file_path, mask_file_path):
img = cv2.imread(img_file_path)
mask = np.zeros_like(img, dtype=np.uint8)
cv2.imwrite(mask_file_path, mask)
def fun(json_file_path, img_file_path, mask_file_path, class_dict):
with open(json_file_path) as f:
data = f.read()
if data == '':
print(json_file_path, "is error")
return
data = json.loads(data)
img = cv2.imread(img_file_path)
mask = np.zeros_like(img, dtype=np.uint8)
objs = data["shapes"]
if len(objs) == 0:
create_background_mask(img_file_path, mask_file_path)
return
for obj in objs:
label = obj["label"]
points = np.array(obj["points"], dtype=np.int32)
value = class_dict[label]
cv2.fillPoly(mask, [points], (value, value, value))
for obj in objs:
label = obj["label"]
if "background" in label:
points = np.array(obj["points"], dtype=np.int32)
value = class_dict[label]
cv2.fillPoly(mask, [points], (value, value, value))
cv2.imwrite(mask_file_path, mask)
os.makedirs(mask_path, exist_ok=True)
img_files = os.listdir(img_path)
json_files = os.listdir(json_path)
for img_file in tqdm(img_files):
json_file = img_file.split('.')[0] + '.json'
mask_file = img_file.split('.')[0] + '.' + 'png'
if json_file not in json_files:
img_file_path = os.path.join(img_path, img_file)
mask_file_path = os.path.join(mask_path, mask_file)
create_background_mask(img_file_path, mask_file_path)
else:
json_file_path = os.path.join(json_path, json_file)
img_file_path = os.path.join(img_path, img_file)
mask_file_path = os.path.join(mask_path, mask_file)
fun(json_file_path, img_file_path, mask_file_path, class_dict)
语义分割数据处理相关代码
于 2023-05-26 11:00:05 首次发布