Ai图像分割模型PaddleSeg——自定义数据集处理

#!/usr/bin/python

# -*- coding: UTF-8 -*-

import argparse

import json

import os

import os.path as osp

import base64

import warnings

import PIL.Image

import yaml

from labelme import utils

import numpy as np

from skimage import img_as_ubyte

def main():

    parser = argparse.ArgumentParser()

    parser.add_argument('json_file')

    parser.add_argument('-o', '--out', default=None)

    args = parser.parse_args()

    json_file = args.json_file

    list_path = os.listdir(json_file)

    for i in range(0, len(list_path)):

        if list_path[i].endswith('.json'):

            path = os.path.join(json_file, list_path[i])

            if os.path.isfile(path):

                data = json.load(open(path))

                img = utils.img_b64_to_arr(data['imageData'])

                lbl, lbl_names = utils.labelme_shapes_to_label(img.shape, data['shapes'])

                captions = ['%d: %s' % (l, name) for l, name in enumerate(lbl_names)]

                lbl_viz = utils.draw_label(lbl, img, captions)

                save_file_name = osp.basename(path).replace('.', '_')

                out_dir1 = osp.join(osp.dirname(path), 'labelme_results')

                if not osp.exists(out_dir1):

                    os.mkdir(out_dir1)

                out_dir1 = osp.join(out_dir1, save_file_name)

                if not osp.exists(out_dir1):

                    os.mkdir(out_dir1)

                PIL.Image.fromarray(img).save(out_dir1 + '\\' + save_file_name + '_img.png')

                PIL.Image.fromarray(lbl).save(out_dir1 + '\\' + save_file_name + '_label.png')

                PIL.Image.fromarray(lbl_viz).save(out_dir1 + '\\' + save_file_name +

                                                  '_label_viz.png')

                images_dir = osp.join(json_file, 'images_dir')

                if not osp.exists(images_dir):

                    os.mkdir(images_dir)

                labels_dir = osp.join(json_file, 'labels_dir')

                if not osp.exists(labels_dir):

                    os.mkdir(labels_dir)

                PIL.Image.fromarray(img).save(osp.join(images_dir, '{}_img.png'.format(save_file_name)))

                PIL.Image.fromarray(lbl).save(osp.join(labels_dir, '{}_label.png'.format(save_file_name)))

                with open(osp.join(out_dir1, 'label_names.txt'), 'w') as f:

                    for lbl_name in lbl_names:

                        f.write(lbl_name + '\n')

                info = dict(label_names=lbl_names)

                with open(osp.join(out_dir1, 'info.yaml'), 'w') as f:

                    yaml.safe_dump(info, f, default_flow_style=False)

                print('Saved to: %s' % out_dir1)

if __name__ == '__main__':

    # base64path = argv[1]

    main()