{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZjdQE0kKcqjA"
   },
   "source": [
    "<p align=\"center\">\n",
    "  <img src=\"https://user-images.githubusercontent.com/14334509/179359809-bd45566a-486d-418f-83fa-67bbbba8c45c.png\" height=120>\n",
    "</p>\n",
    "\n",
    "# CodeFormer Inference Demo \n",
    "## Towards Robust Blind Face Restoration with Codebook Lookup Transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_U5Bu-qie_aH",
    "tags": []
   },
   "source": [
    "# 1. Preparations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6HnTAAlTfFCY",
    "tags": []
   },
   "source": [
    "### Initalize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8SG9AcLQO_FQ"
   },
   "outputs": [],
   "source": [
    "# Clone CodeFormer and enter the CodeFormer folder\n",
    "%cd /content\n",
    "!rm -rf CodeFormer\n",
    "!git clone https://github.com/sczhou/CodeFormer.git\n",
    "%cd CodeFormer\n",
    "\n",
    "# Set up the environment\n",
    "# Install python dependencies\n",
    "!pip install -r requirements.txt\n",
    "# Install basicsr\n",
    "!python basicsr/setup.py develop\n",
    "\n",
    "# Download the pre-trained model\n",
    "!python scripts/download_pretrained_models.py facelib\n",
    "!python scripts/download_pretrained_models.py CodeFormer\n",
    "\n",
    "# Visualization function\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "def display(img1, img2):\n",
    "  fig = plt.figure(figsize=(25, 10))\n",
    "  ax1 = fig.add_subplot(1, 2, 1) \n",
    "  plt.title('Input', fontsize=16)\n",
    "  ax1.axis('off')\n",
    "  ax2 = fig.add_subplot(1, 2, 2)\n",
    "  plt.title('CodeFormer', fontsize=16)\n",
    "  ax2.axis('off')\n",
    "  ax1.imshow(img1)\n",
    "  ax2.imshow(img2)\n",
    "def imread(img_path):\n",
    "  img = cv2.imread(img_path)\n",
    "  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "  return img\n",
    "!mkdir /content/CodeFormer/inputs/user_upload"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ClNIcJqxSN2e"
   },
   "source": [
    "# Run!\n",
    "\n",
    "Upload your images to `/content/CodeFormer/inputs/user_upload`. The AI will process them all in bulk and tell you where it put the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "Cj2YQGg3J0TQ"
   },
   "outputs": [],
   "source": [
    "# Config\n",
    "\n",
    "CODEFORMER_FIDELITY = 0.44 # min:0, max:1 (try steps of 0.01)\n",
    "\n",
    "BACKGROUND_ENHANCE = True # Enhance background image with Real-ESRGAN\n",
    "\n",
    "FACE_UPSAMPLE = True # Upsample restored faces for high-resolution AI-created\n",
    "\n",
    "# =====================================================================\n",
    "\n",
    "if BACKGROUND_ENHANCE:\n",
    "  if FACE_UPSAMPLE:\n",
    "    !python inference_codeformer.py --w $CODEFORMER_FIDELITY --test_path inputs/user_upload --bg_upsampler realesrgan --face_upsample\n",
    "  else:\n",
    "    !python inference_codeformer.py --w $CODEFORMER_FIDELITY --test_path inputs/user_upload --bg_upsampler realesrgan\n",
    "else:\n",
    "  !python inference_codeformer.py --w $CODEFORMER_FIDELITY --test_path inputs/user_upload"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bn5EyOJJSgUa"
   },
   "source": [
    "# Tools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ddFQzqHWV1BM"
   },
   "source": [
    "### Delete uploaded images and results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rcm30PlmSh4O"
   },
   "outputs": [],
   "source": [
    "!rm -rf /notebooks/CodeFormer/inputs/user_upload/*\n",
    "!rm -rf /notebooks/CodeFormer/results/*"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "6HnTAAlTfFCY",
    "kScNto3vV4qD",
    "ddFQzqHWV1BM",
    "OxiB23_sdG8m"
   ],
   "private_outputs": true,
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}