mttr_interactive_demo.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "mttr_interactive_demo",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "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/cedro3/others2/blob/main/mttr_interactive_demo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c7ZB0W8mIOzx"
      },
      "source": [
        "**セットアップ**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hbdn_2C-RO5Y"
      },
      "source": [
        "# moviepyインストール\n",
        "# %%capture\n",
        "!pip install av moviepy yt-dlp ruamel.yaml einops timm transformers\n",
        "\n",
        "# ライブラリのインポート\n",
        "import torch\n",
        "import torchvision\n",
        "import torchvision.transforms.functional as F\n",
        "from einops import rearrange\n",
        "import numpy as np\n",
        "from PIL import Image, ImageDraw, ImageOps, ImageFont\n",
        "from yt_dlp import YoutubeDL\n",
        "from moviepy.editor import VideoFileClip, AudioFileClip, ImageSequenceClip\n",
        "from moviepy.video.io.ffmpeg_tools import ffmpeg_extract_subclip\n",
        "from IPython.display import HTML\n",
        "from base64 import b64encode\n",
        "from tqdm.notebook import trange, tqdm\n",
        "from transformers import logging\n",
        "# logging.set_verbosity_error()\n",
        "\n",
        "# MTTRモデル初期化\n",
        "model, postprocessor = torch.hub.load('mttr2021/MTTR:main','mttr_refer_youtube_vos', force_reload=True)\n",
        "model = model.cuda()\n",
        "\n",
        "# 関数定義\n",
        "class NestedTensor(object):\n",
        "    def __init__(self, tensors, mask):\n",
        "        self.tensors = tensors\n",
        "        self.mask = mask\n",
        "\n",
        "def nested_tensor_from_videos_list(videos_list):\n",
        "    def _max_by_axis(the_list):\n",
        "      maxes = the_list[0]\n",
        "      for sublist in the_list[1:]:\n",
        "          for index, item in enumerate(sublist):\n",
        "              maxes[index] = max(maxes[index], item)\n",
        "      return maxes\n",
        "\n",
        "    max_size = _max_by_axis([list(img.shape) for img in videos_list])\n",
        "    padded_batch_shape = [len(videos_list)] + max_size\n",
        "    b, t, c, h, w = padded_batch_shape\n",
        "    dtype = videos_list[0].dtype\n",
        "    device = videos_list[0].device\n",
        "    padded_videos = torch.zeros(padded_batch_shape, dtype=dtype, device=device)\n",
        "    videos_pad_masks = torch.ones((b, t, h, w), dtype=torch.bool, device=device)\n",
        "    for vid_frames, pad_vid_frames, vid_pad_m in zip(videos_list, padded_videos, videos_pad_masks):\n",
        "        pad_vid_frames[:vid_frames.shape[0], :, :vid_frames.shape[2], :vid_frames.shape[3]].copy_(vid_frames)\n",
        "        vid_pad_m[:vid_frames.shape[0], :vid_frames.shape[2], :vid_frames.shape[3]] = False\n",
        "    return NestedTensor(padded_videos.transpose(0, 1), videos_pad_masks.transpose(0, 1))\n",
        "\n",
        "def apply_mask(image, mask, color, transparency=0.7):\n",
        "    mask = mask[..., np.newaxis].repeat(repeats=3, axis=2)\n",
        "    mask = mask * transparency\n",
        "    color_matrix = np.ones(image.shape, dtype=np.float) * color\n",
        "    out_image = color_matrix * mask + image * (1.0 - mask)\n",
        "    return out_image\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**動画選択とクエリ指定**"
      ],
      "metadata": {
        "id": "I6g9MPbj8P2l"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XbAe8Djicw20"
      },
      "source": [
        "# Choose (by un-commenting) one of the following:\n",
        "# video_url, (start_pt, end_pt), text_queries =  'https://www.youtube.com/watch?v=YThX7_8I3m0', (233, 243), ['guy in black performing tricks on a bike', 'a black bike used to perform tricks']\n",
        "# video_url, (start_pt, end_pt), text_queries = 'https://www.youtube.com/watch?v=hwLo7aU1Aas', (1144, 1152), ['a man riding a surfboard', 'a black and white surfboard']\n",
        "# video_url, (start_pt, end_pt), text_queries =  'https://www.youtube.com/watch?v=yvJDHbrumak', (48, 55), ['a red ball thrown in the air', 'a black horse playing with a person']\n",
        "# video_url, (start_pt, end_pt), text_queries =  'https://www.youtube.com/watch?v=L-Wd4A8ESyk', (289, 297), ['a guy performing tricks on a skateboard', 'a black skateboard']\n",
        "# video_url, (start_pt, end_pt), text_queries =  'https://www.youtube.com/watch?v=4iTiRvk4FHY', (24, 34), ['man in red shirt playing tennis', 'white tennis racket held by a man in a red shirt']\n",
        "# video_url, (start_pt, end_pt), text_queries =  'https://www.youtube.com/watch?v=ZHwlmvuW4NY', (115, 125), ['white dog playing', 'brown and black dog playing']\n",
        "video_url, (start_pt, end_pt), text_queries =  'https://www.youtube.com/watch?v=YThX7_8I3m0', (67, 77), ['guy in white shirt performing tricks on a bike', 'a black bike used to perform tricks']\n",
        "# video_url, (start_pt, end_pt), text_queries =  'https://www.youtube.com/watch?v=C7TCH927--g', (3, 13), ['a dog to the right', 'a cat to the left']\n",
        "# video_url, (start_pt, end_pt), text_queries =  'https://www.youtube.com/watch?v=0Z_WAF1GKfk', (143.5, 147.5), ['a dog to the left playing with a toy', 'a dog to the right playing with a toy']\n",
        "# video_url, (start_pt, end_pt), text_queries =  'https://www.youtube.com/watch?v=aEJJmebTLEs', (70, 80), ['a person hugging a dog', 'a white dog sitting']\n",
        "# video_url, (start_pt, end_pt), text_queries =  'https://www.youtube.com/watch?v=8sDF8lflCTs' ,(15, 23), ['person in blue riding a bike']\n",
        "# video_url, (start_pt, end_pt), text_queries =  'https://www.youtube.com/watch?v=dQw4w9WgXcQ', (2.5, 7.5), ['a person dancing']\n",
        "\n",
        "\n",
        "#OR - try your using own input in the following format: (but keep in mind that performance may be limited!)\n",
        "# video_url, (start_pt, end_pt), text_queries =  f'https://www.youtube.com/watch?v=???????' ,(start_pnt, end_pnt), ['text query 1', 'text query 2']\n",
        "\n",
        "assert  0 < end_pt - start_pt <= 10, 'error - the subclip length must be 0-10 seconds long'\n",
        "assert  1 <= len(text_queries) <= 2, 'error - 1-2 input text queries are expected'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HqVBgST0N7NQ"
      },
      "source": [
        "**動画のダウンロード**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BKkfRZZHdzJg"
      },
      "source": [
        "download_resolution = 360\n",
        "full_video_path = 'full_video.mp4'\n",
        "input_clip_path = 'input_clip.mp4'\n",
        "\n",
        "# download parameters:\n",
        "ydl_opts = {'format': f'best[height<={download_resolution}]', 'overwrites': True, 'outtmpl': full_video_path}\n",
        "# download the whole video:\n",
        "with YoutubeDL(ydl_opts) as ydl:\n",
        "    ydl.download([video_url])\n",
        "\n",
        "# extract the relevant subclip:\n",
        "with VideoFileClip(full_video_path) as video:\n",
        "    subclip = video.subclip(start_pt, end_pt)\n",
        "    subclip.write_videofile(input_clip_path)\n",
        "    \n",
        "# visualize the input clip:\n",
        "input_clip = open(input_clip_path,'rb').read()\n",
        "data_url = \"data:video/mp4;base64,\" + b64encode(input_clip).decode()\n",
        "HTML(\"\"\"<video width=720 controls><source src=\"%s\" type=\"video/mp4\"></video>\"\"\" % data_url)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vtF7RoBRojcR"
      },
      "source": [
        "**インスタンスマスクの生成**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WS4FhVRnwtur"
      },
      "source": [
        "window_length = 24  # length of window during inference\n",
        "window_overlap = 6  # overlap (in frames) between consecutive windows\n",
        "\n",
        "with torch.inference_mode():\n",
        "  # read and preprocess the video clip:\n",
        "  video, audio, meta = torchvision.io.read_video(filename=input_clip_path)\n",
        "  video = rearrange(video, 't h w c -> t c h w')\n",
        "  input_video = F.resize(video, size=360, max_size=640).cuda()\n",
        "  input_video = input_video.to(torch.float).div_(255)\n",
        "  input_video = F.normalize(input_video, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
        "  video_metadata = {'resized_frame_size': input_video.shape[-2:], 'original_frame_size': video.shape[-2:]}\n",
        "  \n",
        "  # partition the clip into overlapping windows of frames:\n",
        "  windows = [input_video[i:i+window_length] for i in range(0, len(input_video), window_length - window_overlap)]\n",
        "  # clean up the text queries:\n",
        "  text_queries = [\" \".join(q.lower().split()) for q in text_queries]\n",
        "\n",
        "  pred_masks_per_query = []\n",
        "  t, _, h, w = video.shape\n",
        "  for text_query in tqdm(text_queries, desc='text queries'):\n",
        "    pred_masks = torch.zeros(size=(t, 1, h, w))\n",
        "    for i, window in enumerate(tqdm(windows, desc='windows')):\n",
        "      window = nested_tensor_from_videos_list([window])\n",
        "      valid_indices = torch.arange(len(window.tensors)).cuda()\n",
        "      outputs = model(window, valid_indices, [text_query])\n",
        "      window_masks = postprocessor(outputs, [video_metadata], window.tensors.shape[-2:])[0]['pred_masks']\n",
        "      win_start_idx = i*(window_length-window_overlap)\n",
        "      pred_masks[win_start_idx:win_start_idx + window_length] = window_masks\n",
        "    pred_masks_per_query.append(pred_masks)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Th4tQP0Wo9W6"
      },
      "source": [
        "**動画にインスタンス・マスクとクエリを適用**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "deCqcgXiXUKL"
      },
      "source": [
        "# RGB colors for instance masks:\n",
        "light_blue = (41, 171, 226)\n",
        "purple = (237, 30, 121)\n",
        "dark_green = (35, 161, 90)\n",
        "orange = (255, 148, 59)\n",
        "colors = np.array([light_blue, purple, dark_green, orange])\n",
        "\n",
        "# width (in pixels) of the black strip above the video on which the text queries will be displayed:\n",
        "text_border_height_per_query = 35\n",
        "\n",
        "video_np = rearrange(video, 't c h w -> t h w c').numpy() / 255.0\n",
        "# del video\n",
        "pred_masks_per_frame = rearrange(torch.stack(pred_masks_per_query), 'q t 1 h w -> t q h w').numpy()\n",
        "masked_video = []\n",
        "for vid_frame, frame_masks in tqdm(zip(video_np, pred_masks_per_frame), total=len(video_np), desc='applying masks...'):\n",
        "  # apply the masks:\n",
        "  for inst_mask, color in zip(frame_masks, colors):\n",
        "    vid_frame = apply_mask(vid_frame, inst_mask, color / 255.0)\n",
        "  vid_frame = Image.fromarray((vid_frame * 255).astype(np.uint8))\n",
        "  # visualize the text queries:\n",
        "  vid_frame = ImageOps.expand(vid_frame, border=(0, len(text_queries)*text_border_height_per_query, 0, 0))\n",
        "  W, H = vid_frame.size\n",
        "  draw = ImageDraw.Draw(vid_frame)\n",
        "  font = ImageFont.truetype(font='LiberationSans-Regular.ttf', size=30)\n",
        "  for i, (text_query, color) in enumerate(zip(text_queries, colors), start=1):\n",
        "      w, h = draw.textsize(text_query, font=font)\n",
        "      draw.text(((W - w) / 2, (text_border_height_per_query * i) - h - 3),\n",
        "                text_query, fill=tuple(color) + (255,), font=font)\n",
        "  masked_video.append(np.array(vid_frame))\n",
        "\n",
        "# generate and save the output clip:\n",
        "output_clip_path = 'output_clip.mp4'\n",
        "clip = ImageSequenceClip(sequence=masked_video, fps=meta['video_fps'])\n",
        "clip = clip.set_audio(AudioFileClip(input_clip_path))\n",
        "clip.write_videofile(output_clip_path, fps=meta['video_fps'], audio=True)\n",
        "del masked_video\n",
        "\n",
        "# visualize the output clip:\n",
        "output_clip = open(output_clip_path,'rb').read()\n",
        "data_url = \"data:video/mp4;base64,\" + b64encode(output_clip).decode()\n",
        "HTML(\"\"\"<video width=720 controls><source src=\"%s\" type=\"video/mp4\"></video>\"\"\" % data_url)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}