SpA_Former_shadow_removal_demo.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G9fueCPLFKin"
      },
      "source": [
        "GitHub<br>\n",
        "https://github.com/zhangbaijin/SpA-Former-shadow-removal<br>\n",
        "論文<br>\n",
        "https://arxiv.org/abs/2206.10910v1<br>\n",
        "<br>\n",
        "<a href=\"https://colab.research.google.com/github/kaz12tech/ai_demos/blob/master/SpA_Former_shadow_removal_demo.ipynb\" target=\"_blank\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "894c0QeEFKit"
      },
      "source": [
        "# 環境セットアップ"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iWBmluQQx-MK"
      },
      "source": [
        "## GPU確認"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qyLiFmW6x_3w"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TmL3LoMMGcNF"
      },
      "source": [
        "## Githubからソースコード取得"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tthUF7D6FKiu"
      },
      "outputs": [],
      "source": [
        "%cd /content\n",
        "\n",
        "!git clone https://github.com/zhangbaijin/SpA-Former-shadow-removal.git"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2-2Fh1sOHdnn"
      },
      "source": [
        "## ライブラリのインストール"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AP8SsHYvHfkR"
      },
      "outputs": [],
      "source": [
        "%cd /content/SpA-Former-shadow-removal\n",
        "\n",
        "!pip install gdown==4.5.1\n",
        "!pip install einops"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mgANPXUSy_MJ"
      },
      "source": [
        "## ライブラリのインポート"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hq7ZoY0oh260"
      },
      "outputs": [],
      "source": [
        "%cd /content/SpA-Former-shadow-removal\n",
        "\n",
        "import os\n",
        "import glob"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jRBoab25xtgl"
      },
      "source": [
        "# 学習済みモデルのセットアップ"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5cSMMFMwx2P5"
      },
      "outputs": [],
      "source": [
        "%cd /content/SpA-Former-shadow-removal\n",
        "!mkdir -p ./checkpoint\n",
        "\n",
        "if not os.path.exists('checkpoint/gen_model_epoch_160.pth'):\n",
        "  !gdown https://drive.google.com/uc?id=1gLOu4jkqslu_fWpQUd0N4hcNz4qCD8Xc -O checkpoint/gen_model_epoch_160.pth\n",
        "if not os.path.exists('checkpoint/dis_model_epoch_160.pth'):\n",
        "  !gdown https://drive.google.com/uc?id=1AcJSAV4oHiYneYUxCwaV0M_W1QU4MdZX -O checkpoint/dis_model_epoch_160.pth"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_BB4TY8OyFkw"
      },
      "source": [
        "# テスト画像のセットアップ"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6ELWON680lXC"
      },
      "outputs": [],
      "source": [
        "%cd /content/SpA-Former-shadow-removal\n",
        "!rm -rf ./test_imgs\n",
        "!mkdir -p ./test_imgs\n",
        "!mkdir -p ./datasets\n",
        "\n",
        "!wget -c https://cdn.pixabay.com/photo/2016/04/11/03/00/run-1321278_960_720.jpg \\\n",
        "      -O ./test_imgs/test_01.jpg\n",
        "\n",
        "!wget -c https://cdn.pixabay.com/photo/2013/10/20/22/05/shadow-198682_960_720.jpg \\\n",
        "      -O ./test_imgs/test_02.jpg\n",
        "\n",
        "!wget -c https://cdn.pixabay.com/photo/2015/09/09/20/33/travel-933171_960_720.jpg \\\n",
        "      -O ./test_imgs/test_03.jpg\n",
        "\n",
        "!wget -c https://cdn.pixabay.com/photo/2021/08/02/18/21/stairs-6517488_960_720.jpg \\\n",
        "      -O ./test_imgs/test_04.jpg\n",
        "\n",
        "if not os.path.exists('datasets/ISTD_Dataset.rar'):\n",
        "  !gdown https://drive.google.com/uc?id=1I0qw-65KBA6np8vIZzO6oeiOvcDBttAY -O datasets/ISTD_Dataset.rar"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jsOB0JxsI3o0"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "\n",
        "# https://drive.google.com/file/d/1I0qw-65KBA6np8vIZzO6oeiOvcDBttAY/view\n",
        "# 上記ファイルのショートカットを自身のGoogle Driveに追加した場合のみ以下を実行\n",
        "\n",
        "!unrar x /content/drive/MyDrive/ISTD_Dataset.rar /content/SpA-Former-shadow-removal/datasets > /dev/null"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dOmybPEWnefE"
      },
      "source": [
        "# Shadow Remove\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jqp2Ev6NnhXW"
      },
      "outputs": [],
      "source": [
        "%cd /content/SpA-Former-shadow-removal\n",
        "\n",
        "import numpy as np\n",
        "import argparse\n",
        "from cv2 import cv2\n",
        "import matplotlib.pyplot as plt\n",
        "import time\n",
        "\n",
        "import torch\n",
        "from torch.autograd import Variable\n",
        "\n",
        "from utils import gpu_manage, heatmap\n",
        "from SpA_Former import Generator\n",
        "from google.colab.patches import cv2_imshow\n",
        "\n",
        "def predict(args):\n",
        "\n",
        "    gpu_manage(args)\n",
        "    ### MODELS LOAD ###\n",
        "    print('===> Loading models')\n",
        "\n",
        "    gen = Generator(gpu_ids=args.gpu_ids)\n",
        "\n",
        "    param = torch.load(args.pretrained)\n",
        "    gen.load_state_dict(param)\n",
        "\n",
        "    if args.cuda:\n",
        "        gen = gen.cuda(0)\n",
        "\n",
        "    print ('<=== Model loaded')\n",
        "\n",
        "    print('===> Loading test image')\n",
        "    img = cv2.imread(args.test_filepath, 1).astype(np.float32)\n",
        "    img = img / 255\n",
        "    img = img.transpose(2, 0, 1)\n",
        "    img = img[None]\n",
        "    print ('<=== test image loaded')\n",
        "\n",
        "    with torch.no_grad():\n",
        "        x = torch.from_numpy(img)\n",
        "        if args.cuda:\n",
        "            x = x.cuda()\n",
        "        \n",
        "        print('===> Removing the cloud...')\n",
        "        start_time = time.time()\n",
        "        att, out = gen(x)\n",
        "        print('<=== finish! %.3fs cost.' % (time.time()-start_time))\n",
        "\n",
        "        x_ = x.cpu().numpy()[0]\n",
        "        x_rgb = x_ * 255\n",
        "        x_rgb = x_rgb.transpose(1, 2, 0).astype('uint8')\n",
        "        out_ = out.cpu().numpy()[0]\n",
        "        out_rgb = np.clip(out_[:3], 0, 1) * 255\n",
        "        out_rgb = out_rgb.transpose(1, 2, 0).astype('uint8')\n",
        "        att_ = att.cpu().numpy()[0] * 255\n",
        "        att_heatmap = heatmap(att_.astype('uint8'))[0]\n",
        "        att_heatmap = att_heatmap.transpose(1, 2, 0)\n",
        "\n",
        "        allim = np.hstack((x_rgb, out_rgb, att_heatmap))\n",
        "\n",
        "        cv2_imshow(allim)\n",
        "\n",
        "images = glob.glob('./test_imgs/*.jpg')        \n",
        "\n",
        "args = argparse.ArgumentParser()\n",
        "args.pretrained = './checkpoint/gen_model_epoch_160.pth'\n",
        "args.cuda = True\n",
        "args.gpu_ids = [0]\n",
        "args.manualSeed = 12\n",
        "\n",
        "for img in images:\n",
        "  args.test_filepath = img\n",
        "  predict(args)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dWE2u-bfFv-7"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "\n",
        "datasets = glob.glob('/content/SpA-Former-shadow-removal/datasets/ISTD_Dataset/test/test_A/*.png')\n",
        "\n",
        "tests = random.sample(datasets, 5)\n",
        "\n",
        "for img in tests:\n",
        "  args.test_filepath = img\n",
        "  predict(args)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": " SpA-Former-shadow-removal_demo.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}