mirror of
https://gitee.com/nanjing-yimao-information/ieemoo-ai-gift.git
synced 2025-08-18 13:20:25 +00:00
70 lines
2.1 KiB
Python
70 lines
2.1 KiB
Python
# 加入负样本，根据图片数量划分数据集
|
||
import os
|
||
import random
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument('--img_path', default='/home/lc/data_center/gift/ori_image/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')
|
||
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()
|