语义分割数据处理相关代码

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)