Color-Compression.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r3zz6K82pkMJ"
      },
      "source": [
        "# **PROJECT 1: COLOR COMPRESSION** <br>\n",
        "**Subject:** Applied Mathematics and Statistics <br>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6PJeXRoQxr81"
      },
      "source": [
        "## **Table of contents**\n",
        "* [Import librabries and designate a path](#c1)\n",
        "* [Image preparation and processing](#c2)\n",
        "* [KMeans function](#c3)\n",
        "* [Test program for k = {3, 5, 7}](#c4)\n",
        "* [Main function](#c5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iWv4zLSmq_ST"
      },
      "source": [
        "### **Import librabries and designate a path**<a class=\"achor\" id=\"c1\"></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "H8JRZ4Y7quSR",
        "outputId": "1ab8ac9a-f6c9-4d3b-f9ae-9023c89f9c21"
      },
      "outputs": [],
      "source": [
        "# Import librabries\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from PIL import Image\n",
        "init_centroidsT = ['random', 'in_pixels'] # For centroids choices"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AEjfrQgPqWr8"
      },
      "source": [
        "### **Image preparation and processing**<a class=\"achor\" id=\"c2\"></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "Aoc9Ns9VqZ2-"
      },
      "outputs": [],
      "source": [
        "# Flatten image\n",
        "def flatten_img(img):\n",
        "  height, width = img.size\n",
        "  return np.reshape(img, (width * height, 3)).astype(int)\n",
        "\n",
        "# Compress image base on the provided centroids and label\n",
        "def compress_img(centroids, labels, img):\n",
        "  height, width = img.size\n",
        "  compressed = centroids[labels.astype(int)]\n",
        "  compressed = compressed.reshape((width, height, 3)).astype(int)\n",
        "  return compressed\n",
        "\n",
        "# Show image function\n",
        "def show_img(img, img_name, img_centroids, num_cluster):\n",
        "  plt.title(f'{img_name} {img_centroids} with k cluster = {num_cluster}')\n",
        "  plt.imshow(img)\n",
        "  plt.axis('off') # Remove the axis\n",
        "  plt.show"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5PO48BizSfRh"
      },
      "source": [
        "### **KMeans function**<a class=\"achor\" id=\"c3\"></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "9kLaSNuBUK2k"
      },
      "outputs": [],
      "source": [
        "def kmeans(img_1d, k_clusters, max_iter, init_centroids='random'):\n",
        "  # Choose k random centroids\n",
        "  length, dim = img_1d.shape\n",
        "  if init_centroids == 'random':\n",
        "    centroids = np.random.choice(256, size = (k_clusters, dim), replace=False)\n",
        "  elif init_centroids == 'in_pixels':\n",
        "    centroids = img_1d[np.random.choice(length, size = k_clusters, replace=False)]\n",
        "\n",
        "  labels = np.zeros(shape=(length))\n",
        "\n",
        "  while max_iter:\n",
        "    for i in range(length):\n",
        "      min_dist = float('inf')\n",
        "\n",
        "      # Distance between each pixel and centroids\n",
        "      for j in range(k_clusters):\n",
        "        mean_dist = np.sqrt(np.sum((img_1d[i] - centroids[j])**2, axis=0))\n",
        "        if mean_dist < min_dist:\n",
        "          min_dist = mean_dist\n",
        "          labels[i] = j\n",
        "\n",
        "    # Update centroids\n",
        "    for i in range(k_clusters):\n",
        "      pixels = img_1d[labels == i]\n",
        "      if len(pixels):\n",
        "        centroids[i] = np.mean(pixels, axis = 0)\n",
        "    max_iter -= 1\n",
        "\n",
        "  return centroids, labels"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qKxscS0KK9H7"
      },
      "source": [
        "### **Test program for k = {3, 5, 7}**<a class=\"achor\" id=\"c4\"></a>\n",
        "Test function to test for kmean in two centroids cases <br>\n",
        "Let max_iteration = 10"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bu41nHZuLYAu"
      },
      "source": [
        "#### **centroids =  random**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "SkmGstpkK83G"
      },
      "outputs": [],
      "source": [
        "def random_test(img):\n",
        "  output_img = []\n",
        "  output_img.append(img)\n",
        "  for k_cluster in [3, 5, 7]:\n",
        "    flat_img = flatten_img(img)\n",
        "    centroids, labels = kmeans(flat_img, k_cluster, 10, 'random')\n",
        "    final_img = compress_img(centroids, labels, init_img)\n",
        "\n",
        "    final_img = Image.fromarray(final_img.astype('uint8'), 'RGB')\n",
        "    output_img.append(final_img)\n",
        "\n",
        "  # Creating subplot\n",
        "  plot_iter = 1\n",
        "  fig, axis = plt.subplots(1, 3, figsize=(12, 8))\n",
        "\n",
        "  for i,k in [(0,3), (1,5), (2,7)]:\n",
        "    axis[i].set_title(f'Random with k cluster = {k}')\n",
        "    axis[i].imshow(output_img[plot_iter])\n",
        "    axis[i].axis('off')\n",
        "    plot_iter += 1\n",
        "  plt.tight_layout();"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zJdcWsOASXNE"
      },
      "source": [
        "#### **centroids = in_pixels**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "j5bii2KQSbjO"
      },
      "outputs": [],
      "source": [
        "def in_pixels_test(img):\n",
        "  output_img = []\n",
        "  output_img.append(img)\n",
        "  for k_cluster in [3, 5, 7]:\n",
        "    flat_img = flatten_img(img)\n",
        "    centroids, labels = kmeans(flat_img, k_cluster, 10, 'in_pixels')\n",
        "    final_img = compress_img(centroids, labels, init_img)\n",
        "\n",
        "    final_img = Image.fromarray(final_img.astype('uint8'), 'RGB')\n",
        "    output_img.append(final_img)\n",
        "\n",
        "  # Creating subplot\n",
        "  plot_iter = 1\n",
        "  fig, axis = plt.subplots(1, 3, figsize=(12, 8))\n",
        "\n",
        "  for i,k in [(0,3), (1,5), (2,7)]:\n",
        "    axis[i].set_title(f'In_pixels with k cluster = {k}')\n",
        "    axis[i].imshow(output_img[plot_iter])\n",
        "    axis[i].axis('off')\n",
        "    plot_iter += 1\n",
        "  plt.tight_layout();"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sCGQO-ux1FwJ"
      },
      "source": [
        "### **Main function** <a class=\"achor\" id=\"c5\"></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "LYY8mf111JQq"
      },
      "outputs": [],
      "source": [
        "def main():\n",
        "  # Input parameters\n",
        "  img_path = input(\"Enter image's path (use C:/ instead of C:\\): \")\n",
        "  max_iter = int(input(\"Enter max iteration: \"))\n",
        "  k_cluster = int(input(\"Number of k cluster: \"))\n",
        "  centroidsT = int(input(\"Initial centroid: \\n0) random \\n1) in_pixels \\nYour choice: \"))\n",
        "\n",
        "  # File name add-on\n",
        "  if centroidsT == 0:\n",
        "    centroids_type = 'random'\n",
        "  elif centroidsT == 1:\n",
        "    centroids_type = 'in_pixels'\n",
        "  else:\n",
        "    print('Invalid centroid initialization method')\n",
        "    return\n",
        "\n",
        "  # Open the image\n",
        "  init_img = Image.open(img_path).convert('RGB')\n",
        "\n",
        "  # Flatten image for processing\n",
        "  flat_img = flatten_img(init_img)\n",
        "\n",
        "  # Kmeans processing\n",
        "  centroids, labels = kmeans(flat_img, k_cluster, max_iter, init_centroidsT[centroidsT])\n",
        "  final_img = compress_img(centroids, labels, init_img)\n",
        "  final_img = Image.fromarray(final_img.astype('uint8'), 'RGB')\n",
        "\n",
        "  # Process image name\n",
        "  ## Remove the front directory and the .<extension>\n",
        "  output_img = (img_path.split('/')[-1]).split('.')[0]\n",
        "\n",
        "  # Get the directory path for saving purposes\n",
        "  directory_path = img_path.replace('/' + img_path.split('/')[-1], '')\n",
        "\n",
        "  # Output file type choice\n",
        "  save_choice = int(input('Enter output file type: \\n1) PNG \\n2) JPG \\n3) PDF \\nYour choice: '))\n",
        "\n",
        "  # Output file check and save image\n",
        "  if save_choice == 1:\n",
        "    final_img.save(f\"{directory_path}/{output_img}_{centroids_type}_{k_cluster}_compressed.png\",\"PNG\")\n",
        "    print(f\"Image saved at {directory_path}/{output_img}_{centroids_type}_{k_cluster}_compressed.png\")\n",
        "  elif save_choice == 2:\n",
        "    final_img.save(f\"{directory_path}/{output_img}_{centroids_type}_{k_cluster}_compressed.jpg\",\"JPEG\")\n",
        "    print(f\"Image saved at {directory_path}/{output_img}_{centroids_type}_{k_cluster}_compressed.jpg\")\n",
        "  elif save_choice == 3:\n",
        "    final_img.save(f\"{directory_path}/{output_img}_{centroids_type}_{k_cluster}_compressed.pdf\",\"PDF\")\n",
        "    print(f\"Image saved at {directory_path}/{output_img}_{centroids_type}_{k_cluster}_compressed.pdf\")\n",
        "  else:\n",
        "    print(\"Invalid file save type!\")\n",
        "    return\n",
        "\n",
        "  # Show image\n",
        "  show_img(final_img, output_img, init_centroidsT[centroidsT], k_cluster)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 615
        },
        "id": "S-yZqh-4stZj",
        "outputId": "879671a2-d7a1-4fa3-acde-946eac6407f4"
      },
      "outputs": [],
      "source": [
        "main()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "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
}