# 加入负样本，根据图片数量划分数据集
import os
import random
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--img_path', default='/home/lc/data_center/gift/v2/images', type=str,
                    help='input xml label path')  # 图片存放地址
# 数据集的划分，地址选择自己数据下的ImageSets/Main
# parser.add_argument('--txt_path', default='/home/lc/data_center/gift/yolov10_data/Main', type=str, help='output txt label path')
parser.add_argument('--txt_path', default='/home/lc/data_center/gift/yolov10_data/Main', type=str,
                    help='output txt label path')
opt = parser.parse_args()

trainval_percent = 1.0
train_percent = 0.8
val_percent = 1.0
# train_percent = 1.0
# val_percent = 0.0

imgfilepath = opt.img_path
txtsavepath = opt.txt_path
total_img = os.listdir(imgfilepath)
if not os.path.exists(txtsavepath):
    os.makedirs(txtsavepath)

num = len(total_img)
print("all num:", num)
list_index = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list_index, tv)
train = random.sample(trainval, tr)

val_test = []
for i in trainval:
    if i not in train:
        val_test.append(i)
num_ = len(val_test)
print("val-test num:", num_)
list_index_ = range(num_)
va = int(num_ * val_percent)
val = random.sample(val_test, va)

file_trainval = open(txtsavepath + '/trainval.txt', 'w')
file_test = open(txtsavepath + '/test.txt', 'w')
file_train = open(txtsavepath + '/train.txt', 'w')
file_val = open(txtsavepath + '/val.txt', 'w')

addtrain_path = ""
for i in list_index:
    name = total_img[i][:-4] + '\n'
    addimg_name = name.strip() + ".jpg"
    # print(addimg_name,type(addimg_name),len(addimg_name))
    if i in trainval:
        file_trainval.write(name)
        # if addimg_name in os.listdir(addtrain_path):#把某些数据加入训练集中
        #     print("addimg_name:",addimg_name)
        #     file_train.write(name)
        if i in train:
            file_train.write(name)
        if i in val:
            file_val.write(name)
        if (i not in train) and (i not in val):
            file_test.write(name)

file_trainval.close()
file_train.close()
file_val.close()
file_test.close()