import os
import re
import numpy as np
import argparse
from skimage import io


# 读取生成的训练集+验证集图片文件绝对路径，并与yolo格式标注文件路径一一对应
def main(train_file: str, output_file: str):
    image_path_list = np.loadtxt(train_file, dtype=str)  # 生成的训练集+验证集图片文件绝对路径
    # print(type(image_path_list))
    if not os.path.exists('data'):
        os.makedirs('data')

    # 对image_path_list中的路径列表进行替换，生成训练集+验证集标注文件绝对路径列表
    annotation_list = list(image_path_list)
    annotation_list = [re.sub(r'JPEGImages', 'labels', s) for s in annotation_list]
    annotation_list = [re.sub(r'.jpg', '.txt', s) for s in annotation_list]

    # lines: 图片文件绝对路径，该图片对应的标注文件中的标注内容，图片的高和宽
    lines = np.array([
        np.array([
            image_path_list[i],
            np.loadtxt(annotation_list[i], dtype=float, ndmin=2),
            np.array(io.imread(image_path_list[i]).shape[0:2])]  # 顺序h，w
        ) for i in range(len(annotation_list))])
    print(lines.shape)
    np.save(output_file, lines)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--train_file', type=str, default='train.txt', help='train.txt file path')
    parser.add_argument('--output_file', type=str, default='data/voc_img_ann.npy', help='output file path')

    opt = parser.parse_args()

    main(opt.train_file, opt.output_file)