提交 4d816b61 编写于 作者: T tink2123

make label for paddleocr

上级 d31effaf
import os
import argparse
def gen_rec_label(input_path, out_label):
out_file = open(out_label, 'w')
with open(input_path, 'r') as f:
for line in f.readlines():
tmp = line.strip('\n').replace(" ", "").split(',')
img_path, label = tmp[0], tmp[1]
label = label.replace("\"", "")
out_file.write(img_path + '\t' + label + '\n')
out_file.close()
def gen_det_label(input_dir, out_label):
root_path = ""
if "training" in input_dir:
root_path = "icdar_c4_train_imgs/"
elif "test" in input_dir:
root_path = "ch4_test_images/"
out_file = open(out_label, 'w')
for label_file in os.listdir(input_dir):
img_path = root_path + label_file[3:-4] + ".jpg"
label = []
with open(os.path.join(input_dir, label_file), 'r') as f:
for line in f.readlines():
tmp = line.strip("\n\r").replace("\xef\xbb\xbf", "").split(',')
points = tmp[:-2]
s = []
for i in range(0, len(points), 2):
b = points[i:i + 2]
s.append(b)
result = {"transcription": tmp[-1], "points": s}
label.append(result)
out_file.write(img_path + '\t' + str(label) + '\n')
out_file.close()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
'--mode',
type=str,
default="rec",
help='Generate rec_label or det_label, can be set rec or det')
parser.add_argument(
'--input_path',
type=str,
default=".",
help='Input_label or input path to be converted')
parser.add_argument(
'--output_label',
type=str,
default="out_label.txt",
help='Output file name')
args = parser.parse_args()
if args.mode == "rec":
print("Generate rec label")
gen_rec_label(args.input_path, args.output_label)
elif args.mode == "det":
gen_det_label(args.input_path, args.output_label)
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册