Spaces:
Sleeping
Sleeping
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) | |