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 外匯跟單gendan5.com+ '\\' + 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()