bootstraping_augment_data.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "cf23271f-e679-46bd-8bbe-fd4d17bab2f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.datasets import ImageFolder\n",
    "from torch.utils.data import DataLoader\n",
    "from PIL import Image\n",
    "import random\n",
    "import shutil\n",
    "\n",
    "source_dir = r'/run/media/magnusjsc/T7/Classification-and-3D-reconstruction-of-archaeological-artifacts_DATA/Splitted data by era - GAN project/test_iron'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "47a64122-0dd3-4a57-8d0f-b0b592c5f46e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of images in the original dataset: 1779\n",
      "Intended number of bootstrap images: 7116\n",
      "Bootstrap indicies:  7116\n",
      "Number of images copied: 7116\n",
      "Done bootstrapping the images! :D\n"
     ]
    }
   ],
   "source": [
    "# Bootstrap data - Random select images and copy them\n",
    "image_size = 128\n",
    "batch_size = 128\n",
    "workers = (4 * 1)  # Multiply with the number of GPUs available\n",
    "\n",
    "target_dir_boot = os.path.abspath(os.path.join(source_dir, \"..\", \"bootstrap_iron\"))\n",
    "os.makedirs(target_dir_boot, exist_ok = True)\n",
    "\n",
    "# Create a subdirectory within target_dir_boot to store the bootstrapped images\n",
    "bootstrap_subdir = os.path.join(target_dir_boot, \"bootstrapped_images\")\n",
    "os.makedirs(bootstrap_subdir, exist_ok = True)\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((image_size, image_size)),\n",
    "    transforms.ToTensor(),\n",
    "])\n",
    "\n",
    "dataset = ImageFolder(root = source_dir, transform = transform)\n",
    "# Image indices\n",
    "indices = list(range(len(dataset)))\n",
    "num_images_original = len(dataset)\n",
    "\n",
    "print(\"Number of images in the original dataset:\", num_images_original)\n",
    "\n",
    "# n times the size of the original dataset\n",
    "bootstrap_ratio = 4.0\n",
    "num_bootstrap_images = int(bootstrap_ratio * num_images_original)\n",
    "\n",
    "print(\"Intended number of bootstrap images:\", num_bootstrap_images)\n",
    "\n",
    "# Select a random sequence of images to bootstrap - With replacement\n",
    "bootstrap_indices = random.choices(indices, k = num_bootstrap_images)\n",
    "print(\"Bootstrap indicies: \", len(bootstrap_indices))\n",
    "\n",
    "i = 0\n",
    "for index in bootstrap_indices:\n",
    "    image_path, _ = dataset.samples[index]\n",
    "    image_name = str(i) + \"_\" + os.path.basename(image_path)\n",
    "    target_image_path = os.path.join(bootstrap_subdir, image_name)  # Save in the subdirectory\n",
    "    shutil.copy(image_path, target_image_path)\n",
    "    i += 1\n",
    "\n",
    "num_copied_images = len(os.listdir(bootstrap_subdir))\n",
    "print(\"Number of images copied:\", num_copied_images)\n",
    "print(\"Done bootstrapping the images! :D\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "08f91f1b-7766-4015-9a5e-184ea2907228",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done augmenting the images! :D\n"
     ]
    }
   ],
   "source": [
    "# Augment data - Augment the data and add them \n",
    "target_dir_aug = os.path.abspath(os.path.join(source_dir, \"..\", \"augmented_iron\"))\n",
    "os.makedirs(target_dir_aug, exist_ok=True)\n",
    "\n",
    "# Create a subdirectory within target_dir_aug to store the augmented images\n",
    "augmented_subdir = os.path.join(target_dir_aug, \"augmented_images\")\n",
    "os.makedirs(augmented_subdir, exist_ok = True)\n",
    "\n",
    "augment_transform = transforms.Compose([\n",
    "    transforms.Resize((image_size, image_size)),\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.RandomRotation(90), # 90 degree rotation\n",
    "    transforms.RandomAffine(\n",
    "        degrees = 0, \n",
    "        scale = (0.8, 1.5), # 0.8 is a bit zoom out, 1.2 is a bit zoom in, and 1 is original\n",
    "        translate = None,\n",
    "        shear = None,\n",
    "        center = None, \n",
    "        fill = 0, \n",
    "        # interpolation=InterpolationMode.NEAREST\n",
    "    ), \n",
    "    transforms.ToTensor()\n",
    "])\n",
    "\n",
    "original_transform = transforms.Compose([\n",
    "    transforms.Resize((image_size, image_size)), \n",
    "    transforms.ToTensor(),  \n",
    "])\n",
    "\n",
    "# Load dataset\n",
    "dataset = ImageFolder(root = source_dir, transform = original_transform)\n",
    "\n",
    "def save_augmented_images(images, labels, batch_index):\n",
    "    for i in range(len(images)):\n",
    "        # Convert tensor to PIL image\n",
    "        image = transforms.ToPILImage()(images[i])\n",
    "        image_name = f\"augmented_iron_{labels[i]}_{batch_index * batch_size + i}.jpg\"\n",
    "        image_path = os.path.join(augmented_subdir, image_name)  # Save in the subdirectory\n",
    "        image.save(image_path)\n",
    "\n",
    "# The coins in the iron folder will get the label 0 by default since it's the first folder \n",
    "augmented_data_loader = DataLoader(dataset, batch_size, shuffle = True, num_workers = workers)\n",
    "\n",
    "for batch_index, (images, labels) in enumerate(augmented_data_loader): \n",
    "    # Convert tensors to PIL images\n",
    "    pil_images = [transforms.ToPILImage()(image) for image in images]\n",
    "    # Apply augmentation\n",
    "    augmented_images = [augment_transform(image) for image in pil_images]\n",
    "    augmented_images = torch.stack(augmented_images)\n",
    "    # Save the augmented images \n",
    "    save_augmented_images(augmented_images, labels, batch_index)\n",
    "\n",
    "    # Stop after first run - Debug\n",
    "    # break \n",
    "\n",
    "print(\"Done augmenting the images! :D\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "f27718c6-daa4-4426-9795-da6a831e5a7a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done augmenting the bootstrapped images! :D\n"
     ]
    }
   ],
   "source": [
    "# Augment the bootstraped images\n",
    "source_dir = r'/run/media/magnusjsc/T7/Classification-and-3D-reconstruction-of-archaeological-artifacts_DATA/Splitted data by era - GAN project/test_iron'\n",
    "target_dir_aug_boot = os.path.abspath(os.path.join(source_dir, \"..\", \"augmented_boot_iron\"))\n",
    "os.makedirs(target_dir_aug_boot, exist_ok = True)\n",
    "\n",
    "dataset_aug_boot = ImageFolder(root = target_dir_boot, transform = original_transform)\n",
    "\n",
    "bootstrap_data_loader = DataLoader(\n",
    "    dataset_aug_boot, \n",
    "    batch_size = batch_size, \n",
    "    shuffle = True, \n",
    "    num_workers = workers\n",
    ")\n",
    "\n",
    "def save_bootstrapped_augmented_images(images, labels, index):\n",
    "    for i in range(len(images)):\n",
    "        image = transforms.ToPILImage()(images[i])\n",
    "        image_name = f\"bootstrap_augmented_iron_{labels[i]}_{index * batch_size + i}.jpg\"\n",
    "        image_path = os.path.join(target_dir_aug_boot, image_name)\n",
    "        image.save(image_path)\n",
    "\n",
    "\n",
    "for index, (images, labels) in enumerate(bootstrap_data_loader):\n",
    "    pil_images = [transforms.ToPILImage()(image) for image in images]\n",
    "    # Apply augmentation\n",
    "    augmented_images = [augment_transform(image) for image in pil_images]\n",
    "    # Convert augmented images back to tensors and save\n",
    "    augmented_images = torch.stack(augmented_images)\n",
    "    save_bootstrapped_augmented_images(augmented_images, labels, index)\n",
    "\n",
    "    # Stop after first run - Debug\n",
    "    # break\n",
    "\n",
    "\n",
    "print(\"Done augmenting the bootstrapped images! :D\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fc71fb1-70cf-42ec-87c6-673f06f74f30",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}