document_colorizer.ipynb

{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### **Document Colorizer**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### **Import necessary libraries**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "import cv2\n",
        "import random"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### **Parse a set of documents as background images to paste checkboxes onto**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "##### **RVL-CDIP**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The dataset contains 16 types of documents:\n",
        "0 letter\n",
        "1 form\n",
        "2 email\n",
        "3 handwritten\n",
        "4 advertisement\n",
        "5 scientific report\n",
        "6 scientific publication\n",
        "7 specification\n",
        "8 file folder\n",
        "9 news article\n",
        "10 budget\n",
        "11 invoice\n",
        "12 presentation\n",
        "13 questionnaire\n",
        "14 resume\n",
        "15 memo\n",
        "\n",
        "For the purpose of reducing as much as noise as possible for the training set, we might look at category 0, 2, 3, 6, 9, 12"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [],
      "source": [
        "DOC_CATEGORIES = {\n",
        "    0: \"letter\",\n",
        "    1: \"form\",\n",
        "    2: \"email\",\n",
        "    3: \"handwritten\",\n",
        "    4: \"advertisement\",\n",
        "    5: \"scientific_report\",\n",
        "    6: \"scientific_publication\",\n",
        "    7: \"specification\",\n",
        "    8: \"file_folder\",\n",
        "    9: \"news_article\",\n",
        "    10: \"budget\",\n",
        "    11: \"invoice\",\n",
        "    12: \"presentation\",\n",
        "    13: \"questionnaire\",\n",
        "    14: \"resume\",\n",
        "    15: \"memo\"\n",
        "}\n",
        "categories = [0, 2, 3, 6, 9]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [],
      "source": [
        "train_dir = \"datasets/rvl-cdip/images\"\n",
        "train_txt = \"datasets/rvl-cdip/labels/train.txt\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "count = 0\n",
        "\n",
        "with open(train_txt, 'r') as train_folder:\n",
        "    for line in train_folder:\n",
        "        line_array = line.split(\" \")\n",
        "        cat_num = int(line_array[1])\n",
        "        source_path = os.path.join(train_dir, line_array[0])\n",
        "        # Obtain about 10000 images in the abovementioned categories which will serve as source documents to paste checkboxes onto\n",
        "        if cat_num in categories and count <= 10000:\n",
        "            new_path = f\"datasets/train_val/rvl_boxes/train/images/{str(count) + '_' + DOC_CATEGORIES[cat_num] + '.jpg'}\"\n",
        "            shutil.copy(source_path, new_path) # Copy the image\n",
        "            count += 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "##### **Colorize**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Since the existing documents are entirely grayscale images, which differ greatly from real-life documents which can be in various settings, we will apply a set of colorization processes on a portion of the whole data. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Define helper functions and constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [],
      "source": [
        "COLOR_COMBINATIONS = [ # Text, Background\n",
        "    (\"#2F3C7E\", \"#FBEAEB\"),\n",
        "    (\"#317773\", \"#E2D1F9\"),\n",
        "    (\"#990011\", \"#FCF6F5\"),\n",
        "    (\"#FFFFFF\", \"#8AAAE5\"), \n",
        "    (\"#EE4E34\", \"#FCEDDA\"),\n",
        "    (\"#3A6B35\", \"#CBD18F\"),\n",
        "    (\"#00008B\", \"#ADD8E6\"),\n",
        "    (\"#101820\", \"#FBF8BE\"),\n",
        "    (\"#000000\", \"#E7E8D1\"),\n",
        "    (\"#101820\", \"#FCF6F5\"),\n",
        "    (\"#2C5F2D\", \"#FCF6F5\"),\n",
        "    (\"#000000\", \"#FCEDDA\"),\n",
        "    (\"#A41681\", \"#B2FAE6\"),\n",
        "    (\"#B2FAE6\", \"#A41681\"),\n",
        "    (\"#081B26\", \"#1E90FF\"),\n",
        "    (\"#008080\", \"#FDFEFD\")\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [],
      "source": [
        "def convertHexToRGB(value):\n",
        "    value = value.lstrip('#')\n",
        "    lv = len(value)\n",
        "    return [int(value[i:i + lv // 3], 16) for i in range(0, lv, lv // 3)]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [],
      "source": [
        "document_image_paths = next(os.walk(\"datasets/train_val/rvl_boxes/train/images/\"))[2]\n",
        "# Get all unedited images\n",
        "complete_document_image_paths = [os.path.join(\"datasets/train_val/rvl_boxes/train/images/\", x) for x in document_image_paths if x.find(\".DS_Store\") < 0]\n",
        "# Get a random group to apply colorization\n",
        "transform_group = random.choices(complete_document_image_paths, k=4000)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Randomly change background, foreground color of images"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 65,
      "metadata": {},
      "outputs": [],
      "source": [
        "for image_path in transform_group:\n",
        "    color_choice = random.choices(COLOR_COMBINATIONS)[0]\n",
        "    img = cv2.imread(image_path)\n",
        "    if img is None:\n",
        "        continue\n",
        "    # Get all black pixels\n",
        "    darker_pixels = np.where(\n",
        "        (img[:, :, 0] <= 128) & \n",
        "        (img[:, :, 1] <= 128) & \n",
        "        (img[:, :, 2] <= 128)\n",
        "    )\n",
        "\n",
        "    lighter_pixels = np.where(\n",
        "        (img[:, :, 0] > 128) & \n",
        "        (img[:, :, 1] > 128) & \n",
        "        (img[:, :, 2] > 128)\n",
        "    )\n",
        "\n",
        "    # set those pixels to text color in the color choice tuple\n",
        "    img[darker_pixels] = convertHexToRGB(color_choice[0])\n",
        "    img[lighter_pixels] = convertHexToRGB(color_choice[1])\n",
        "    cv2.imwrite(image_path, img) # Overwrite the original black-and-white image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 66,
      "metadata": {},
      "outputs": [],
      "source": [
        "document_image_paths = next(os.walk(\"datasets/train_val/rvl_boxes/train/images/\"))[2]\n",
        "# Get all unedited images\n",
        "complete_document_image_paths = [os.path.join(\"datasets/train_val/rvl_boxes/train/images/\", x) for x in document_image_paths if x.find(\".DS_Store\") < 0]"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "6fr7wOFs2VCi"
      ],
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.11"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}