Spaces:
Sleeping
Sleeping
File size: 3,150 Bytes
908e980 |
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 |
import os
import shutil
from tqdm import tqdm
from sklearn.model_selection import train_test_split
PROJECT_DIR = os.path.dirname(os.path.dirname(__file__))
def create_output_folders(train_path, test_path, val_path):
os.makedirs(train_path, exist_ok=True)
os.makedirs(test_path, exist_ok=True)
os.makedirs(val_path, exist_ok=True)
os.makedirs(os.path.join(train_path, "images"), exist_ok=True)
os.makedirs(os.path.join(train_path, "labels"), exist_ok=True)
os.makedirs(os.path.join(test_path, "images"), exist_ok=True)
os.makedirs(os.path.join(test_path, "labels"), exist_ok=True)
os.makedirs(os.path.join(val_path, "images"), exist_ok=True)
os.makedirs(os.path.join(val_path, "labels"), exist_ok=True)
def copy_images_and_labels(src_path, dst_path, labels_path, folder_name, image_filenames):
print(f"Copying {folder_name} images and labels...")
for image_filename in tqdm(image_filenames):
# Copy the image file to the folder
src_img_path = os.path.join(src_path, image_filename)
dst_img_path = os.path.join(dst_path, "images", image_filename)
shutil.copy(src_img_path, dst_img_path)
# Copy the corresponding label file to the folder with the same name
label_filename = os.path.splitext(image_filename)[0] + ".txt"
src_label_path = os.path.join(labels_path, label_filename)
dst_label_path = os.path.join(dst_path, "labels", label_filename)
shutil.copy(src_label_path, dst_label_path)
def split_data(images_path, labels_path, train_path, test_path, val_path, test_size=0.1, val_size=0.05, shuffle=True):
# Set the paths for the train, test, and validation folders
create_output_folders(train_path, test_path, val_path)
# Get a list of all image filenames in the images folder
image_filenames = [f for f in os.listdir(images_path) if os.path.isfile(os.path.join(images_path, f))]
# Split the image filenames into train, test, and validation sets
train_image_filenames, test_image_filenames = train_test_split(image_filenames, test_size=test_size, shuffle=shuffle)
train_image_filenames, val_image_filenames = train_test_split(train_image_filenames, test_size=val_size, shuffle=shuffle)
# Copy train images and labels
copy_images_and_labels(images_path, train_path, labels_path, "train", train_image_filenames)
# Copy test images and labels
copy_images_and_labels(images_path, test_path, labels_path, "test", test_image_filenames)
# Copy validation images and labels
copy_images_and_labels(images_path, val_path, labels_path, "validation", val_image_filenames)
if __name__ == "__main__":
images_path = os.path.join(PROJECT_DIR, 'data', 'yolo_format', 'images')
labels_path = os.path.join(PROJECT_DIR, 'data', 'yolo_format', 'labels')
train_path = os.path.join(PROJECT_DIR, 'data', 'yolo_format', 'dataset', 'train')
test_path = os.path.join(PROJECT_DIR, 'data', 'yolo_format', 'dataset', 'test')
val_path = os.path.join(PROJECT_DIR, 'data', 'yolo_format', 'dataset', 'val')
split_data(images_path, labels_path, train_path, test_path, val_path)
|