File size: 2,982 Bytes
eec42bd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os, sys, argparse
import glob
import numpy as np
from PIL import Image
from tqdm import tqdm
from collections import OrderedDict
from labelme.utils import lblsave as label_save
sys.path.append(os.path.join(os.path.dirname(os.path.realpath(__file__)), '..', '..', '..'))
from common.utils import get_classes
def label_convert(input_label_path, output_label_path, class_names=None):
if not os.path.isdir(input_label_path):
raise ValueError('Input path does not exist!\n')
os.makedirs(output_label_path, exist_ok=True)
if class_names:
# count class item number
class_count = OrderedDict([(item, 0) for item in class_names])
label_files = glob.glob(os.path.join(input_label_path, '*.png'))
pbar = tqdm(total=len(label_files), desc='Label converting')
for label_file in label_files:
label_array = np.array(Image.open(label_file))
if class_names:
# count object class for statistic
label_list = list(np.unique(label_array))
for label in label_list:
if label >= len(class_names):
continue
class_name = class_names[label]
class_count[class_name] = class_count[class_name] + 1
# save numpy label array as png label image,
# using labelme utils function
png_file_name = os.path.basename(label_file).split('.')[0]+'.png'
label_save(os.path.join(output_label_path, png_file_name), label_array)
pbar.update(1)
pbar.close()
print('total number of converted images: ', len(label_files))
if class_names:
# show item number statistic
print('Image number for each class:')
for (class_name, number) in class_count.items():
if class_name == 'background':
continue
print('%s: %d' % (class_name, number))
def main():
parser = argparse.ArgumentParser(argument_default=argparse.SUPPRESS, description='convert onboard gray PNG label to PascalVOC PNG label')
parser.add_argument('--input_label_path', type=str, required=True, help='input path of gray label files')
parser.add_argument('--output_label_path', type=str, required=True, help='output path of converted png label files')
parser.add_argument('--classes_path', type=str, required=False, help='path to class definitions, optional', default=None)
args = parser.parse_args()
if args.classes_path:
# add background class to match model & GT
class_names = get_classes(args.classes_path)
assert len(class_names) < 254, 'PNG image label only support less than 254 classes.'
class_names = ['background'] + class_names
else:
class_names = None
label_convert(args.input_label_path, args.output_label_path, class_names)
if __name__ == '__main__':
main()
|