nafnet.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "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/detektor777/colab_list/blob/main/nafnet.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "https://drive.google.com/drive"
      ],
      "metadata": {
        "id": "3dZa2XzOxAc_"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7igSG2WsruEG"
      },
      "outputs": [],
      "source": [
        "#@title ##**Install** { display-mode: \"form\" }\n",
        "%%capture\n",
        "!nvidia-smi\n",
        "%cd /content\n",
        "\n",
        "!git clone https://github.com/megvii-research/NAFNet\n",
        "\n",
        "%cd /content/NAFNet\n",
        "\n",
        "!pip install -r requirements.txt\n",
        "!pip install --upgrade --no-cache-dir gdown\n",
        "!python3 setup.py develop --no_cuda_ext\n",
        "!pip install Pillow\n",
        "!apt-get install -y ffmpeg\n",
        "!pip install moviepy imageio imageio-ffmpeg --quiet\n",
        "\n",
        "%cd /content/NAFNet\n",
        "\n",
        "import gdown\n",
        "\n",
        "import torch\n",
        "\n",
        "from basicsr.models import create_model\n",
        "from basicsr.utils import img2tensor as _img2tensor, tensor2img, imwrite\n",
        "from basicsr.utils.options import parse\n",
        "import numpy as np\n",
        "import cv2\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import os\n",
        "from google.colab import files\n",
        "import shutil\n",
        "import glob\n",
        "\n",
        "from google.colab import files\n",
        "import shutil\n",
        "from google.colab import files\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "\n",
        "from tqdm import tqdm\n",
        "import sys\n",
        "import re\n",
        "from PIL import Image, ImageEnhance\n",
        "\n",
        "import PIL.Image\n",
        "import numpy as np\n",
        "from IPython.display import display\n",
        "from IPython.display import display, clear_output\n",
        "\n",
        "import ipywidgets as widgets\n",
        "from IPython.display import display\n",
        "\n",
        "%cd /content/NAFNet\n",
        "\n",
        "# Deblur\n",
        "if not os.path.exists(\"./experiments/pretrained_models/NAFNet-REDS-width64.pth\"):\n",
        "  gdown.download('https://drive.google.com/uc?id=14D4V4raNYIOhETfcuuLI3bGLB-OYIv6X', \"./experiments/pretrained_models/\", quiet=False)\n",
        "\n",
        "def imread(img_path):\n",
        "  img = cv2.imread(img_path)\n",
        "  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
        "  return img\n",
        "def img2tensor(img, bgr2rgb=False, float32=True):\n",
        "    img = img.astype(np.float32) / 255.\n",
        "    return _img2tensor(img, bgr2rgb=bgr2rgb, float32=float32)\n",
        "\n",
        "def display(img1, img2):\n",
        "  fig = plt.figure(figsize=(25, 10))\n",
        "  ax1 = fig.add_subplot(1, 2, 1)\n",
        "  plt.title('Input image', fontsize=16)\n",
        "  ax1.axis('off')\n",
        "  ax2 = fig.add_subplot(1, 2, 2)\n",
        "  plt.title('NAFNet output', fontsize=16)\n",
        "  ax2.axis('off')\n",
        "  ax1.imshow(img1)\n",
        "  ax2.imshow(img2)\n",
        "\n",
        "def single_image_inference(model, img, save_path):\n",
        "      model.feed_data(data={'lq': img.unsqueeze(dim=0)})\n",
        "\n",
        "      if model.opt['val'].get('grids', False):\n",
        "          model.grids()\n",
        "\n",
        "      model.test()\n",
        "\n",
        "      if model.opt['val'].get('grids', False):\n",
        "          model.grids_inverse()\n",
        "\n",
        "      visuals = model.get_current_visuals()\n",
        "      sr_img = tensor2img([visuals['result']])\n",
        "      imwrite(sr_img, save_path)\n",
        "\n",
        "def sr_display(LR_l, LR_r, SR_l, SR_r):\n",
        "  h,w = SR_l.shape[:2]\n",
        "  LR_l = cv2.resize(LR_l, (w,h), interpolation=cv2.INTER_CUBIC)\n",
        "  LR_r = cv2.resize(LR_r, (w,h), interpolation=cv2.INTER_CUBIC)\n",
        "  fig = plt.figure(figsize=(w//40, h//40))\n",
        "  ax1 = fig.add_subplot(2, 2, 1)\n",
        "  plt.title('Input image (Left)', fontsize=16)\n",
        "  ax1.axis('off')\n",
        "  ax2 = fig.add_subplot(2, 2, 2)\n",
        "  plt.title('NAFSSR output (Left)', fontsize=16)\n",
        "  ax2.axis('off')\n",
        "  ax1.imshow(LR_l)\n",
        "  ax2.imshow(SR_l)\n",
        "\n",
        "  ax3 = fig.add_subplot(2, 2, 3)\n",
        "  plt.title('Input image (Right)', fontsize=16)\n",
        "  ax3.axis('off')\n",
        "  ax4 = fig.add_subplot(2, 2, 4)\n",
        "  plt.title('NAFSSR output (Right)', fontsize=16)\n",
        "  ax4.axis('off')\n",
        "  ax3.imshow(LR_r)\n",
        "  ax4.imshow(SR_r)\n",
        "\n",
        "  plt.subplots_adjust(wspace=0.04, hspace=0.04)\n",
        "\n",
        "def stereo_image_inference(model, img_l, img_r, save_path):\n",
        "      img = torch.cat([img_l, img_r], dim=0)\n",
        "      model.feed_data(data={'lq': img.unsqueeze(dim=0)})\n",
        "\n",
        "      if model.opt['val'].get('grids', False):\n",
        "          model.grids()\n",
        "\n",
        "      model.test()\n",
        "\n",
        "      if model.opt['val'].get('grids', False):\n",
        "          model.grids_inverse()\n",
        "\n",
        "      visuals = model.get_current_visuals()\n",
        "      img_L = visuals['result'][:,:3]\n",
        "      img_R = visuals['result'][:,3:]\n",
        "      img_L, img_R = tensor2img([img_L, img_R])\n",
        "\n",
        "      imwrite(img_L, save_path.format('L'))\n",
        "      imwrite(img_R, save_path.format('R'))\n",
        "\n",
        "%cd /content/NAFNet\n",
        "\n",
        "opt_path = 'options/test/REDS/NAFNet-width64.yml'\n",
        "opt = parse(opt_path, is_train=False)\n",
        "opt['dist'] = False\n",
        "NAFNet = create_model(opt)\n",
        "clear_output(wait=True)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##**Google drive** { display-mode: \"form\" }\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "id": "7XfKuRUrs_Iz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##**Upload** { display-mode: \"form\" }\n",
        "video_source = \"google drive\" #@param [\"pc\", \"google drive\"]\n",
        "\n",
        "if video_source == \"pc\":\n",
        "    from google.colab import files\n",
        "    uploaded = files.upload()\n",
        "    # Get the name of uploaded file\n",
        "    file_name = list(uploaded.keys())[0]\n",
        "    print(f\"Selected file: {file_name}\")\n",
        "elif video_source == \"google drive\":\n",
        "    from google.colab import drive\n",
        "    drive.mount('/content/drive')\n",
        "    import os\n",
        "    import ipywidgets as widgets\n",
        "    from IPython.display import display\n",
        "\n",
        "    # Search for video files in Google Drive\n",
        "    video_extensions = ('.mp4', '.avi', '.mov', '.mkv', '.flv', '.wmv')\n",
        "\n",
        "    video_files = []\n",
        "    for root, dirs, files in os.walk('/content/drive/MyDrive'):\n",
        "        for file in files:\n",
        "            if file.lower().endswith(video_extensions):\n",
        "                file_path = os.path.join(root, file)\n",
        "                video_files.append(file_path)\n",
        "\n",
        "    if not video_files:\n",
        "        print(\"No video files found in Google Drive.\")\n",
        "    else:\n",
        "        # Create dropdown list for file selection\n",
        "        video_dropdown = widgets.Dropdown(\n",
        "            options=video_files,\n",
        "            description='Select video file:',\n",
        "            style={'description_width': 'initial'},\n",
        "            layout=widgets.Layout(width='90%')\n",
        "        )\n",
        "\n",
        "        # Button for confirming selection\n",
        "        select_button = widgets.Button(description='Select file')\n",
        "        output = widgets.Output()\n",
        "\n",
        "        def on_button_clicked(b):\n",
        "            with output:\n",
        "                global file_name\n",
        "                file_name = video_dropdown.value\n",
        "                print(f\"Selected file: {file_name}\")\n",
        "\n",
        "        select_button.on_click(on_button_clicked)\n",
        "\n",
        "        display(video_dropdown, select_button, output)"
      ],
      "metadata": {
        "id": "ODhKxj4Nr7b0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##**Config** { display-mode: \"form\" }\n",
        "\n",
        "continue_status = \"new\" #@param [\"continue\",\"new\"]\n",
        "opacity = \"20\" #@param [10, 20, 30, 40, 50, 60, 70, 80, 90, 100]\n",
        "\n",
        "sharpness_and_upscale = \"yes\" #@param [\"yes\",\"no\"]\n",
        "sharpness_factor = 3  # @param {type:\"slider\", min:0.0, max:7.0, step:0.1}\n",
        "upscale_option = \"height 720px\"  # @param [\"no\", \"height 720px\", \"height 1080px\", \"height 2160px\", \"2x\", \"4x\"]\n",
        "\n",
        "from PIL import ImageEnhance\n",
        "\n",
        "%cd /content/NAFNet\n",
        "\n",
        "path = '/content/drive/MyDrive/'\n",
        "deblur_upload_folder = '/content/drive/MyDrive/deblur_input'\n",
        "deblur_result_folder = '/content/drive/MyDrive/deblur_output'\n",
        "deblur_result_folder_overlay = '/content/drive/MyDrive/deblur_output_overlay'\n",
        "deblur_result_folder_sharp = '/content/drive/MyDrive/deblur_output_overlay_sharp'\n",
        "\n",
        "base_result_folder = deblur_upload_folder\n",
        "\n",
        "if not os.path.exists(deblur_upload_folder):\n",
        "    os.makedirs(deblur_upload_folder)\n",
        "\n",
        "if not os.path.exists(deblur_result_folder):\n",
        "    os.makedirs(deblur_result_folder)\n",
        "\n",
        "if not os.path.exists(deblur_result_folder_overlay):\n",
        "    os.makedirs(deblur_result_folder_overlay)\n",
        "\n",
        "if not os.path.exists(deblur_result_folder_sharp):\n",
        "    os.makedirs(deblur_result_folder_sharp)\n",
        "\n",
        "if continue_status == \"new\":\n",
        "  if os.path.isdir(deblur_upload_folder):\n",
        "      shutil.rmtree(deblur_upload_folder)\n",
        "\n",
        "  os.makedirs(deblur_upload_folder)\n",
        "\n",
        "  if os.path.isdir(deblur_result_folder):\n",
        "      shutil.rmtree(deblur_result_folder)\n",
        "\n",
        "  os.makedirs(deblur_result_folder)\n",
        "\n",
        "  if os.path.isdir(deblur_result_folder_overlay):\n",
        "      shutil.rmtree(deblur_result_folder_overlay)\n",
        "\n",
        "  os.makedirs(deblur_result_folder_overlay)\n",
        "\n",
        "  if os.path.isdir(deblur_result_folder_sharp):\n",
        "      shutil.rmtree(deblur_result_folder_sharp)\n",
        "\n",
        "  os.makedirs(deblur_result_folder_sharp)\n"
      ],
      "metadata": {
        "id": "Ge3K-TBd0p7l"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##**Sequence** { display-mode: \"form\" }\n",
        "import cv2\n",
        "import imageio\n",
        "import os\n",
        "import tqdm\n",
        "import subprocess\n",
        "import numpy as np\n",
        "import time\n",
        "\n",
        "video_convertion = \"no\" #@param [\"yes\",\"no\"]\n",
        "if (video_convertion == \"yes\"):\n",
        "\n",
        "  output_file = os.path.splitext(file_name)[0] + \"_converted.mp4\"\n",
        "\n",
        "  !ffmpeg -i \"{file_name}\" -c:v libx264 -preset slow -crf 18 -c:a aac -b:a 192k -strict experimental -y \"{output_file}\"\n",
        "\n",
        "  if os.path.exists(output_file):\n",
        "      file_name = output_file\n",
        "      print(f\"Video successfully converted: {file_name}\")\n",
        "  else:\n",
        "      print(\"Video conversion error\")\n",
        "\n",
        "library = \"imageio\" #@param [\"cv2\",\"pyav\",\"imageio\",\"ffmpeg\",\"skvideo\",\"scipy\",\"moviepy\"]\n",
        "delay = \"0.05\" #@param [0, 0.01, 0.05, 0.05, 0.1, 0.2, 0.3, 0.4, 0.5]\n",
        "\n",
        "full_path = os.path.join(path, file_name)\n",
        "\n",
        "if (library == \"cv2\"):\n",
        "    cmd = 'ffprobe -v error -select_streams v:0 -count_packets -show_entries stream=nb_read_packets -of csv=p=0 ' + full_path\n",
        "    output = subprocess.check_output(cmd, shell=True).decode('utf-8').strip()\n",
        "\n",
        "    cmd = 'ffprobe -v error -select_streams v:0 -show_entries stream=r_frame_rate -of csv=p=0 ' + full_path\n",
        "    output = subprocess.check_output(cmd, shell=True).decode('utf-8').strip()\n",
        "    vidcap = cv2.VideoCapture(full_path)\n",
        "    fps = vidcap.get(cv2.CAP_PROP_FPS)\n",
        "    frame_count = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
        "\n",
        "    duration = frame_count / fps\n",
        "\n",
        "    print(\"FPS: \", fps)\n",
        "    print(\"Duration: \", duration)\n",
        "    print(\"Frames: \", frame_count)\n",
        "\n",
        "    vidcap = cv2.VideoCapture(full_path)\n",
        "    success, image = vidcap.read()\n",
        "\n",
        "    tqdm.tqdm.set_lock(tqdm.tqdm.get_lock())\n",
        "    pbar_cv2 = tqdm.tqdm(total=frame_count, ncols=100, position=0, leave=True)\n",
        "    while success:\n",
        "        try:\n",
        "            cv2.imwrite(f\"{deblur_upload_folder}/{int(vidcap.get(cv2.CAP_PROP_POS_FRAMES)):09d}.jpg\", image)\n",
        "        except cv2.error:\n",
        "            print(\"Error writing to disk. Retrying...\")\n",
        "            continue\n",
        "        pbar_cv2.update(1)\n",
        "        time.sleep(float(delay))\n",
        "        success, image = vidcap.read()\n",
        "\n",
        "    pbar_cv2.close()\n",
        "    vidcap.release()\n",
        "\n",
        "elif (library == \"pyav\"):\n",
        "    !pip install av\n",
        "    import av\n",
        "    video = av.open(full_path)\n",
        "\n",
        "    fps_value = video.streams.video[0].average_rate\n",
        "    frame_count = video.streams.video[0].frames\n",
        "\n",
        "    if fps_value is not None:\n",
        "        duration = frame_count / fps_value\n",
        "\n",
        "        print(\"FPS: \", fps_value)\n",
        "        print(\"Duration: \", duration)\n",
        "        print(\"Frames: \", frame_count)\n",
        "\n",
        "        tqdm.tqdm.set_lock(tqdm.tqdm.get_lock())\n",
        "        pbar_pyav = tqdm.tqdm(total=frame_count, ncols=100, position=0, leave=True)\n",
        "        for packet in video.demux():\n",
        "            for frame in packet.decode():\n",
        "                img = frame.to_ndarray(format='rgb24')\n",
        "                img = img.astype('uint8')\n",
        "                img_path = f\"{deblur_upload_folder}/{packet.pts:09d}.jpg\"\n",
        "                imageio.imwrite(img_path, img)\n",
        "                pbar_pyav.update(1)\n",
        "                time.sleep(float(delay))\n",
        "\n",
        "        pbar_pyav.close()\n",
        "        video.close()\n",
        "    else:\n",
        "        print(\"Error: FPS value is None.\")\n",
        "\n",
        "elif (library == \"imageio\"):\n",
        "    video = imageio.get_reader(full_path)\n",
        "\n",
        "    fps = video.get_meta_data()['fps']\n",
        "    frame_count = video.count_frames()\n",
        "    duration = frame_count / fps\n",
        "\n",
        "    print(\"FPS: \", fps)\n",
        "    print(\"Duration: \", duration)\n",
        "    print(\"Frames: \", frame_count)\n",
        "\n",
        "    tqdm.tqdm.set_lock(tqdm.tqdm.get_lock())\n",
        "    pbar_imageio = tqdm.tqdm(total=frame_count, ncols=100, position=0, leave=True)\n",
        "    for i, frame in enumerate(video):\n",
        "      img_path = f\"{deblur_upload_folder}/{i+1:09d}.jpg\"\n",
        "      if os.path.isfile(img_path):\n",
        "          pbar_imageio.update(1)\n",
        "          continue\n",
        "      while True:\n",
        "          try:\n",
        "              imageio.imwrite(img_path, frame)\n",
        "          except Exception as e:\n",
        "              print(f\"Error writing to disk: {str(e)}. Retrying...\")\n",
        "              continue\n",
        "          break\n",
        "      pbar_imageio.update(1)\n",
        "      time.sleep(float(delay))\n",
        "\n",
        "    pbar_imageio.close()\n",
        "    video.close()\n",
        "\n",
        "elif (library == \"ffmpeg\"):\n",
        "    !pip install ffmpeg-python\n",
        "    import ffmpeg\n",
        "\n",
        "    probe = ffmpeg.probe(full_path)\n",
        "    video_info = next(stream for stream in probe['streams'] if stream['codec_type'] == 'video')\n",
        "    fps = video_info['r_frame_rate']\n",
        "    duration = float(video_info['duration'])\n",
        "    frame_count = int(video_info['nb_frames'])\n",
        "\n",
        "    print(\"FPS: \", fps)\n",
        "    print(\"Duration: \", duration)\n",
        "    print(\"Frames: \", frame_count)\n",
        "\n",
        "    pbar_ffmpeg = tqdm.tqdm(total=frame_count, ncols=100, position=0, leave=True)\n",
        "    process = (\n",
        "        ffmpeg\n",
        "        .input(full_path)\n",
        "        .output('pipe:', format='rawvideo', pix_fmt='rgb24', qscale=0)\n",
        "        .run_async(pipe_stdout=True)\n",
        "    )\n",
        "\n",
        "    for i in range(frame_count):\n",
        "        try:\n",
        "            raw_video = process.stdout.read(video_info['width'] * video_info['height'] * 3)\n",
        "            frame = np.frombuffer(raw_video, dtype='uint8').reshape((video_info['height'], video_info['width'], 3))\n",
        "            frame_path = f\"{deblur_upload_folder}/{i:09d}.jpg\"\n",
        "            if os.path.isfile(frame_path):\n",
        "              pbar_ffmpeg.update(1)\n",
        "              continue\n",
        "            imageio.imwrite(frame_path, frame)\n",
        "        except Exception as e:\n",
        "            print(f\"Error writing to disk: {str(e)}. Retrying...\")\n",
        "            continue\n",
        "        pbar_ffmpeg.update(1)\n",
        "        time.sleep(float(delay))\n",
        "\n",
        "    pbar_ffmpeg.close()\n",
        "    process.wait()\n",
        "\n",
        "elif (library == \"skvideo\"):\n",
        "    !pip install scikit-video\n",
        "    import skvideo.io\n",
        "\n",
        "    video = skvideo.io.vread(full_path)\n",
        "\n",
        "    fps_str = skvideo.io.ffprobe(full_path)['video']['@avg_frame_rate']\n",
        "    fps_parts = fps_str.split('/')\n",
        "    fps = int(fps_parts[0]) / int(fps_parts[1])\n",
        "    frame_count = len(video)\n",
        "    duration = frame_count / fps\n",
        "\n",
        "    print(\"FPS: \", fps)\n",
        "    print(\"Duration: \", duration)\n",
        "    print(\"Frames: \", frame_count)\n",
        "\n",
        "    tqdm.tqdm.set_lock(tqdm.tqdm.get_lock())\n",
        "    outputdict = {'-q:1': '1'}\n",
        "    pbar_skvideo = tqdm.tqdm(total=frame_count, ncols=100, position=0, leave=True)\n",
        "    for i, frame in enumerate(video):\n",
        "        while True:\n",
        "            try:\n",
        "                skvideo.io.vwrite(f\"{deblur_upload_folder}/{i:09d}.jpg\", frame, outputdict=outputdict)\n",
        "            except Exception as e:\n",
        "                print(f\"Error writing to disk: {str(e)}. Retrying...\")\n",
        "                continue\n",
        "            break\n",
        "        pbar_skvideo.update(1)\n",
        "        time.sleep(float(delay))\n",
        "\n",
        "    pbar_skvideo.close()\n",
        "\n",
        "elif (library == \"scipy\"):\n",
        "    import scipy.misc\n",
        "    import scipy.ndimage\n",
        "\n",
        "    video = imageio.get_reader(full_path)\n",
        "    fps_value = video.get_meta_data()['fps']\n",
        "    frame_count = video.get_length()\n",
        "\n",
        "    if fps_value is not None:\n",
        "        duration = frame_count / fps_value\n",
        "\n",
        "        print(\"FPS: \", fps_value)\n",
        "\n",
        "        tqdm.tqdm.set_lock(tqdm.tqdm.get_lock())\n",
        "        pbar_scipy = tqdm.tqdm(total=frame_count, ncols=100, position=0, leave=True)\n",
        "        for i, im in enumerate(video):\n",
        "            img_path = f\"{deblur_upload_folder}/{i+1:09d}.jpg\"\n",
        "            im = Image.fromarray(im)\n",
        "            im.save(img_path, quality=100)\n",
        "            pbar_scipy.update(1)\n",
        "            time.sleep(float(delay))\n",
        "\n",
        "        pbar_scipy.close()\n",
        "    else:\n",
        "        print(\"Error: FPS value is None.\")\n",
        "\n",
        "elif library == \"moviepy\":\n",
        "    from moviepy.editor import VideoFileClip\n",
        "    from PIL import Image\n",
        "    video = VideoFileClip(full_path)\n",
        "    fps_value = video.fps\n",
        "    frame_count = int(video.duration * fps_value)\n",
        "\n",
        "    print(\"FPS: \", fps_value)\n",
        "\n",
        "    tqdm.tqdm.set_lock(tqdm.tqdm.get_lock())\n",
        "    pbar_moviepy = tqdm.tqdm(total=frame_count, ncols=100, position=0, leave=True)\n",
        "    for i, frame in enumerate(video.iter_frames()):\n",
        "        img_path = f\"{deblur_upload_folder}/{i+1:09d}.jpg\"\n",
        "        im = Image.fromarray(frame)\n",
        "        im.save(img_path, quality=100)\n",
        "        pbar_moviepy.update(1)\n",
        "        time.sleep(float(delay))\n",
        "\n",
        "    pbar_moviepy.close()"
      ],
      "metadata": {
        "id": "gZPx11AL8_Od"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##**Check for missing frames(optional)** { display-mode: \"form\" }\n",
        "import os\n",
        "import time\n",
        "\n",
        "frame_dir = deblur_upload_folder\n",
        "\n",
        "while True:\n",
        "    try:\n",
        "        frames = [int(f.split('.')[0]) for f in os.listdir(frame_dir) if f.endswith('.jpg')]\n",
        "        min_frame = min(frames)\n",
        "        max_frame = max(frames)\n",
        "        print(min_frame)\n",
        "        print(max_frame)\n",
        "        break  # Exit the loop if successful\n",
        "    except Exception as e:\n",
        "        print(f\"Error reading the list of files: {e}. Retrying in 2 seconds\")\n",
        "        time.sleep(2)\n",
        "    except BaseException as e:\n",
        "        # Handle exceptions like KeyboardInterrupt and SystemExit\n",
        "        print(f\"Critical error: {e}. Retrying in 2 seconds\")\n",
        "        time.sleep(2)\n"
      ],
      "metadata": {
        "id": "scVEYtVH89on"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##**Run deblur** { display-mode: \"form\" }\n",
        "%cd /content/NAFNet\n",
        "import glob\n",
        "import os\n",
        "import time\n",
        "from tqdm import tqdm\n",
        "\n",
        "\n",
        "upload_files = set()\n",
        "while not upload_files:\n",
        "    try:\n",
        "        upload_files = set(os.listdir(deblur_upload_folder))\n",
        "    except OSError as e:\n",
        "        print(f\"Error: {e}, retrying...\")\n",
        "        time.sleep(2)\n",
        "\n",
        "result_files = set(os.listdir(deblur_result_folder))\n",
        "\n",
        "input_list = [file_path for file_path in glob.glob(os.path.join(deblur_upload_folder, '*'))\n",
        "              if os.path.basename(file_path) not in result_files and os.path.basename(file_path) in upload_files]\n",
        "\n",
        "\n",
        "for input_path in tqdm(input_list, desc=\"Processing images\"):\n",
        "    img_input = imread(input_path)\n",
        "    inp = img2tensor(img_input)\n",
        "    output_path = os.path.join(deblur_result_folder, os.path.basename(input_path))\n",
        "    single_image_inference(NAFNet, inp, output_path)\n",
        "\n",
        "\n",
        "#check missing frames\n",
        "frames = None\n",
        "while frames is None:\n",
        "    try:\n",
        "        frames = [int(f.split('.')[0]) for f in os.listdir(deblur_result_folder) if f.endswith('.jpg')]\n",
        "    except:\n",
        "        print(\"Error reading the list of files. Retrying in 2 seconds\")\n",
        "        time.sleep(2)\n",
        "\n",
        "min_frame = min(frames)\n",
        "max_frame = max(frames)\n",
        "print(min_frame)\n",
        "print(max_frame)\n",
        "\n",
        "missing_frames = []\n",
        "for i in range(min_frame, max_frame + 1):\n",
        "    if i not in frames:\n",
        "        missing_frames.append(i)\n",
        "\n",
        "if len(missing_frames) > 0:\n",
        "    print(f\"Missing frames: {missing_frames}\")\n",
        "else:\n",
        "    print(\"All frames present\")\n",
        "base_result_folder = deblur_result_folder"
      ],
      "metadata": {
        "id": "ZBSVjJ4jDx8s"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##**Run overlay** { display-mode: \"form\" }\n",
        "if opacity != \"100\":\n",
        "  base_result_folder = deblur_result_folder_overlay\n",
        "\n",
        "  if os.path.isdir(deblur_result_folder_overlay):\n",
        "      shutil.rmtree(deblur_result_folder_overlay)\n",
        "  os.makedirs(deblur_result_folder_overlay)\n",
        "\n",
        "  max_retries = 10\n",
        "  retries = 0\n",
        "  while retries < max_retries:\n",
        "      try:\n",
        "          filenames = os.listdir(deblur_upload_folder)\n",
        "          break\n",
        "      except:\n",
        "          print(\"Error getting file list. Trying\", retries+1, \"from\", max_retries)\n",
        "          retries += 1\n",
        "          time.sleep(2)\n",
        "\n",
        "  for filename in filenames:\n",
        "\n",
        "      try:\n",
        "          img1 = Image.open(os.path.join(deblur_upload_folder, filename))\n",
        "          img2 = Image.open(os.path.join(deblur_result_folder, filename))\n",
        "      except:\n",
        "          print(f\"Error opening file {filename}\")\n",
        "          time.sleep(2)\n",
        "          continue\n",
        "\n",
        "\n",
        "      img1 = img1.convert(\"RGBA\")\n",
        "      img2 = img2.convert(\"RGBA\")\n",
        "\n",
        "      width1, height1 = img1.size\n",
        "      width2, height2 = img2.size\n",
        "\n",
        "      img2.putalpha(int(float(opacity) * 2.55))\n",
        "\n",
        "      if width1 * height1 > width2 * height2:\n",
        "          img2 = img2.resize((width1, height1))\n",
        "      else:\n",
        "          img1 = img1.resize((width2, height2))\n",
        "\n",
        "      result = Image.alpha_composite(img1, img2)\n",
        "\n",
        "      result = result.convert(\"RGB\")\n",
        "\n",
        "      result.save(os.path.join(deblur_result_folder_overlay, filename), quality=100)\n",
        "deblur_result_folder = '/content/drive/MyDrive/deblur_output_overlay'"
      ],
      "metadata": {
        "id": "Uw57pKFFKAwy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##**Run sharpness** { display-mode: \"form\" }\n",
        "if sharpness_and_upscale == \"yes\":\n",
        "  import shutil\n",
        "  from tqdm import tqdm\n",
        "  import os\n",
        "  import sys\n",
        "  import re\n",
        "  from PIL import Image, ImageEnhance\n",
        "\n",
        "  source_folder = base_result_folder\n",
        "  destination_folder = deblur_result_folder_sharp\n",
        "\n",
        "  if not os.path.exists(destination_folder):\n",
        "      os.makedirs(destination_folder)\n",
        "\n",
        "  file_list = os.listdir(source_folder)\n",
        "  file_list.sort()\n",
        "  frames = [int(f.split('.')[0]) for f in file_list if f.endswith('.jpg')]\n",
        "  min_frame = min(frames)\n",
        "\n",
        "  dest_files = os.listdir(destination_folder)\n",
        "  if dest_files:\n",
        "      dest_frames = [int(re.findall(r'(\\d+)\\.jpg', f)[0]) for f in dest_files if re.match(r'\\d+\\.jpg', f)]\n",
        "      start_frame = max(dest_frames) + 1\n",
        "  else:\n",
        "      start_frame = min_frame\n",
        "\n",
        "  max_frame = frames[-1]\n",
        "  print(\"Max frame:\", max_frame)\n",
        "  files_to_copy = [f\"{source_folder}/{frame:09d}.jpg\" for frame in range(start_frame, max_frame+1) if f\"{frame:09d}.jpg\" in file_list]\n",
        "\n",
        "  total_files = len(files_to_copy)\n",
        "  batch_size = 10\n",
        "  num_iterations = (total_files + batch_size - 1) // batch_size\n",
        "\n",
        "  print(\"Start frame:\", start_frame)\n",
        "  print(\"Min frame:\", min_frame)\n",
        "  print(\"Count files:\", total_files)\n",
        "  print(\"Count iterations:\", num_iterations)\n",
        "\n",
        "  with tqdm(total=num_iterations) as pbar:\n",
        "      for i in range(0, total_files, batch_size):\n",
        "          files_to_copy_batch = files_to_copy[i:i+batch_size]\n",
        "\n",
        "          for file in files_to_copy_batch:\n",
        "              filename = os.path.basename(file)\n",
        "              dest_path = os.path.join(destination_folder, filename)\n",
        "\n",
        "              original_image = Image.open(file)\n",
        "\n",
        "              enhancer = ImageEnhance.Sharpness(original_image)\n",
        "              processed_image = enhancer.enhance(sharpness_factor)\n",
        "\n",
        "              if upscale_option != \"no\":\n",
        "                  width, height = processed_image.size\n",
        "\n",
        "                  if upscale_option == \"height 720px\":\n",
        "                      new_height = 720\n",
        "                      new_width = int(width * (new_height / height))\n",
        "                  elif upscale_option == \"height 1080px\":\n",
        "                      new_height = 1080\n",
        "                      new_width = int(width * (new_height / height))\n",
        "                  elif upscale_option == \"height 2160px\":\n",
        "                      new_height = 2160\n",
        "                      new_width = int(width * (new_height / height))\n",
        "                  elif upscale_option == \"2x\":\n",
        "                      new_width = width * 2\n",
        "                      new_height = height * 2\n",
        "                  elif upscale_option == \"4x\":\n",
        "                      new_width = width * 4\n",
        "                      new_height = height * 4\n",
        "                  else:\n",
        "                      new_width, new_height = width, height\n",
        "\n",
        "                  processed_image = processed_image.resize((new_width, new_height), Image.LANCZOS)\n",
        "\n",
        "              copied = False\n",
        "              while not copied:\n",
        "                  try:\n",
        "                      processed_image.save(dest_path, quality=100)\n",
        "                      copied = True\n",
        "                  except:\n",
        "                      print(f\"Failed to save file {dest_path}. Trying again...\")\n",
        "                      continue\n",
        "\n",
        "          pbar.update(1)\n",
        "\n",
        "  print(\"Processing complete\")\n",
        "  base_result_folder = deblur_result_folder_sharp"
      ],
      "metadata": {
        "id": "E_zIftrMMmp_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##**Create Video** { display-mode: \"form\" }\n",
        "\n",
        "import os\n",
        "from moviepy.editor import VideoFileClip, ImageSequenceClip\n",
        "\n",
        "quality = 18  # @param {type:\"slider\", min:0.0, max:50, step:1}\n",
        "\n",
        "output_file = os.path.join(path, f\"{os.path.splitext(file_name)[0]}_output.mp4\")\n",
        "full_path = os.path.join(path, file_name)\n",
        "\n",
        "print(\"Original video:\", full_path)\n",
        "\n",
        "original_video = VideoFileClip(full_path)\n",
        "fps_of_video = original_video.fps\n",
        "audio_clip = original_video.audio\n",
        "\n",
        "print(\"Frame rate:\", fps_of_video)\n",
        "\n",
        "img_files = os.listdir(base_result_folder)\n",
        "img_files.sort()\n",
        "\n",
        "if not img_files:\n",
        "    print(\"There are no frames to process in the folder:\", destination_folder)\n",
        "else:\n",
        "    img_paths = [os.path.join(destination_folder, img_file) for img_file in img_files]\n",
        "\n",
        "    video_clip = ImageSequenceClip(img_paths, fps=fps_of_video)\n",
        "    video_clip = video_clip.set_audio(audio_clip)\n",
        "\n",
        "    video_clip.write_videofile(\n",
        "        output_file,\n",
        "        codec=\"libx264\",\n",
        "        audio_codec=\"aac\",\n",
        "        ffmpeg_params=[\"-crf\", str(quality)]\n",
        "    )\n",
        "\n",
        "    video_clip.close()\n",
        "    original_video.close()\n",
        "    del video_clip\n",
        "    del original_video\n",
        "\n",
        "    print('Assembly complete. Final video saved to path:', output_file)\n"
      ],
      "metadata": {
        "id": "SpdOLKcKPKSo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##**Download file (optional, you can download it faster from Google Drive)** { display-mode: \"form\" }\n",
        "from google.colab import files\n",
        "files.download(output_file)"
      ],
      "metadata": {
        "id": "CpO1bEBfRs9b"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}