Test_train_split_an_image_set_with_metadata.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/emcdona1/fmnh_scripts/blob/main/Test_train_split_an_image_set_with_metadata.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hqgaEXlkQV0P"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from pathlib import Path\n",
        "import shutil\n",
        "import time\n",
        "from random import randint"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uy6ckb9wdVuv"
      },
      "source": [
        "**In the cell below, adjust the values as necessary.**\n",
        "\n",
        "`SPLIT` is the test split.\n",
        "\n",
        "`image_dir` is the directory of the images to be split.  A new folder with the same name + `_split` will be generated in the same parent directory.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SfqCY2jVQzLK"
      },
      "outputs": [],
      "source": [
        "SPLIT = 0.1\n",
        "image_dir = Path('<<put your image folder path here>>')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QN_h4wBjcsR2",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ed312e75-24bf-4fd8-e398-4cf3e0806451"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "WARNING: These directories already exist!\n",
            "[Errno 17] File exists: 'drive/MyDrive/Machine Learning projects - Grainger-funded/3. Handwriting recognition project - Beth McDonald + Keshab Panthi/image_sets/IAM_words_split/train'\n"
          ]
        }
      ],
      "source": [
        "copy_image_dir = Path(image_dir.parent, f'{image_dir.name}_split')\n",
        "train_dir = Path(copy_image_dir, 'train')\n",
        "test_dir = Path(copy_image_dir, 'test')\n",
        "try:\n",
        "    os.makedirs(train_dir)\n",
        "    os.makedirs(test_dir)\n",
        "except FileExistsError as e:\n",
        "    print('WARNING: These directories already exist!')\n",
        "    print(e)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DDSRwpR7wGK_"
      },
      "source": [
        "## 1. Even sampling (one directory of images with a metadata file)\n",
        "\n",
        "Only CSV and Markdown files will be copied to both train and test directories.\n",
        "\n",
        "All other files will be evenly distributed (e.g. every 10th file for 10% train/test split)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M3tWjepLSvqI"
      },
      "outputs": [],
      "source": [
        "images = os.listdir(image_dir)\n",
        "metadata = [f for f in images if Path(f).suffix == '.csv' or Path(f).suffix == '.md']\n",
        "for m in metadata:\n",
        "  images.remove(m)\n",
        "\n",
        "test_set_size = int(len(images) * SPLIT)\n",
        "test_idx = sorted([int(i * (1 // SPLIT)) for i in list(range(test_set_size))], reverse=True)\n",
        "\n",
        "test_images = list()\n",
        "for idx, img in enumerate(test_idx):\n",
        "    test_images.append(images.pop(idx))\n",
        "train_images = images"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BZiMXjiGQzsI"
      },
      "outputs": [],
      "source": [
        "for img in train_images:\n",
        "    shutil.copyfile(Path(image_dir, img), Path(train_dir, img))\n",
        "for img in test_images:\n",
        "    shutil.copyfile(Path(image_dir, img), Path(test_dir, img))\n",
        "for file in metadata:\n",
        "    shutil.copyfile(Path(image_dir, file), Path(train_dir, file))\n",
        "    shutil.copyfile(Path(image_dir, file), Path(test_dir, file))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZssaJYrYwDTm"
      },
      "source": [
        "## 2. Stratified sampling (multiple sub directories of images)\n",
        "\n",
        "Any files in the root folder will be copied to BOTH train and test folders (e.g. metadata, readmes, any data preprocessing)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZgPMQX2-gxlm"
      },
      "outputs": [],
      "source": [
        "for parent, sub_folders, files in os.walk(image_dir):\n",
        "    print(parent)\n",
        "    if Path(parent).relative_to(image_dir) == Path('.'):\n",
        "        for f in files:\n",
        "            src = str(Path(parent, f)).replace(' ', r'\\ ')\n",
        "            dest = str(Path(train_dir, f)).replace(' ', r'\\ ')\n",
        "            !cp $src $dest -u\n",
        "            # shutil.copy(src, dest)\n",
        "            dest = str(Path(test_dir, f)).replace(' ', r'\\ ')\n",
        "            !cp $src $dest -u\n",
        "            # shutil.copy(src, dest)\n",
        "    else:\n",
        "        test_image_indicies = sorted(list(range(int(len(files) * SPLIT))), reverse=True)\n",
        "        test_image_indicies = [int(i * (1 // SPLIT)) for i in test_image_indicies]\n",
        "        test_images = [files[i] for i in test_image_indicies]\n",
        "        subdir = Path(parent).relative_to(image_dir)\n",
        "        for idx in test_image_indicies:\n",
        "            files.pop(idx)\n",
        "        train_images = files\n",
        "        for image in train_images:\n",
        "            src = str(Path(parent, image)).replace(' ', r'\\ ')\n",
        "            dest = str(Path(train_dir, subdir, image)).replace(' ', r'\\ ')\n",
        "            !cp $src $dest -u\n",
        "            # shutil.copy(src, dest)\n",
        "        for image in test_images:\n",
        "            src = str(Path(parent, image)).replace(' ', r'\\ ')\n",
        "            dest = str(Path(test_dir, subdir, image)).replace(' ', r'\\ ')\n",
        "            !cp $src $dest -u\n",
        "            # shutil.copy(src, dest)\n",
        "    for child_folder in sub_folders:\n",
        "        new_folder = Path(parent, child_folder).relative_to(image_dir)\n",
        "        try:\n",
        "            os.makedirs(Path(train_dir, new_folder))\n",
        "            os.makedirs(Path(test_dir, new_folder))\n",
        "        except:\n",
        "            pass"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "Quick utility to split an image set into train/test",
      "provenance": [],
      "mount_file_id": "1fJlOhPTvLVOD4UW5MMnYXo9rvLc1cO3X",
      "authorship_tag": "ABX9TyMiDUcK8Rjj0bRYG8PdJDJT",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}