隨機劃分訓練集、測試集
阿新 • • 發佈:2021-01-30
技術標籤:好用的python小程式
只需要把root_path 的路徑修改即可
# -*- coding: utf-8 -*-
from __future__ import division, print_function, absolute_import
import sys
sys.path.append('../../')
import shutil
import os
import random
import math
def mkdir(path):
if not os.path.exists(path):
os.makedirs(path)
divide_rate = 0.82
root_path = '根路徑'
image_path = root_path + '/jpg'
xml_path = root_path + '/xml'
image_list = os.listdir(image_path)
# image_name = [n.split('.')[0] for n in image_list]
image_name = image_list
random.shuffle(image_name)
train_image = image_name[:int(math.ceil(len(image_name)) * divide_rate)]
test_image = image_name[int(math.ceil(len(image_name)) * divide_rate):]
image_output_train = os.path.join(root_path, 'train/img')
# mkdir(image_output_train)
image_output_test = os.path.join(root_path, 'test/img')
# mkdir(image_output_test)
xml_train = os.path.join(root_path, 'train/xml')
# mkdir(xml_train)
xml_test = os.path.join(root_path, 'test/xml')
# mkdir(xml_test)
count = 0
for i in train_image:
shutil.copy(os.path.join(image_path, i), image_output_train)
imgname = i.split('.')[0]
shutil.copy(os.path.join(xml_path, imgname + '.xml'), xml_train)
if count % 1000 == 0:
print("process step {}".format(count))
count += 1
for i in test_image:
shutil.copy(os.path.join(image_path, i), image_output_test)
imgname = i.split('.')[0]
shutil.copy(os.path.join(xml_path, imgname + '.xml'), xml_test)
if count % 1000 == 0:
print("process step {}".format(count))
count += 1