Car_VS_Rest / code.txt
Nekshay's picture
Update code.txt
16c3ce8 verified
raw
history blame
2.04 kB
import os
import random
import shutil
from concurrent.futures import ThreadPoolExecutor
# Define paths
dataset_folder = 'path/to/dataset'
train_folder = os.path.join(dataset_folder, 'train')
val_folder = os.path.join(dataset_folder, 'validation')
# Create validation folder if it doesn't exist
os.makedirs(val_folder, exist_ok=True)
# Get all label folders inside train folder
label_folders = [f for f in os.listdir(train_folder) if os.path.isdir(os.path.join(train_folder, f))]
# Function to move images from a specific label folder
def process_label_folder(label_folder, num_threads):
train_label_folder = os.path.join(train_folder, label_folder)
val_label_folder = os.path.join(val_folder, label_folder)
# Create corresponding validation label folder
os.makedirs(val_label_folder, exist_ok=True)
# Get all images in the train/label_folder
all_images = os.listdir(train_label_folder)
total_images = len(all_images)
# Calculate 20% of images for validation
val_size = int(total_images * 0.2)
# Randomly select 20% of the images for validation
val_images = random.sample(all_images, val_size)
# Function to move a single image
def move_image(image):
src = os.path.join(train_label_folder, image)
dest = os.path.join(val_label_folder, image)
shutil.move(src, dest)
# Use ThreadPoolExecutor to move images in parallel
with ThreadPoolExecutor(max_workers=num_threads) as executor:
executor.map(move_image, val_images)
print(f"Moved {val_size} images from {label_folder} to validation folder.")
# Main function to get user input for number of threads and process folders
def main():
# Ask user for the number of threads
num_threads = int(input("Enter the number of threads to use: "))
# Process each label folder using the input number of threads
for label_folder in label_folders:
process_label_folder(label_folder, num_threads)
print("Validation dataset created.")
if __name__ == "__main__":
main()