create_transformed_images.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "60000\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "if not os.path.exists(\"../MNIST_dataset\"):\n",
    "    os.system(\"git clone https://github.com/DeepTrackAI/MNIST_dataset\")\n",
    "\n",
    "train_path = os.path.join(\"..\",\"MNIST_dataset\", \"mnist\", \"train\")\n",
    "train_images_files = sorted(os.listdir(train_path))\n",
    "\n",
    "print(len(train_images_files))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "from PIL import Image, ImageOps, ImageEnhance, ImageChops, ImageFilter\n",
    "\n",
    "def number_to_string(number):\n",
    "    num_to_str = {\n",
    "        0: 'zero',\n",
    "        1: 'one',\n",
    "        2: 'two',\n",
    "        3: 'three',\n",
    "        4: 'four',\n",
    "        5: 'five',\n",
    "        6: 'six',\n",
    "        7: 'seven',\n",
    "        8: 'eight',\n",
    "        9: 'nine',\n",
    "        10: 'ten',\n",
    "        30: 'thirty',\n",
    "        45: 'forty-five',\n",
    "        60: 'sixty',\n",
    "        90: 'ninety',\n",
    "        120: 'one hundred-twenty',\n",
    "        135: 'one hundred-thirty-five',\n",
    "        150: 'one hundred-fifty',\n",
    "        180: 'ninety',\n",
    "    }\n",
    "    return num_to_str[number]\n",
    "\n",
    "\n",
    "def number_representation(number):\n",
    "    if random.choice([True, False]):\n",
    "        return str(number)\n",
    "    else:\n",
    "        return number_to_string(number)\n",
    "\n",
    "\n",
    "def apply_color_transform(image, color, bkg_color):\n",
    "    # # Apply a slight blur to the mask to smooth the edges\n",
    "    # image = image.filter(ImageFilter.GaussianBlur(1))\n",
    "    # Create a new image with the background color\n",
    "    bkg_image = Image.new(\"RGB\", image.size, bkg_color)\n",
    "    # Convert the original image to RGB\n",
    "    image = image.convert(\"RGB\")\n",
    "    # Create a mask where the digit is\n",
    "    mask = image.convert(\"L\").point(lambda p: p > 128 and 255)\n",
    "    # # Apply a slight blur to the mask to smooth the edges\n",
    "    # mask = mask.filter(ImageFilter.GaussianBlur(1))\n",
    "    # Apply the color to the digit\n",
    "    colored_image = ImageChops.multiply(image, Image.new(\"RGB\", image.size, color))\n",
    "    # Combine the colored digit with the background\n",
    "    final_image = Image.composite(colored_image, bkg_image, mask)\n",
    "    return final_image\n",
    "\n",
    "def transform_image(image, label):\n",
    "    # Convert the image to a PIL Image\n",
    "    pil_image = Image.fromarray(np.uint8(image * 255))\n",
    "\n",
    "    # Define possible transformations\n",
    "    transformations = []\n",
    "    description_parts = []\n",
    "\n",
    "    # Colors\n",
    "    colors = ['black', 'red', 'blue', 'green', 'white']#, 'orange', 'purple', 'yellow', 'gray']\n",
    "    color = random.choice(colors)\n",
    "    colors.remove(color)\n",
    "    bkg_color = random.choice(colors)\n",
    "    description_parts.append(f'in {color} on a {bkg_color} background')\n",
    "\n",
    "\n",
    "    # lims = (4,5)\n",
    "    # zoom = random.choice(['in', 'out', 'none'])\n",
    "    # if zoom == 'in':\n",
    "    #     transformations.append(lambda img: img.crop((3, 3, 25, 25)))\n",
    "    #     description_parts.append(random.choice(['zoomed in', 'large']))\n",
    "    # elif zoom == 'out':\n",
    "    #     transformations.append(lambda img: ImageOps.expand(img, border=6, fill=0))\n",
    "    #     description_parts.append(random.choice(['zoomed out', 'small']))\n",
    "    #     lims=(5,10)\n",
    "\n",
    "    ### add more angles and flips\n",
    "    # Rotation\n",
    "    rotate = random.choice(['mirror','flip','rotate','nothing',])\n",
    "    #rotate = random.choice(['mirror','flip','nothing','nothing_','nothing__'])\n",
    "    if rotate=='rotate':\n",
    "        rot_angle = random.choice([-90, 90, 180])\n",
    "        # rot_angle = random.choice([-150, -135, -120, -90, -60, -45, -30, 30, 45, 60, 90, 120, 135, 150, 180])\n",
    "        transformations.append(lambda img: img.rotate(rot_angle))\n",
    "        if rot_angle>0:\n",
    "            description_parts.append(f'rotated by {number_representation(rot_angle)} degrees clockwise')\n",
    "        elif rot_angle<0:\n",
    "            description_parts.append(f'rotated by {number_representation(-rot_angle)} degrees anticlockwise')\n",
    "    elif rotate=='mirror':\n",
    "        transformations.append(lambda img: ImageOps.mirror(img))\n",
    "        description_parts.append(random.choice(['flipped horizontally', 'mirrored', 'reflected horizontally']) )\n",
    "    elif rotate=='flip':\n",
    "        transformations.append(lambda img: ImageOps.flip(img))\n",
    "        description_parts.append(random.choice(['flipped upside-down', 'reflected vertically']))\n",
    "\n",
    "    # if zoom!='in':\n",
    "    #     # Shifts\n",
    "    #     shift_lr = random.choice([True, False])\n",
    "    #     shift_ud = random.choice([True, False])\n",
    "    \n",
    "    #     if shift_lr:\n",
    "    #         delta_shift_lr = random.randint(*lims)*random.choice([-1,1])\n",
    "    #         transformations.append(lambda img: img.transform(img.size, Image.AFFINE, (1, 0, delta_shift_lr, 0, 1, 0)))\n",
    "    #         if delta_shift_lr>0:\n",
    "    #             description_parts.append(f'shifted left by {number_representation(delta_shift_lr)} pixels')\n",
    "    #         else:   description_parts.append(f'shifted right by {number_representation(-delta_shift_lr)} pixels')\n",
    "    #     else:   description_parts.append('centered horizontally')\n",
    "    #     if shift_ud:\n",
    "    #         delta_shift_ud = random.randint(*lims)*random.choice([-1,1])\n",
    "    #         transformations.append(lambda img: img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, delta_shift_ud)))\n",
    "    #         if delta_shift_ud>0:\n",
    "    #             description_parts.append(f'shifted up by {number_representation(delta_shift_ud)} pixels')\n",
    "    #         else:  description_parts.append(f'shifted down by {number_representation(-delta_shift_ud)} pixels')\n",
    "    #     else:   description_parts.append('centered vertically')\n",
    "\n",
    "\n",
    "    # Apply transformations\n",
    "    trans_image = pil_image\n",
    "    for transform in transformations:\n",
    "        trans_image = transform(trans_image)\n",
    "\n",
    "     # Apply color transformation\n",
    "    trans_image = apply_color_transform(trans_image, color, bkg_color)\n",
    "    \n",
    "    # # Resize the image to 28x28\n",
    "    trans_image = trans_image.resize((28, 28))\n",
    "    # trans_image = trans_image.resize((28, 28), Image.Resampling.LANCZOS)\n",
    "\n",
    "    # Normalize the image to range 0-255 using autocontrast\n",
    "    trans_image = ImageOps.autocontrast(trans_image, cutoff=0)\n",
    "\n",
    "    # # Convert the transformed image back to a numpy array\n",
    "    # trans_image0 = np.array(trans_image)# / 255.0\n",
    "    # print(trans_image0.min(),trans_image0.max())\n",
    "\n",
    "    # Shuffle the description parts before creating the sentence\n",
    "    # random.shuffle(description_parts)\n",
    "\n",
    "    # Create the description sentence\n",
    "    start_sent = random.choice(['Create an image of', 'Display the digit', 'Represent the handwritten digit', 'Generate an illustration of',\n",
    "                'Show the number',  'Depict the handwritten numeral', 'Produce an image of',  'Illustrate the digit', 'Craft a visual of',\n",
    "                'Render the number', 'Display the handwritten number', 'Create a depiction of', 'Form an image of'])\n",
    "    digit_repr = number_representation(label)\n",
    "    sentence = f\"{start_sent} {digit_repr} {', '.join(description_parts)}.\"\n",
    "\n",
    "    return trans_image, sentence\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "viz = False\n",
    "MNIST_images, MNIST_labels, MNIST_sentences = [], [],[]\n",
    "\n",
    "\n",
    "# Create a dictionary to store MNIST filenames and sentences\n",
    "mnist_descriptions_dict = {}\n",
    "\n",
    "for file in train_images_files:#random.sample(train_images_files,30):\n",
    "    image = plt.imread(os.path.join(train_path, file))\n",
    "\n",
    "    filename = os.path.basename(file)\n",
    "    label = int(filename[0])\n",
    "\n",
    "    trans_image, description = transform_image(image, label)\n",
    "    \n",
    "    mnist_descriptions_dict[filename] = description\n",
    "\n",
    "    # image = Image.fromarray(np.uint8(rgb_array))\n",
    "    # Save the image as a PNG file\n",
    "    trans_image.save(os.path.join('../text2image_dataset/train/images',filename))\n",
    "\n",
    "    if viz:\n",
    "        fig, axes = plt.subplots(1, 2, figsize=(8, 4))\n",
    "        fig.suptitle(description, fontsize=12)\n",
    "        # Plot original image\n",
    "        axes[0].imshow(image, cmap='gray')\n",
    "        axes[0].set_title('Original Image')\n",
    "        axes[0].axis('off')\n",
    "        \n",
    "        # Plot transformed image\n",
    "        axes[1].imshow(trans_image)\n",
    "        axes[1].set_title('Transformed Image')\n",
    "        axes[1].axis('off')\n",
    "        \n",
    "        plt.show()\n",
    "\n",
    "# Save the dictionary as a JSON file\n",
    "with open('../text2image_dataset/train/image_descriptions.json', 'w') as jsonfile:\n",
    "    json.dump(mnist_descriptions_dict, jsonfile)\n",
    "\n",
    "    # break\n",
    "    # MNIST_images.append(trans(image))\n",
    "    \n",
    "    # filename = os.path.basename(file)\n",
    "    # label = int(filename[0])\n",
    "    # MNIST_labels.append(label)\n",
    "\n",
    "    # sentence = mnist_sentences.get(filename, \"\")\n",
    "    # MNIST_sentences.append(sentence)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "deeplay_env",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}