coloring.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true,
      "gpuType": "T4",
      "mount_file_id": "11_YAyYTaSvFg3diBXcJR2Ww6_T2AHyvB",
      "authorship_tag": "ABX9TyOCu39dNr1Wi3Hl24WRaUKB",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/lifeisbeautifu1/graduate-work/blob/main/coloring.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Neural Network for coloring grey-scale images"
      ],
      "metadata": {
        "id": "qSsiu6yubKvQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# check for gpu\n",
        "!nvidia-smi -L"
      ],
      "metadata": {
        "id": "Lf2wkcMUZsjC",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "cd216bf5-54ef-46f7-d6c4-fab87faed952"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "GPU 0: Tesla T4 (UUID: GPU-25920366-4fb9-4418-064b-a9f3f337df34)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Alpha version\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "FhjSzB2cMik1"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Imports"
      ],
      "metadata": {
        "id": "maptLKifPlqn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from keras.layers import Conv2D, UpSampling2D, InputLayer, Conv2DTranspose\n",
        "from keras.layers import Activation, Dense, Dropout, Flatten\n",
        "from keras.layers import BatchNormalization\n",
        "from keras.models import Sequential\n",
        "from keras.preprocessing.image import ImageDataGenerator\n",
        "from keras.utils import array_to_img, img_to_array, load_img\n",
        "from skimage.color import rgb2lab, lab2rgb, rgb2gray, xyz2lab\n",
        "from skimage.io import imsave\n",
        "import numpy as np\n",
        "import os\n",
        "import random\n",
        "import tensorflow as tf"
      ],
      "metadata": {
        "id": "-mIZTdCOkBSF"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "PATH_TO_DRIVE = \"drive/MyDrive/Colab Notebooks/\"\n",
        "IMAGE_SHAPE = 228\n",
        "\n",
        "image = img_to_array(load_img(PATH_TO_DRIVE + 'woman.jpg'))\n",
        "image = np.array(image, dtype=float)"
      ],
      "metadata": {
        "id": "aEupwPD1kBUs"
      },
      "execution_count": 47,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "X = rgb2lab(1.0/255*image)[:,:,0]\n",
        "Y = rgb2lab(1.0/255*image)[:,:,1:]\n",
        "Y /= 128\n",
        "X = X.reshape(1,  400, 400, 1)\n",
        "Y = Y.reshape(1,  400, 400, 2)"
      ],
      "metadata": {
        "id": "aLq_L6tmkBX1"
      },
      "execution_count": 48,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = Sequential()\n",
        "model.add(InputLayer(input_shape=(None, None, 1)))\n",
        "model.add(Conv2D(8, (3, 3), activation='relu', padding='same', strides=2))\n",
        "model.add(Conv2D(8, (3, 3), activation='relu', padding='same'))\n",
        "model.add(Conv2D(16, (3, 3), activation='relu', padding='same'))\n",
        "model.add(Conv2D(16, (3, 3), activation='relu', padding='same', strides=2))\n",
        "model.add(Conv2D(32, (3, 3), activation='relu', padding='same'))\n",
        "model.add(Conv2D(32, (3, 3), activation='relu', padding='same', strides=2))\n",
        "model.add(UpSampling2D((2, 2)))\n",
        "model.add(Conv2D(32, (3, 3), activation='relu', padding='same'))\n",
        "model.add(UpSampling2D((2, 2)))\n",
        "model.add(Conv2D(16, (3, 3), activation='relu', padding='same'))\n",
        "model.add(UpSampling2D((2, 2)))\n",
        "model.add(Conv2D(2, (3, 3), activation='tanh', padding='same'))"
      ],
      "metadata": {
        "id": "MzNI_wUNkBbN"
      },
      "execution_count": 49,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model.compile(optimizer='rmsprop',loss='mse')"
      ],
      "metadata": {
        "id": "WCHoiQUfkBeE"
      },
      "execution_count": 50,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model.fit(x=X,\n",
        "    y=Y,\n",
        "    batch_size=1,\n",
        "    verbose=0,\n",
        "    epochs=1000)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IkE9OU8BkBhe",
        "outputId": "74b8e7ff-bc5a-46ca-e333-8f01136c8418"
      },
      "execution_count": 51,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<keras.callbacks.History at 0x7938b94ae140>"
            ]
          },
          "metadata": {},
          "execution_count": 51
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(model.evaluate(X, Y, batch_size=1))\n",
        "output = model.predict(X)\n",
        "output *= 128\n",
        "# Output colorizations\n",
        "cur = np.zeros((400, 400, 3))\n",
        "cur[:,:,0] = X[0][:,:,0]\n",
        "cur[:,:,1:] = output[0]\n",
        "\n",
        "# cur.shape\n",
        "image_rgb = lab2rgb(cur)\n",
        "image_rgb *= 255\n",
        "imsave(PATH_TO_DRIVE + \"img_result.png\", image_rgb.astype(np.uint8))\n",
        "imsave(PATH_TO_DRIVE + \"img_gray_version.png\", rgb2gray(image_rgb).astype(np.uint8))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Kb5tnrqeWJSo",
        "outputId": "57a7b6a9-2961-4739-aaa3-7b8f3cbb0370"
      },
      "execution_count": 52,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1/1 [==============================] - 0s 140ms/step - loss: 9.9724e-04\n",
            "0.000997238210402429\n",
            "1/1 [==============================] - 0s 103ms/step\n"
          ]
        }
      ]
    }
  ]
}