Add k-fold cross-validation split (closes #50)

New kfold() function splits data into k folds for cross-validation,
with each fold containing train/ and val/ subdirectories. Defaults to
symlinks to avoid k× disk usage. Also adds tests for symlink and
format filtering features from merged PRs.

Author: jfilter

README.md

@@ -91,6 +91,12 @@ splitfolders.ratio("input_folder", output="output",
 # Set 3 values, e.g. `(300, 100, 100)`, to limit the number of training values.
 splitfolders.fixed("input_folder", output="output",
     seed=1337, fixed=(100, 100), oversample=False, group_prefix=None, formats=None, move=False) # default values
+
+# Split into k folds for cross-validation.
+# Each fold directory contains train/ and val/ subdirectories.
+# Uses symlinks by default to avoid k× disk usage.
+splitfolders.kfold("input_folder", output="output",
+    seed=1337, k=5, group_prefix=None, formats=None, move="symlink") # default values
 ```
 
 Occasionally, you may have things that comprise more than a single file (e.g. picture (.png) + annotation (.txt)).
@@ -109,13 +115,14 @@ Set
 
 ```
 Usage:
-    splitfolders [--output] [--ratio] [--fixed] [--seed] [--oversample] [--group_prefix] [--formats] [--move] folder_with_images
+    splitfolders [--output] [--ratio] [--fixed] [--kfold] [--seed] [--oversample] [--group_prefix] [--formats] [--move] folder_with_images
 Options:
     --output        path to the output folder. defaults to `output`. Get created if non-existent.
     --ratio         the ratio to split. e.g. for train/val/test `.8 .1 .1 --` or for train/val `.8 .2 --`.
     --fixed         set the absolute number of items per validation/test set. The remaining items constitute
                     the training set. e.g. for train/val/test `100 100` or for train/val `100`.
                     Set 3 values, e.g. `300 100 100`, to limit the number of training values.
+    --kfold         split into k folds for cross-validation. e.g. `5` for 5-fold CV. Uses symlinks by default.
     --seed          set seed value for shuffling the items. defaults to 1337.
     --oversample    enable oversampling of imbalanced datasets, works only with --fixed.
     --group_prefix  split files into equally-sized groups based on their prefix
@@ -124,6 +131,7 @@ Options:
     --symlink       symlink(create shortcut) the files instead of copying
 Example:
     splitfolders --ratio .8 .1 .1 -- folder_with_images
+    splitfolders --kfold 5 folder_with_images
 ```
 
 Because of some [Python quirks](https://github.com/jfilter/split-folders/issues/19) you have to prepend ` --` after using `--ratio`.

splitfolders/cli.py

@@ -1,6 +1,6 @@
 import argparse
 
-from .split import fixed, ratio
+from .split import fixed, kfold, ratio
 
 
 def run():
@@ -38,6 +38,12 @@ def run():
             " Set 3 values, e.g. `300 100 100`, to limit the number of training values."
         ),
     )
+    parser.add_argument(
+        "--kfold",
+        type=int,
+        default=None,
+        help="split into k folds for cross-validation. e.g. `5` for 5-fold CV. Each fold has train/ and val/ subdirs.",
+    )
     parser.add_argument(
         "--oversample",
         action="store_true",
@@ -83,20 +89,26 @@ def run():
 
     if args.ratio:
         ratio(args.input, args.output, args.seed, args.ratio, args.group_prefix, args.move, args.formats)
+    elif args.fixed:
+        fixed(
+            args.input,
+            args.output,
+            args.seed,
+            args.fixed,
+            args.oversample,
+            args.group_prefix,
+            args.move,
+            args.formats,
+        )
+    elif args.kfold:
+        kfold(
+            args.input,
+            args.output,
+            args.seed,
+            args.kfold,
+            args.group_prefix,
+            args.move if args.move else "symlink",
+            args.formats,
+        )
     else:
-        if args.fixed:
-            fixed(
-                args.input,
-                args.output,
-                args.seed,
-                args.fixed,
-                args.oversample,
-                args.group_prefix,
-                args.move,
-                args.formats,
-            )
-        else:
-            print(
-                "Please specify either your `--ratio` or your `--fixed` number of items"
-                " for the split. see -h for more help."
-            )
+        print("Please specify either your `--ratio`, `--fixed`, or `--kfold` for the split. see -h for more help.")

splitfolders/split.py

@@ -192,6 +192,49 @@ def fixed(
                     shutil.copy2(str(f_orig), str(f_dest))
 
 
+def kfold(input, output="output", seed=1337, k=5, group_prefix=None, move="symlink", formats=None):
+    if k < 2:
+        raise ValueError("`k` must be 2 or greater.")
+
+    check_input_format(input)
+    valid_extensions(formats)
+
+    if use_tqdm:
+        prog_bar = tqdm(desc="Copying files", unit=" files")
+
+    for class_dir in list_dirs(input):
+        split_class_dir_kfold(class_dir, output, k, seed, prog_bar if use_tqdm else None, group_prefix, move, formats)
+
+    if use_tqdm:
+        prog_bar.close()
+
+
+def split_class_dir_kfold(class_dir, output, k, seed, prog_bar, group_prefix, move, formats):
+    """
+    Splits a class folder into k folds for cross-validation.
+    Each fold directory gets train/ and val/ subdirectories.
+    """
+    files = setup_files(class_dir, seed, group_prefix, formats)
+
+    # Partition files into k roughly equal chunks
+    fold_size = len(files) // k
+    remainder = len(files) % k
+    partitions = []
+    idx = 0
+    for i in range(k):
+        size = fold_size + (1 if i < remainder else 0)
+        partitions.append(files[idx : idx + size])
+        idx += size
+
+    # For each fold, val = partition i, train = all other partitions
+    for i in range(k):
+        val_files = partitions[i]
+        train_files = [f for j, part in enumerate(partitions) if j != i for f in part]
+        fold_output = str(Path(output) / f"fold_{i + 1}")
+        li = [(train_files, "train"), (val_files, "val")]
+        copy_files(li, class_dir, fold_output, prog_bar, move)
+
+
 def group_by_prefix(files, len_pairs):
     """
     Split files into groups of len `len_pairs` based on their prefix.

tests/test_split.py

@@ -4,7 +4,7 @@
 
 import pytest
 
-from splitfolders import fixed, ratio
+from splitfolders import fixed, kfold, ratio
 
 
 def test_second_package():
@@ -251,6 +251,95 @@ def test_split_fixed_oversample_prefix():
     assert a != b
 
 
+def test_split_ratio_symlink():
+    input_dir = os.path.join(os.path.dirname(__file__), "imgs")
+    output_dir = os.path.join(os.path.dirname(__file__), "output")
+
+    shutil.rmtree(output_dir, ignore_errors=True)
+
+    ratio(input_dir, output_dir, move="symlink")
+
+    # ensure the number of pics is the same
+    a = len(list(pathlib.Path(input_dir).glob("**/*.jpg")))
+    b = len(list(pathlib.Path(output_dir).glob("**/*.jpg")))
+    assert a == b
+
+    # ensure all output files are symlinks
+    for f in pathlib.Path(output_dir).rglob("*.jpg"):
+        assert f.is_symlink()
+
+
+def test_split_fixed_symlink():
+    input_dir = os.path.join(os.path.dirname(__file__), "imgs")
+    output_dir = os.path.join(os.path.dirname(__file__), "output")
+
+    shutil.rmtree(output_dir, ignore_errors=True)
+
+    fixed(input_dir, output_dir, fixed=(2, 2), move="symlink")
+
+    # ensure the number of pics is the same
+    a = len(list(pathlib.Path(input_dir).glob("**/*.jpg")))
+    b = len(list(pathlib.Path(output_dir).glob("**/*.jpg")))
+    assert a == b
+
+    # ensure all output files are symlinks
+    for f in pathlib.Path(output_dir).rglob("*.jpg"):
+        assert f.is_symlink()
+
+
+def test_split_ratio_formats():
+    input_dir = os.path.join(os.path.dirname(__file__), "imgs_texts")
+    output_dir = os.path.join(os.path.dirname(__file__), "output")
+
+    shutil.rmtree(output_dir, ignore_errors=True)
+
+    ratio(input_dir, output_dir, formats=[".jpg"])
+
+    # only jpg files should be in output
+    jpg_count = len(list(pathlib.Path(output_dir).glob("**/*.jpg")))
+    txt_count = len(list(pathlib.Path(output_dir).glob("**/*.txt")))
+    assert jpg_count > 0
+    assert txt_count == 0
+
+
+def test_split_fixed_formats():
+    input_dir = os.path.join(os.path.dirname(__file__), "imgs_texts")
+    output_dir = os.path.join(os.path.dirname(__file__), "output")
+
+    shutil.rmtree(output_dir, ignore_errors=True)
+
+    fixed(input_dir, output_dir, fixed=(1, 1), formats=[".txt"])
+
+    # only txt files should be in output
+    jpg_count = len(list(pathlib.Path(output_dir).glob("**/*.jpg")))
+    txt_count = len(list(pathlib.Path(output_dir).glob("**/*.txt")))
+    assert txt_count > 0
+    assert jpg_count == 0
+
+
+def test_split_ratio_formats_multiple():
+    input_dir = os.path.join(os.path.dirname(__file__), "imgs_texts")
+    output_dir = os.path.join(os.path.dirname(__file__), "output")
+
+    shutil.rmtree(output_dir, ignore_errors=True)
+
+    ratio(input_dir, output_dir, formats=[".jpg", ".txt"])
+
+    # both jpg and txt files should be in output
+    jpg_count = len(list(pathlib.Path(output_dir).glob("**/*.jpg")))
+    txt_count = len(list(pathlib.Path(output_dir).glob("**/*.txt")))
+    assert jpg_count > 0
+    assert txt_count > 0
+
+
+def test_invalid_formats():
+    input_dir = os.path.join(os.path.dirname(__file__), "imgs")
+    output_dir = os.path.join(os.path.dirname(__file__), "output")
+
+    with pytest.raises(ValueError):
+        ratio(input_dir, output_dir, formats=["jpg"])
+
+
 def test_split_ratio_prefix_error_1():
     input_dir = os.path.join(os.path.dirname(__file__), "imgs_texts_error_1")
     output_dir = os.path.join(os.path.dirname(__file__), "output")
@@ -269,3 +358,93 @@ def test_split_ratio_prefix_error_2():
 
     with pytest.raises(ValueError):
         ratio(input_dir, output_dir, group_prefix=2)
+
+
+def test_kfold_basic():
+    input_dir = os.path.join(os.path.dirname(__file__), "imgs")
+    output_dir = os.path.join(os.path.dirname(__file__), "output")
+
+    shutil.rmtree(output_dir, ignore_errors=True)
+
+    kfold(input_dir, output_dir, k=5, move=False)
+
+    # verify 5 fold directories created
+    fold_dirs = sorted(pathlib.Path(output_dir).iterdir())
+    assert len(fold_dirs) == 5
+
+    input_count = len(list(pathlib.Path(input_dir).glob("**/*.jpg")))
+
+    for fold_dir in fold_dirs:
+        # each fold should have train/ and val/
+        assert (fold_dir / "train").is_dir()
+        assert (fold_dir / "val").is_dir()
+
+        # train + val should equal input count
+        train_count = len(list((fold_dir / "train").rglob("*.jpg")))
+        val_count = len(list((fold_dir / "val").rglob("*.jpg")))
+        assert train_count + val_count == input_count
+
+        # no overlap between train and val
+        train_files = {f.name for f in (fold_dir / "train").rglob("*.jpg")}
+        val_files = {f.name for f in (fold_dir / "val").rglob("*.jpg")}
+        assert train_files.isdisjoint(val_files)
+
+
+def test_kfold_3():
+    input_dir = os.path.join(os.path.dirname(__file__), "imgs")
+    output_dir = os.path.join(os.path.dirname(__file__), "output")
+
+    shutil.rmtree(output_dir, ignore_errors=True)
+
+    kfold(input_dir, output_dir, k=3, move=False)
+
+    fold_dirs = sorted(pathlib.Path(output_dir).iterdir())
+    assert len(fold_dirs) == 3
+
+
+def test_kfold_symlink():
+    input_dir = os.path.join(os.path.dirname(__file__), "imgs")
+    output_dir = os.path.join(os.path.dirname(__file__), "output")
+
+    shutil.rmtree(output_dir, ignore_errors=True)
+
+    kfold(input_dir, output_dir, k=3)
+
+    # default is symlink
+    for f in pathlib.Path(output_dir).rglob("*.jpg"):
+        assert f.is_symlink()
+
+
+def test_kfold_copy():
+    input_dir = os.path.join(os.path.dirname(__file__), "imgs")
+    output_dir = os.path.join(os.path.dirname(__file__), "output")
+
+    shutil.rmtree(output_dir, ignore_errors=True)
+
+    kfold(input_dir, output_dir, k=3, move=False)
+
+    # move=False means real copies, not symlinks
+    for f in pathlib.Path(output_dir).rglob("*.jpg"):
+        assert not f.is_symlink()
+
+
+def test_kfold_formats():
+    input_dir = os.path.join(os.path.dirname(__file__), "imgs_texts")
+    output_dir = os.path.join(os.path.dirname(__file__), "output")
+
+    shutil.rmtree(output_dir, ignore_errors=True)
+
+    kfold(input_dir, output_dir, k=3, move=False, formats=[".jpg"])
+
+    jpg_count = len(list(pathlib.Path(output_dir).rglob("*.jpg")))
+    txt_count = len(list(pathlib.Path(output_dir).rglob("*.txt")))
+    assert jpg_count > 0
+    assert txt_count == 0
+
+
+def test_kfold_invalid():
+    input_dir = os.path.join(os.path.dirname(__file__), "imgs")
+    output_dir = os.path.join(os.path.dirname(__file__), "output")
+
+    with pytest.raises(ValueError):
+        kfold(input_dir, output_dir, k=1)