data_wrangling.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data selection & augmentation\n",
    "\n",
    "TODO:\n",
    "- speed up..."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import re\n",
    "import shutil\n",
    "from glob import glob\n",
    "from pathlib import Path\n",
    "from typing import Tuple\n",
    "from warnings import warn\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "from PIL import Image\n",
    "\n",
    "# from torch.utils.data import Subset\n",
    "from torchvision import transforms\n",
    "from torchvision.datasets import ImageFolder\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "from src.utils_Img2Img import print_grid"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data selection\n",
    "\n",
    "Performs the following:\n",
    "- selects *the same amount* (number of images) of data, at random, from each *super*-class of the original dataset, where a super-class is a first-level class, ie:\n",
    "```\n",
    "- latrunculin_B_high_conc     <- super-class\n",
    "      - latrunculin B 30      <- not a super-class\n",
    "            - file1\n",
    "            - file2\n",
    "      - latrunculin B 10      <- not a super-class either\n",
    "            - file\n",
    "```\n",
    "- splits the data into train and test sets"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Original dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the original dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "augmentations = transforms.Compose(\n",
    "    [\n",
    "        transforms.Resize(128, interpolation=transforms.InterpolationMode.BILINEAR),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.5], [0.5]),  # map to [-1, 1] for SiLU\n",
    "    ]\n",
    ")\n",
    "\n",
    "data_root_path = \"/projects/imagesets/Golgi/128x128/\"\n",
    "\n",
    "dataset = ImageFolder(\n",
    "    root=data_root_path,\n",
    "    transform=lambda x: augmentations(x.convert(\"RGB\")),\n",
    "    target_transform=lambda y: torch.tensor(y).long(),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset ImageFolder\n",
       "    Number of datapoints: 175020\n",
       "    Root location: /projects/imagesets/Golgi/128x128/\n",
       "    StandardTransform\n",
       "Transform: <function <lambda> at 0x7f12a788d870>\n",
       "Target transform: <function <lambda> at 0x7f12a788d990>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check original dataset stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Class names and indices: [('0', 0), ('1', 1)]\n",
      "Counts: {0: 89446, 1: 85574}\n",
      "Common number of samples to select per class: 85574\n"
     ]
    }
   ],
   "source": [
    "unique, counts = np.unique(dataset.targets, return_counts=True)\n",
    "print(\"Class names and indices:\", [\n",
    "      (cl, dataset.class_to_idx[cl]) for cl in dataset.classes])\n",
    "print(\"Counts:\", dict(zip(unique, counts)))\n",
    "\n",
    "common_nb_samples = np.min(counts)\n",
    "print(\"Common number of samples to select per class:\", common_nb_samples)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define train/test split sizes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "split_approx_props = {\"train\": 0.5, \"test\": 0.5}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "split_sizes: {'train': 42787, 'test': 42787}\n"
     ]
    }
   ],
   "source": [
    "assert sum(split_approx_props.values()) == 1.0\n",
    "\n",
    "split_sizes = dict.fromkeys(split_approx_props.keys(), 0)\n",
    "\n",
    "cumsum = 0\n",
    "for idx, (split_name, prop) in enumerate(split_approx_props.items()):\n",
    "    if idx == 1: # fill the test split with the remaining samples\n",
    "        assert split_name == \"test\"\n",
    "        split_size = common_nb_samples - cumsum\n",
    "    else:\n",
    "        split_size = int(prop * common_nb_samples)\n",
    "    cumsum += split_size\n",
    "    split_sizes[split_name] = split_size\n",
    "assert cumsum == common_nb_samples\n",
    "\n",
    "print(\"split_sizes:\", split_sizes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Helper function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_full_names(selected_files: list[str], src_dir: Path) -> list[Tuple[Path, str]]:\n",
    "    \"\"\"Returns the (source, target) file path pairs,\n",
    "    with possibly the intermediary directory name prefixed to it. eg:\n",
    "\n",
    "    - DMSO (= src_dir)\n",
    "        - file1\n",
    "        - file2\n",
    "\n",
    "    would give [(file1, file1), (file2, file2)], but:\n",
    "\n",
    "    - latrunculin_B_high_conc (= src_dir)\n",
    "        - latrunculin B 30\n",
    "            - file1\n",
    "            - file2\n",
    "      - latrunculin B 10\n",
    "            - file1\n",
    "\n",
    "    would give [\n",
    "        (latrunculin B 30/file1, latrunculin B 30_file1), \n",
    "        (latrunculin B 30/file2, latrunculin B 30_file2), \n",
    "        (latrunculin B 10/file1, latrunculin B 10_file1)\n",
    "    ]                    ^                       ^     \n",
    "                         |                       |\n",
    "    (notice the         '/'          vs         '_')\n",
    "    \"\"\"\n",
    "    selected_files_full_names = []\n",
    "\n",
    "    for file in selected_files:\n",
    "        # find the last part(s) (2 max)\n",
    "        parts = Path(file).parts\n",
    "        min_len = min(len(parts), len(src_dir.parts))\n",
    "        idx = 0\n",
    "        while idx < min_len and parts[idx] == src_dir.parts[idx]:\n",
    "            idx += 1\n",
    "        last_parts = parts[idx:]\n",
    "        assert len(last_parts) <= 2\n",
    "        # form the (Path, str) pair\n",
    "        to_append = (Path(*last_parts), '_'.join(last_parts))\n",
    "        selected_files_full_names.append(to_append)\n",
    "\n",
    "    return selected_files_full_names"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Target directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "aug_data_root_path = \"/projects/deepdevpath/Thomas/data/Golgi\"  # train/test split done in the copy script"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "notebookRunGroups": {
     "groupValue": ""
    }
   },
   "source": [
    "# Data selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# do_copy=False means dry-run (no data will be written)\n",
    "do_copy = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Select & copy data into splits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3fcf9340c0b54e0490c8d522b0542207",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "C:   0%|          | 0/171148 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Destination directory: /projects/deepdevpath/Thomas/data/Golgi\n",
      "Creating /projects/deepdevpath/Thomas/data/Golgi\n",
      "\n",
      "===========================> 0\n",
      "Source directory:      /projects/imagesets/Golgi/128x128/0\n",
      "Found 89446 PNG files\n",
      "Creating /projects/deepdevpath/Thomas/data/Golgi/train/0\n",
      "Creating /projects/deepdevpath/Thomas/data/Golgi/test/0\n",
      "\n",
      "===========================> 1\n",
      "Source directory:      /projects/imagesets/Golgi/128x128/1\n",
      "Found 85574 PNG files\n",
      "Creating /projects/deepdevpath/Thomas/data/Golgi/train/1\n",
      "Creating /projects/deepdevpath/Thomas/data/Golgi/test/1\n"
     ]
    }
   ],
   "source": [
    "desc = \"Copying files\" if do_copy else \"*Not* copying files\"\n",
    "\n",
    "pbar = tqdm(\n",
    "    total=common_nb_samples * len(dataset.classes),\n",
    "    desc=desc,\n",
    ")\n",
    "\n",
    "# Create general destination directory\n",
    "dst_dir = Path(aug_data_root_path)\n",
    "print(f\"Destination directory: {dst_dir}\")\n",
    "if dst_dir.exists():\n",
    "    if do_copy:\n",
    "        raise RuntimeError(f\"{dst_dir} already exists!\")\n",
    "    else:\n",
    "        warn(f\"{dst_dir} already exists! But no copy will be made...\")\n",
    "if do_copy:\n",
    "    print(f\"Creating {dst_dir}\")\n",
    "    dst_dir.mkdir(parents=True, exist_ok=False)\n",
    "else:\n",
    "    print(f\"Would create {dst_dir}\")\n",
    "\n",
    "# fill the class splits\n",
    "for class_name in dataset.classes:\n",
    "    print(f\"\\n===========================> {class_name}\")\n",
    "    # Set the source and destination directories\n",
    "    src_dir = Path(data_root_path, class_name)\n",
    "    print(f\"Source directory:      {src_dir}\")\n",
    "\n",
    "    # Get a list of all PNG files in the source directory\n",
    "    # A single class might have multiple subdirectories\n",
    "    # so we need some glob magic\n",
    "    pathname = src_dir.as_posix() + \"/**/*.png\"\n",
    "    png_file_paths = list(glob(pathname, recursive=True))\n",
    "    print(f\"Found {len(png_file_paths)} PNG files\")\n",
    "\n",
    "    # Check if there are enough PNG files to select from\n",
    "    if len(png_file_paths) < common_nb_samples:\n",
    "        raise ValueError(f\"Not enough PNG files in {src_dir} to select from.\")\n",
    "\n",
    "    # Select a random subset of `common_nb_samples` images files\n",
    "    selected_files = random.sample(png_file_paths, common_nb_samples)\n",
    "\n",
    "    # Split this selection into train/test splits\n",
    "    train_files = selected_files[: split_sizes[\"train\"]]\n",
    "    test_files = selected_files[split_sizes[\"train\"] :]\n",
    "    assert len(train_files) + len(test_files) == common_nb_samples\n",
    "    files_dict = {\"train\": train_files, \"test\": test_files}\n",
    "\n",
    "    for split, files in files_dict.items():\n",
    "        # Copy the selected files to the destination directory\n",
    "        # first get their name (+ possibly intermediary class)\n",
    "        selected_files_full_names = get_full_names(files, src_dir)\n",
    "\n",
    "        split_class_dir = Path(dst_dir, split, class_name)\n",
    "        if not split_class_dir.exists():\n",
    "            if do_copy:\n",
    "                print(f\"Creating {split_class_dir}\")\n",
    "                split_class_dir.mkdir(parents=True, exist_ok=False)\n",
    "            else:\n",
    "                print(f\"Would create {split_class_dir}\")\n",
    "\n",
    "        # then copy\n",
    "        for srcfilename, trgtfilename in selected_files_full_names:\n",
    "            src_path = Path(src_dir, srcfilename)\n",
    "            dst_path = Path(dst_dir, split, class_name, trgtfilename)\n",
    "            pbar.set_postfix_str(f\"{srcfilename} -> {split}/{trgtfilename}\")\n",
    "            if do_copy:\n",
    "                if dst_path.exists():\n",
    "                    raise RuntimeError(f\"{dst_path} already exists!\")\n",
    "                shutil.copy(src_path, dst_path)\n",
    "            pbar.update()\n",
    "        if not do_copy:\n",
    "            print(\n",
    "                f\"    Would have copied {len(selected_files_full_names)} files to {split_class_dir}\"\n",
    "            )\n",
    "\n",
    "pbar.close()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data augmentation\n",
    "\n",
    "The symmetry group of a square $\\mathrm{Dih}_4$ is of order 8, so given the assumed semantic invariance of our 2D squared images (no up and down or right and left, no bottom and above), we can achieve up to 8x times data augmentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "OPS = (\n",
    "    \"rot1\",\n",
    "    \"rot2\",\n",
    "    \"flip2rot2\",\n",
    "    \"flip1rot2\",\n",
    "    \"rot3\",\n",
    "    \"flip2rot3\",\n",
    "    \"flip1rot3\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def perform_data_aug_op(image: np.ndarray, op_code: str):\n",
    "    \"\"\"\n",
    "    Performs the data augmentation operation given by op_code on the image.\n",
    "\n",
    "    Arguments\n",
    "    =========\n",
    "    - image: array, shape (3, res, res)\n",
    "\n",
    "    - op_code: str\n",
    "\n",
    "    Returns\n",
    "    =======\n",
    "    - modified_image: array, same shape\n",
    "    \"\"\"\n",
    "\n",
    "    assert op_code in OPS, \"op_code not in OPS_LIST\"\n",
    "\n",
    "    image = image.copy()\n",
    "\n",
    "    if op_code == \"id\":\n",
    "        modified_image = image\n",
    "    elif re.match(r\"rot\\d\", op_code):\n",
    "        res = re.match(r\"rot\\d\", op_code).group()\n",
    "        r = int(res[-1])\n",
    "        modified_image = np.rot90(image, r, (1, 2))\n",
    "    elif re.match(r\"flip\\drot\\d\", op_code):\n",
    "        res = re.match(r\"flip\\drot\\d\", op_code).group()\n",
    "        r, f = int(res[-1]), int(res[-5])\n",
    "        modified_image = np.flip(np.rot90(image, r, (1, 2)), f)\n",
    "    else:\n",
    "        raise ValueError(\"No match!?\")\n",
    "\n",
    "    return modified_image"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check the augmentations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for class_name in dataset.classes:\n",
    "    print(f\"\\n===========================> {class_name}\")\n",
    "    for split in [\"train\", \"test\"]:\n",
    "        # Set the source and destination directories\n",
    "        dir = Path(aug_data_root_path, split, class_name)\n",
    "        print(f\"Directory: {dir}\")\n",
    "\n",
    "        # Get a list of all PNG files in the directory\n",
    "        png_file_paths = list(dir.glob('**/*.png'))\n",
    "        assert len(png_file_paths) == split_sizes[split]\n",
    "        print(f\"Found {len(png_file_paths)} PNG files\")\n",
    "\n",
    "        # Visualize data aug\n",
    "        img_list = []\n",
    "        for file in png_file_paths:\n",
    "            image = np.asarray(Image.open(file))\n",
    "            image = image.transpose(2, 0, 1)\n",
    "            img_list.append(Image.open(file))\n",
    "            for op_code in OPS:\n",
    "                modified_image = perform_data_aug_op(image, op_code)\n",
    "                img_list.append(Image.fromarray(modified_image.transpose(1, 2, 0)))\n",
    "            break\n",
    "\n",
    "        print_grid(img_list)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Augment the images *in-place*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pbar = tqdm(total=common_nb_samples * len(dataset.classes) * len(OPS))\n",
    "\n",
    "for class_name in dataset.classes:\n",
    "    print(f\"\\n===========================> {class_name}\")\n",
    "    for split in [\"train\", \"test\"]:\n",
    "        # Set the source and destination directories\n",
    "        dir = Path(aug_data_root_path, split, class_name)\n",
    "        print(f\"Directory: {dir}\")\n",
    "\n",
    "        # Get a list of all PNG files in the directory\n",
    "        png_file_paths = list(dir.glob('**/*.png'))\n",
    "        assert len(png_file_paths) == split_sizes[split]\n",
    "        print(f\"Found {len(png_file_paths)} PNG files\")\n",
    "\n",
    "        # Perform data aug\n",
    "        for file in png_file_paths:\n",
    "            array = np.asarray(Image.open(file)).transpose(2, 0, 1)\n",
    "\n",
    "            for op_code in OPS:\n",
    "                modified_array = perform_data_aug_op(array, op_code)\n",
    "                modif_filename = file.stem + \"_\" + op_code + \".png\"\n",
    "                modif_filepath = Path(file.parent, modif_filename)\n",
    "                assert not modif_filepath.exists()\n",
    "                img = Image.fromarray(modified_array.transpose(1, 2, 0))\n",
    "                img.save(modif_filepath)\n",
    "                pbar.update()\n",
    "\n",
    "pbar.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Quick dirty check (adapt to correct folder names!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "42788\n",
      "42788\n",
      "42788\n",
      "42788\n"
     ]
    }
   ],
   "source": [
    "!cd $aug_data_root_path && ls -l train/0 | wc -l && ls -l train/1 | wc -l && ls -l test/0 | wc -l && ls -l test/1 | wc -l"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}