图片与标签数据合并 脚本编写

由于官方提供的数据集过大，因此放在了两个文件夹中。本脚本的目的是把label文件合并在一起，把图片合并在一个文件夹中。

import os
import shutil
import argparse

def concat_labels(label1, label2, output):
    label_list = []
    file1 = open(label1, "r")
    file2 = open(label2, "r")
    for line in file1.readlines():
        item = line.strip()
        label_list.append(item)
    file1.close()
    for line in file2.readlines():
        item = line.strip()
        label_list.append(item)
    file2.close()
    
    print(len(label_list))
    save_path = os.path.join(output, "cls_gt_train.txt")
    new_file = open(save_path, "w")
    for i in range(len(label_list)):
        print(label_list[i])
        new_file.write(label_list[i]+ "\n")
    new_file.close()


def concat_images(img_dir, out_dir):
    for file_dir in os.listdir(img_dir):
        son_dir = os.path.join(img_dir, file_dir)
        if not os.path.isdir(son_dir):
            continue
        for _, _, files in os.walk(son_dir):
            if len(files) > 0:
                for f in files:
                    new_dir = os.path.join(son_dir, f)
                    print(new_dir)
                    shutil.move(new_dir, out_dir)


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--label1", default="./label.txt")
    parser.add_argument("--label2", default="./labelr.txt")
    parser.add_argument("--output", default="../paddleocr/PaddleOCR/train_data/cls/", help="label path")
    parser.add_argument("--input_dir", default="./train_data/", help="old image dir")
    parser.add_argument("--out_dir", default="../paddleocr/PaddleOCR/train_data/cls/train/", help="new image dir")
    return parser.parse_args()


def main(args):
    if not os.path.isdir(args.output):
        os.makedirs(args.output)
    if not os.path.isdir(args.out_dir):
        os.makedirs(args.out_dir)
    concat_labels(args.label1, args.label2, args.output)
    concat_images(args.input_dir, args.out_dir)


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