Autolens_train_test_split.py

import os
import random
import shutil

# Define folder paths
images_folder = "images_folder"
xml_folder = "xml_folder"
test_images_folder = "test_images"
test_xml_folder = "test_xml"

# Create test folders if they don't exist
os.makedirs(test_images_folder, exist_ok=True)
os.makedirs(test_xml_folder, exist_ok=True)

# Get list of all XML files
xml_files = [f for f in os.listdir(xml_folder) if f.endswith(".xml")]

# Select 20% of files randomly
num_to_move = int(len(xml_files) * 0.2)
selected_files = random.sample(xml_files, num_to_move)

# Move files to test folders
for xml_file in selected_files:
    # Get corresponding image filename
    file_name = os.path.splitext(xml_file)[0]  # Remove extension
    image_path = os.path.join(images_folder, file_name + ".jpg")  # Adjust extension if needed
    xml_path = os.path.join(xml_folder, xml_file)

    # Define destination paths
    test_image_path = os.path.join(test_images_folder, file_name + ".jpg")
    test_xml_path = os.path.join(test_xml_folder, xml_file)

    # Move files if both exist
    if os.path.exists(image_path) and os.path.exists(xml_path):
        shutil.move(image_path, test_image_path)
        shutil.move(xml_path, test_xml_path)
        print(f"Moved: {xml_file} & {file_name}.jpg")
    else:
        print(f"Skipping: {xml_file} (Missing corresponding image)")

print("✅ 20% of files moved successfully!")