02-dataset-preparation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "02-dataset_preparation.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vzNwCqRMxvCo"
      },
      "source": [
        "# Dataset Preparation\n",
        "Dataset processing and augmentation tools\n",
        "\n",
        "Created by [Artem Konevskikh](https://github.com/artem-konevskikh)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "81nfc8JN2EBK",
        "cellView": "form"
      },
      "source": [
        "#@title Import libraries\n",
        "#@markdown Import libraries required for image manipulation\n",
        "!apt-get update > /dev/null 2>&1\n",
        "!apt-get install imagemagick > /dev/null 2>&1\n",
        "\n",
        "import glob\n",
        "import os\n",
        "from math import floor, ceil\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "from tqdm.notebook import tqdm\n",
        "from fastai.vision.utils import verify_images\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "42K-qGO7q7I7"
      },
      "source": [
        "#@title Define functions\n",
        "\n",
        "#@markdown Functions for image manipulations\n",
        "\n",
        "def crop_center(img, ratio):\n",
        "  img = img.copy()\n",
        "  width, height = img.size\n",
        "  crop_size= ceil(ratio * min(img.size))\n",
        " \n",
        "  left = (width - crop_size)/2\n",
        "  top = (height - crop_size)/2\n",
        "  right = (width + crop_size)/2\n",
        "  bottom = (height + crop_size)/2\n",
        "\n",
        "  return img.crop((left, top, right, bottom))\n",
        "\n",
        "def save_crop(img, cnt, path, augment=False):\n",
        "    img = img.copy()\n",
        "\n",
        "    img.save(f'{path}{cnt:08d}.png')\n",
        "    cnt += 1\n",
        "    if augment:\n",
        "        img.transpose(Image.FLIP_LEFT_RIGHT).save(f'{path}/{cnt:08d}.png')\n",
        "        cnt += 1\n",
        "        img.transpose(Image.FLIP_TOP_BOTTOM).save(f'{path}/{cnt:08d}.png')\n",
        "        cnt += 1\n",
        "        img.transpose(Image.ROTATE_180).save(f'{path}/{cnt:08d}.png')\n",
        "        cnt += 1\n",
        "\n",
        "    return cnt\n",
        "\n",
        "def thumbnail(img, size=256):\n",
        "    \"\"\"\n",
        "    resize image so smallest side will be equal to size\n",
        "    \"\"\"\n",
        "    \n",
        "    img = img.copy()\n",
        "\n",
        "    if img.mode not in ('L', 'RGB'):\n",
        "        img = img.convert('RGB')\n",
        "\n",
        "    width, height = img.size\n",
        "\n",
        "    if width == height:\n",
        "        img.thumbnail((size, size), Image.ANTIALIAS)\n",
        "\n",
        "    elif height > width:\n",
        "        ratio = float(height) / float(width)\n",
        "        newheight = ratio * size\n",
        "        img = img.resize((size, int(floor(newheight))), Image.ANTIALIAS)\n",
        "\n",
        "    elif width > height:\n",
        "        ratio = float(width) / float(height)\n",
        "        newwidth = ratio * size\n",
        "        img = img.resize((int(floor(newwidth)), size), Image.ANTIALIAS)\n",
        "\n",
        "\n",
        "    return img\n",
        "\n",
        "def get_crop_bboxes(w, h):\n",
        "    \"\"\"\n",
        "    calculate bounding boxes based on width and height\n",
        "    \"\"\"\n",
        "    n_step = ceil(float(w) / float(h))\n",
        "    shift = (w - h) // n_step\n",
        "    shifting_array = []\n",
        "    for step in range(0, n_step):\n",
        "        if w > h:\n",
        "            shifting = (shift * step, 0, shift * step + h, h)\n",
        "        else:\n",
        "            shifting = (0, shift * step, w, shift * step + w)\n",
        "        shifting_array.append(shifting)\n",
        "    if w > h:\n",
        "        shifting = (w - h, 0, w, h)\n",
        "    else:\n",
        "        shifting = (0, h - w, w, h)\n",
        "    shifting_array.append(shifting)\n",
        "    return shifting_array"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "MtszQjS01CYI",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "7c479bbd-95c9-4cce-ae87-b2fe12b4c7e7"
      },
      "source": [
        "#@title Mount Google Drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Cropping"
      ],
      "metadata": {
        "id": "pIed1r-onumK"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AEkxxHEkxuFy",
        "cellView": "form"
      },
      "source": [
        "#@title Crop center\n",
        "#@markdown This cell will crop the center of each image in input directory, resize it and save result to output directory\n",
        "\n",
        "#@markdown Input directory\n",
        "input_dir = '/content/drive/MyDrive/uibk/Experiment-cropped' #@param {type:'string'}\n",
        "#@markdown Output directory. Will be created if not exist\n",
        "out_dir = '/content/drive/MyDrive/ants-exp' #@param {type:'string'}\n",
        "#@markdown Resize to\n",
        "resize = \"512\" #@param [256, 512, 1024] {allow-input: true}\n",
        "resize = int(resize)\n",
        "#@markdown Crop size ratio. The size of cropped square will be equal to `ratio * min side size`\n",
        "ratio = 1 #@param {type:\"slider\", min:0.1, max:1, step:0.05}\n",
        "#@markdown Augment dataset with rotation\n",
        "augment=False #@param {type:'boolean'}\n",
        "\n",
        "if out_dir[-1] != \"/\":\n",
        "  out_dir += \"/\"\n",
        "\n",
        "if not os.path.exists(out_dir):\n",
        "    os.makedirs(out_dir)\n",
        "\n",
        "img_ext = ['tif', 'tiff', 'TIF', 'TIFF', 'png', 'PNG', 'jpg', 'JPG', 'jpeg', 'JPEG', 'bmp', 'BMP']\n",
        "images = []\n",
        "for ext in img_ext:\n",
        "  images += glob.glob(f'{input_dir}/*.{ext}')\n",
        "\n",
        "cnt=0\n",
        "print(f'Cropping {len(images)} images\\n')\n",
        "for image in tqdm(images):\n",
        "  img = Image.open(image)\n",
        "  cropped = crop_center(img, ratio)\n",
        "  resized = cropped.resize((resize, resize), Image.ANTIALIAS)\n",
        "  cnt=save_crop(resized, cnt, out_dir, augment)\n",
        "print(f'Number of images saved: {cnt}')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "djBeB_42yF0D",
        "cellView": "form"
      },
      "source": [
        "#@title Slice images\n",
        "\n",
        "#@markdown This cell will resize each image in input directory, slice it into several square images and save result to output directory\n",
        "\n",
        "#@markdown Input directory\n",
        "input_dir = \"/content/drive/MyDrive/scraped-images/Lichen\" #@param {type:'string'}\n",
        "#@markdown Output directory. Will be created if not exist\n",
        "out_dir = '/content/drive/MyDrive/lichen-sliced' #@param {type:'string'}\n",
        "#@markdown Resize to\n",
        "resize = \"1024\" #@param [256, 512, 1024] {allow-input: true}\n",
        "resize = int(resize)\n",
        "#@markdown Augment dataset with rotation\n",
        "augment=True #@param {type:'boolean'}\n",
        "\n",
        "if out_dir[-1] != \"/\":\n",
        "  out_dir += \"/\"\n",
        "\n",
        "if not os.path.exists(out_dir):\n",
        "    os.makedirs(out_dir)\n",
        "\n",
        "img_ext = ['tif', 'tiff', 'TIF', 'TIFF', 'png', 'PNG', 'jpg', 'JPG', 'jpeg', 'JPEG', 'bmp', 'BMP']\n",
        "images = []\n",
        "for ext in img_ext:\n",
        "  images += glob.glob(f'{input_dir}/*.{ext}')\n",
        "\n",
        "cnt=0\n",
        "print(f'Cropping {len(images)} images\\n')\n",
        "for image in tqdm(images):\n",
        "    img = Image.open(image)\n",
        "    img = thumbnail(img, size=resize)\n",
        "    w, h = img.size\n",
        "    bboxes = get_crop_bboxes(w, h)\n",
        "    for bbox in bboxes:\n",
        "        img_cropped = img.crop(bbox)\n",
        "        cnt = save_crop(img_cropped, cnt, out_dir, augment)\n",
        "print(f'Number of images saved: {cnt}')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Enhancing"
      ],
      "metadata": {
        "id": "ZoAebhAOnzXM"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WdYlMb15rbRl",
        "cellView": "form"
      },
      "source": [
        "#@title Auto-enhance images\n",
        "\n",
        "\n",
        "\n",
        "#@markdown Input directory\n",
        "input_dir = '/content/drive/MyDrive/lichen-sliced' #@param {type:'string'}\n",
        "#@markdown Output directory. Will be created if not exist\n",
        "out_dir = '/content/drive/MyDrive/lichen-sliced-level-gamma' #@param {type:'string'}\n",
        "#@markdown Enhancing options\n",
        "auto_level = False #@param {type:\"boolean\"}\n",
        "auto_gamma = False #@param {type:\"boolean\"}\n",
        "invert = False #@param {type:\"boolean\"}\n",
        "\n",
        "options = ''\n",
        "options+=' -auto-level' if auto_level else ''\n",
        "options+=' -auto-gamma' if auto_gamma else ''\n",
        "options+=' -negate' if invert else ''\n",
        "\n",
        "if out_dir[-1] != \"/\":\n",
        "  out_dir += \"/\"\n",
        "\n",
        "if not os.path.exists(out_dir):\n",
        "    os.makedirs(out_dir)\n",
        "\n",
        "img_ext = ['tif', 'tiff', 'TIF', 'TIFF', 'png', 'PNG', 'jpg', 'JPG', 'jpeg', 'JPEG', 'bmp', 'BMP']\n",
        "images = []\n",
        "for ext in img_ext:\n",
        "  images += glob.glob(f'{input_dir}/*.{ext}')\n",
        "\n",
        "print(f'Enhancing {len(images)} images\\n')\n",
        "for image in tqdm(images):\n",
        "  prefix = options.replace('-auto','').replace(' ', '')[1:]\n",
        "  outfile = f\"{out_dir}{prefix}-{image.split('/')[-1]}\"\n",
        "  !convert {image} {options} {outfile}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Verifying"
      ],
      "metadata": {
        "id": "pb-cMxkCn39N"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kJmepgICy1cG",
        "cellView": "form"
      },
      "source": [
        "#@title Verify images\n",
        "#@markdown This cell will verify images in input directory. Images will be resized to defined size, converted to RGB and saved to output directory\n",
        "\n",
        "#@markdown Input directory\n",
        "input_dir = '/content/drive/MyDrive/lichen-sliced' #@param {type:'string'}\n",
        "#@markdown Output directory. Will be created if not exist\n",
        "out_dir = '/content/drive/MyDrive/lichen-dataset-256' #@param {type:'string'}\n",
        "#@markdown Resize to\n",
        "resize = \"512\" #@param [256, 512, 1024] {allow-input: true}\n",
        "resize = int(resize)\n",
        "\n",
        "if out_dir[-1] != \"/\":\n",
        "  out_dir += \"/\"\n",
        "\n",
        "if not os.path.exists(out_dir):\n",
        "    os.makedirs(out_dir)\n",
        "\n",
        "img_ext = ['tif', 'tiff', 'TIF', 'TIFF', 'png', 'PNG', 'jpg', 'JPG', 'jpeg', 'JPEG', 'bmp', 'BMP']\n",
        "images = []\n",
        "for ext in img_ext:\n",
        "  images += glob.glob(f'{input_dir}/*.{ext}')\n",
        "\n",
        "cnt=0\n",
        "print(f'Verifying {len(images)} images\\n')\n",
        "for image in tqdm(images):\n",
        "    img = Image.open(image)\n",
        "    if img.mode == 'CMYK':\n",
        "      img = img.convert('RGB')\n",
        "    img = img.resize((resize, resize), Image.ANTIALIAS)\n",
        "    imgarr = np.array(img)\n",
        "    img_channels = 1 if len(imgarr.shape) == 2 else imgarr.shape[2]\n",
        "    if img_channels == 1:\n",
        "      # print(imgarr.shape)\n",
        "      R = np.stack((imgarr, imgarr, imgarr), axis=2)\n",
        "      img = Image.fromarray(R, 'RGB')\n",
        "      imgarr = np.array(img)\n",
        "      # print(imgarr.shape)\n",
        "    if img.mode in ('RGBA', 'LA') or (img.mode == 'P' and 'transparency' in img.info):\n",
        "      bg = Image.new('RGB', img.size, (255, 255, 255))\n",
        "      bg.paste(img, (0, 0), img)\n",
        "      img = bg\n",
        "    img.save(f'{out_dir}{cnt:08d}.png')\n",
        "    cnt+=1\n",
        "    "
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}