create_anno.py
430 Bytes
import os
import pandas as pd
import random
root = '/data1/lxl/data/ocr/generate1108/'
img_path = os.path.join(root, 'img')
img_list = os.listdir(img_path)
random.shuffle(img_list)
train_df = pd.DataFrame(columns=['name'])
val_df = pd.DataFrame(columns=['name'])
train_df['name'] = img_list[:16000]
val_df['name'] = img_list[16000:]
train_df.to_csv(os.path.join(root, 'train.csv'))
val_df.to_csv(os.path.join(root, 'val.csv'))