Colab-traiNNer-inference.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Colab-traiNNer-inference.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "private_outputs": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Colab-traiNNer-inference\n",
        "\n",
        "victorca25's BasicSR fork: [victorca25/traiNNer](https://github.com/victorca25/traiNNer)\n",
        "\n",
        "My fork: [styler00dollar/Colab-traiNNer](https://github.com/styler00dollar/Colab-traiNNer)\n",
        "\n",
        "Currently only a minimalistic example for inpainting. "
      ],
      "metadata": {
        "id": "xphod4VBituY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!nvidia-smi"
      ],
      "metadata": {
        "id": "NJ5tY-v4jYnl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AX-1OUHZiZUy",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Install dependencies (only the most basic ones, install others if needed)\n",
        "%cd /content/\n",
        "# create empty folders\n",
        "!mkdir /content/input\n",
        "!mkdir /content/output\n",
        "\n",
        "!git clone https://github.com/styler00dollar/Colab-traiNNer\n",
        "!pip install omegaconf piq wget adamp efficientnet_pytorch tensorboardX vit-pytorch swin-transformer-pytorch madgrad timm pillow-avif-plugin kornia\n",
        "!pip uninstall pytorch_lightning -y\n",
        "!pip install git+https://github.com/styler00dollar/pytorch-lightning.git@remotes/origin/release/1.4.x\n",
        "\n",
        "!pip install tfrecord tensorboardX -U\n",
        "!pip install crc32c==2.0 # downgrade to surpress tensorboardX warnings"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title inpaint inference (cuda)\n",
        "#@markdown Put files into input folder `input.png / input_mask.png`\n",
        "#@markdown and manually specify model and its parameters\n",
        "import glob\n",
        "import cv2\n",
        "import os\n",
        "from tqdm import tqdm\n",
        "import torch\n",
        "from torchvision.utils import save_image\n",
        "\n",
        "rootdir = \"/content/input\" #@param\n",
        "outdir = \"/content/output/\" #@param\n",
        "fp16 = True #@param {type:\"boolean\"}\n",
        "model_path = \"/content/drive/MyDrive/model.pth\" #@param\n",
        "files = glob.glob(rootdir + '/**/*.png', recursive=True)\n",
        "\n",
        "# select and load a model\n",
        "%cd /content/Colab-traiNNer/code/arch\n",
        "from lama_arch import FFCResNetGenerator\n",
        "netG = FFCResNetGenerator(4, 3)\n",
        "netG.load_state_dict(torch.load(model_path))\n",
        "netG.cuda().eval()\n",
        "if fp16:\n",
        "  netG.half()\n",
        "%cd /content/\n",
        "\n",
        "# create list with only image files\n",
        "inference_files = []\n",
        "for i in files:\n",
        "  if \"_mask\" not in i:\n",
        "    inference_files.append(i)\n",
        "\n",
        "for f in tqdm(inference_files):\n",
        "  image = cv2.imread(f)\n",
        "  image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
        "  mask = cv2.imread(os.path.join(rootdir, os.path.splitext(os.path.basename(f))[0] + \"_mask.png\"), cv2.IMREAD_GRAYSCALE)\n",
        "  # forcing loaded mask to either be 0 or 255\n",
        "  _,mask = cv2.threshold(mask,127,255,cv2.THRESH_BINARY)\n",
        "  image[mask==255] = [0,0,0]\n",
        "  image = torch.from_numpy(image).cuda().unsqueeze(0).permute(0,3,1,2)/255\n",
        "  mask = torch.from_numpy(mask).cuda().unsqueeze(0).unsqueeze(0)/255\n",
        "  if fp16:\n",
        "    image = image.half()\n",
        "    mask = mask.half()\n",
        "  \n",
        "  mask = 1-mask # inverting mask\n",
        "  out = netG(image, mask)\n",
        "  out = out*(1-mask)+image*mask\n",
        "  save_image(out, os.path.join(outdir, os.path.basename(f)))\n",
        "\n",
        "  # OOM hotfix\n",
        "  torch.cuda.empty_cache()"
      ],
      "metadata": {
        "id": "-pqCKBkJizbe",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title delete input/output\n",
        "%cd /content/\n",
        "!sudo rm -rf /content/input\n",
        "!mkdir /content/input\n",
        "!sudo rm -rf /content/output\n",
        "!mkdir /content/output"
      ],
      "metadata": {
        "cellView": "form",
        "id": "cVbBd4h7oRHY"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}