OCR方向分类模型标签生成 脚本编写

PaddleOCR官网方向分类模型要求正向图片标签为0，旋转180度后的图片标签为180。生成标签的txt格式为：目录+\t+label的格式，如train/123.jpg. 0

import argparse
import os

def get_label(input, output, label):
    with open(output, "w") as f:
        if os.path.isdir(input):
            for filename in os.listdir(input):
                name, suffix = os.path.splitext(filename)
                if suffix.lower() not in [".jpg", ".jpeg", ".png"]:
                    continue
                string = "train/" + filename
                string1 = '\t'.join([string, label, "\n"])
                print(string1)
                f.write(string1)
        f.close()


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_dir", default="./train")
    parser.add_argument("--output", default="./label.txt")
    parser.add_argument("--label", default="0", help="angle")
    return parser.parse_args()


def main(args):
    get_label(args.input_dir, args.output, args.label)


if __name__ == "__main__":
    main(get_args())