build_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build Dataset Structure\n",
    "Copies all images from different dataset folders in a dataset structure with _test_, _train_ and _validation_ subfolders according to a given ratio (70%/20%/10% by default). See example below.  \n",
    "\n",
    "Images naming convention:  ```{shape}_{dataset}_{set}_{id}.{extension}```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Source Datasets -------------------------------------------------------------\n",
    "'''\n",
    "[d1]\n",
    "    [rock]\n",
    "        name1.ext, name2.ext, ...\n",
    "    [paper]\n",
    "        name1.ext, name2.ext, ...\n",
    "    [scissors]\n",
    "        name1.ext, name2.ext, ...\n",
    "\n",
    "[d2]\n",
    "    [rock]\n",
    "        name1.ext, name2.ext, ...\n",
    "    [paper]\n",
    "        name1.ext, name2.ext, ...\n",
    "    [scissors]\n",
    "        name1.ext, name2.ext, ...\n",
    "'''\n",
    "\n",
    "# Target Structure ------------------------------------------------------------\n",
    "'''\n",
    "\n",
    "[train]\n",
    "    [rock]\n",
    "        rock_d1_train_1.ext, rock_d2_train_2.ext, ...\n",
    "    [paper]\n",
    "        paper_d2_train_1.ext, paper_d1_train_2.ext, ...\n",
    "    [scissors]\n",
    "        scissors_d1_train_1.ext, scissors_d2_train_2.ext, ...\n",
    "\n",
    "[test]\n",
    "    [rock]\n",
    "        rock_d1_test_1.ext, rock_d1_test_2.ext, ...\n",
    "    [paper]\n",
    "        paper_d2_test_1.ext, paper_d2_test_2.ext, ...\n",
    "    [scissors]\n",
    "        scissors_d1_test_1.ext, scissors_d2_test_2.ext, ...\n",
    "\n",
    "[validation]\n",
    "    [rock]\n",
    "        rock_d2_validation_1.ext, rock_d1_validation_2.ext, ...\n",
    "    [paper]\n",
    "        paper_d2_validation_1.ext, paper_d2_validation_2.ext, ...\n",
    "    [scissors]\n",
    "        scissors_d1_validation_1.ext, scissors_d1_validation_2.ext, ...\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "import random\n",
    "\n",
    "root_path = 'dataset'\n",
    "dataset_paths = []  # dataset folders to use, all if empty\n",
    "split_ratio = [0.7, 0.2, 0.1]  # train / test / validation ratios\n",
    "classes = ['rock', 'paper', 'scissors']\n",
    "steps = ['train', 'test', 'validation']\n",
    "img_ext = ['jpg', 'jpeg', 'png']\n",
    "\n",
    "# define ids to be used in the new dataset img names, to avoid overwriting\n",
    "train_id = 0\n",
    "test_id = 0\n",
    "validation_id = 0\n",
    "train_img_count = 0\n",
    "test_img_count = 0\n",
    "validation_img_count = 0\n",
    "\n",
    "try:\n",
    "    # list all directories in root_path if dataset_paths is empty\n",
    "    if len(dataset_paths) == 0:\n",
    "        dataset_paths = [d for d in os.listdir(root_path) if os.path.isdir(os.path.join(root_path, d))]\n",
    "\n",
    "    print(f'Loading {len(dataset_paths)} datasets: ', dataset_paths)\n",
    "\n",
    "    # build structure\n",
    "    for s in steps:\n",
    "        for c in classes:\n",
    "            os.makedirs(os.path.join(root_path, s, c))\n",
    "\n",
    "    # copy files\n",
    "    for d in dataset_paths:\n",
    "\n",
    "        for c in classes:\n",
    "\n",
    "            # list imaged in directory\n",
    "            source_dir = os.path.join(root_path, d, c)\n",
    "            img_list = [img for img in os.listdir(source_dir) if img.split('.')[-1].lower() in img_ext]\n",
    "            img_len = len(img_list)\n",
    "\n",
    "            # create 3 random lists of images for train, test and validation based on split ration\n",
    "            train_img_len = int(img_len * split_ratio[0])\n",
    "            train_img_list = random.sample(img_list, train_img_len)\n",
    "            test_img_len = int(img_len * split_ratio[1])\n",
    "            test_img_list = random.sample([_ for _ in img_list if _ not in train_img_list], test_img_len)\n",
    "            validation_img_list = [_ for _ in img_list if _ not in train_img_list and _ not in test_img_list]\n",
    "\n",
    "            # copy train images\n",
    "            train_target_dir = os.path.join(root_path, 'train', c)\n",
    "            for img in train_img_list:\n",
    "                source_filepath = os.path.join(source_dir, img)\n",
    "                ext = img.split('.')[-1]\n",
    "                target_filepath = os.path.join(train_target_dir, f'{c}_{d}_train_{train_id}.{ext}')\n",
    "                shutil.copy(source_filepath, target_filepath)\n",
    "                to_print = f'{source_filepath} -> {target_filepath}'\n",
    "                print(f'{to_print:<200}', end='\\r', flush=True)\n",
    "                train_id += 1\n",
    "                train_img_count += 1\n",
    "\n",
    "            # copy test images\n",
    "            test_target_dir = os.path.join(root_path, 'test', c)\n",
    "            for img in test_img_list:\n",
    "                source_filepath = os.path.join(source_dir, img)\n",
    "                ext = img.split('.')[-1]\n",
    "                target_filepath = os.path.join(test_target_dir, f'{c}_{d}_test_{test_id}.{ext}')\n",
    "                shutil.copy(source_filepath, target_filepath)\n",
    "                to_print = f'{source_filepath} -> {target_filepath}'\n",
    "                print(f'{to_print:<200}', end='\\r', flush=True)\n",
    "                test_id += 1\n",
    "                test_img_count += 1\n",
    "\n",
    "            # copy validation images\n",
    "            validation_target_dir = os.path.join(root_path, 'validation', c)\n",
    "            for img in validation_img_list:\n",
    "                source_filepath = os.path.join(source_dir, img)\n",
    "                ext = img.split('.')[-1]\n",
    "                target_filepath = os.path.join(validation_target_dir, f'{c}_{d}_validation_{validation_id}.{ext}')\n",
    "                shutil.copy(source_filepath, target_filepath)\n",
    "                to_print = f'{source_filepath} -> {target_filepath}'\n",
    "                print(f'{to_print:<200}', end='\\r', flush=True)\n",
    "                validation_id += 1\n",
    "                validation_img_count += 1\n",
    "\n",
    "except Exception as e:\n",
    "    print('\\nError creating dataset structure: ', e)\n",
    "\n",
    "print(f'\\nImages copied successfully using split ratio {split_ratio}.')\n",
    "print(f'Train: {train_img_count} | Test: {test_img_count} | Validation: {validation_img_count} | Total: {train_img_count + test_img_count + validation_img_count}')\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "keras",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}