update

datasets_preprocess/split_train_val_.py

@@ -0,0 +1,69 @@
+# 加入负样本，根据图片数量划分数据集
+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()