File size: 3,303 Bytes
20e841b |
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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 |
import sys
import os
import glob
import argparse
# this script will load class_list.txt and find class names with spaces
# then replace spaces with delimiters inside ground-truth/ and predicted/
parser = argparse.ArgumentParser()
parser.add_argument('-d', '--delimiter', type=str, help="delimiter to replace space (default: '-')", default='-')
parser.add_argument('-y', '--yes', action='store_true', help="force yes confirmation on yes/no query (default: False)", default=False)
args = parser.parse_args()
def query_yes_no(question, default="yes", bypass=False):
"""Ask a yes/no question via raw_input() and return their answer.
"question" is a string that is presented to the user.
"default" is the presumed answer if the user just hits <Enter>.
It must be "yes" (the default), "no" or None (meaning
an answer is required of the user).
The "answer" return value is True for "yes" or False for "no".
"""
valid = {"yes": True, "y": True, "ye": True,
"no": False, "n": False}
if default is None:
prompt = " [y/n] "
elif default == "yes":
prompt = " [Y/n] "
elif default == "no":
prompt = " [y/N] "
else:
raise ValueError("invalid default answer: '%s'" % default)
while True:
sys.stdout.write(question + prompt)
if bypass:
break
if sys.version_info[0] == 3:
choice = input().lower() # if version 3 of Python
else:
choice = raw_input().lower()
if default is not None and choice == '':
return valid[default]
elif choice in valid:
return valid[choice]
else:
sys.stdout.write("Please respond with 'yes' or 'no' "
"(or 'y' or 'n').\n")
def rename_class(current_class_name, new_class_name):
# get list of txt files
file_list = glob.glob('*.txt')
file_list.sort()
# iterate through the txt files
for txt_file in file_list:
class_found = False
# open txt file lines to a list
with open(txt_file) as f:
content = f.readlines()
# remove whitespace characters like `\n` at the end of each line
content = [x.strip() for x in content]
new_content = []
# go through each line of eache file
for line in content:
#class_name = line.split()[0]
if current_class_name in line:
class_found = True
line = line.replace(current_class_name, new_class_name)
new_content.append(line)
if class_found:
# rewrite file
with open(txt_file, 'w') as new_f:
for line in new_content:
new_f.write("%s\n" % line)
with open('../../data/classes/coco.names') as f:
for line in f:
current_class_name = line.rstrip("\n")
new_class_name = line.replace(' ', args.delimiter).rstrip("\n")
if current_class_name == new_class_name:
continue
y_n_message = ("Are you sure you want "
"to rename the class "
"\"" + current_class_name + "\" "
"into \"" + new_class_name + "\"?"
)
if query_yes_no(y_n_message, bypass=args.yes):
os.chdir("../ground-truth")
rename_class(current_class_name, new_class_name)
os.chdir("../predicted")
rename_class(current_class_name, new_class_name)
print('Done!')
|