File size: 2,893 Bytes
1060621 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import os
import shutil
import random
from tqdm import tqdm
def main():
# Set the dataset path
Dataset_Path = 'CVRPDataset'
img_dir = os.path.join(Dataset_Path, 'img_dir')
ann_dir = os.path.join(Dataset_Path, 'ann_dir')
# Set the ratio of training set to test set
test_frac = 0.1
random.seed(123)
# Create directories if not exist
os.makedirs(os.path.join(Dataset_Path, 'train'), exist_ok=True)
os.makedirs(os.path.join(Dataset_Path, 'val'), exist_ok=True)
# Get image file paths
img_paths = os.listdir(img_dir)
random.shuffle(img_paths)
val_number = int(len(img_paths) * test_frac)
train_files = img_paths[val_number:]
val_files = img_paths[:val_number]
print(f"Total images: {len(img_paths)}")
print(f"Training set images: {len(train_files)}")
print(f"Test set images: {len(val_files)}")
# Move the training set images to the train directory
for each in tqdm(train_files, desc="Move the training set images"):
src_path = os.path.join(img_dir, each)
dst_path = os.path.join(Dataset_Path, 'train', each)
shutil.move(src_path, dst_path)
# Move the test set images to the test directory
for each in tqdm(val_files, desc="Move the test set images"):
src_path = os.path.join(img_dir, each)
dst_path = os.path.join(Dataset_Path, 'val', each)
shutil.move(src_path, dst_path)
# Move the train and val directories into img_dir
shutil.move(os.path.join(Dataset_Path, 'train'), os.path.join(img_dir, 'train'))
shutil.move(os.path.join(Dataset_Path, 'val'), os.path.join(img_dir, 'val'))
# Process annotation files
# Ensure the annotation directories exist
os.makedirs(os.path.join(Dataset_Path, 'train'), exist_ok=True)
os.makedirs(os.path.join(Dataset_Path, 'val'), exist_ok=True)
# Move the training set annotation files to the train directory
for each in tqdm(train_files, desc="Move the training set annotations"):
src_path = os.path.join(ann_dir, each.split('.')[0] + '.png')
dst_path = os.path.join(Dataset_Path, 'train', each.split('.')[0] + '.png')
shutil.move(src_path, dst_path)
# Move the test set annotation files to the test directory
for each in tqdm(val_files, desc="Move the test set annotations"):
src_path = os.path.join(ann_dir, each.split('.')[0] + '.png')
dst_path = os.path.join(Dataset_Path, 'val', each.split('.')[0] + '.png')
shutil.move(src_path, dst_path)
# Move the train and val annotation directories into ann_dir
shutil.move(os.path.join(Dataset_Path, 'train'), os.path.join(ann_dir, 'train'))
shutil.move(os.path.join(Dataset_Path, 'val'), os.path.join(ann_dir, 'val'))
if __name__ == '__main__':
main()
|