Spaces:
Sleeping
Sleeping
File size: 3,046 Bytes
a8ea7fe | 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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | import os
import random
import shutil
# Base paths
BASE_DIR = "data"
SOURCE_DIR = os.path.join(BASE_DIR, "images")
TRAIN_DIR = os.path.join(BASE_DIR, "train")
VAL_DIR = os.path.join(BASE_DIR, "val")
TEST_DIR = os.path.join(BASE_DIR, "test")
# Supported image extensions
IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff")
# Train/Val/Test split ratios
SPLIT_RATIOS = {
"train": 0.7,
"val": 0.15,
"test": 0.15,
}
random.seed(42) # For reproducibility
def ensure_dir(path: str):
os.makedirs(path, exist_ok=True)
def is_image(fname: str) -> bool:
return fname.lower().endswith(IMAGE_EXTS)
def check_empty(path: str):
"""Ensure the directory is empty before writing."""
if os.path.exists(path) and os.listdir(path):
raise RuntimeError(f"{path} is not empty. Clear it before running.")
def main():
print(f"Source directory: {os.path.abspath(SOURCE_DIR)}")
if not os.path.isdir(SOURCE_DIR):
raise RuntimeError(f"Source directory not found: {SOURCE_DIR}")
# Prepare train/val/test dirs
for split_dir in [TRAIN_DIR, VAL_DIR, TEST_DIR]:
ensure_dir(split_dir)
check_empty(split_dir)
# Get all class folders
class_names = sorted(
d for d in os.listdir(SOURCE_DIR)
if os.path.isdir(os.path.join(SOURCE_DIR, d))
)
print("\nClasses detected:")
for cname in class_names:
print(" -", cname)
total_src = 0
total_split = {"train": 0, "val": 0, "test": 0}
# Process every class folder
for class_name in class_names:
src_dir = os.path.join(SOURCE_DIR, class_name)
print(f"\nProcessing: {class_name}")
files = [f for f in os.listdir(src_dir) if is_image(f)]
files.sort()
if not files:
print(" -> No images found, skipping.")
continue
random.shuffle(files)
n = len(files)
total_src += n
n_train = int(n * SPLIT_RATIOS["train"])
n_val = int(n * SPLIT_RATIOS["val"])
n_test = n - n_train - n_val # Remaining goes to test
splits = {
"train": files[:n_train],
"val": files[n_train:n_train + n_val],
"test": files[n_train + n_val:],
}
for split_name, split_files in splits.items():
dst_dir = os.path.join(BASE_DIR, split_name, class_name)
ensure_dir(dst_dir)
print(f" {split_name}: {len(split_files)} images")
total_split[split_name] += len(split_files)
for fname in split_files:
shutil.copy2(os.path.join(src_dir, fname),
os.path.join(dst_dir, fname))
print(f"Total images: {total_src}")
print(f"Train: {total_split['train']}")
print(f"Val: {total_split['val']}")
print(f"Test: {total_split['test']}")
print("Dataset split completed.")
if __name__ == "__main__":
main()
|