SubMe.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "SubMe.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [
        "KY40nikm321k"
      ],
      "toc_visible": true,
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ChanJianHao/SubMe/blob/master/SubMe.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uV-cMxksZjct"
      },
      "source": [
        "# SubMe\n",
        "\n",
        "Github URL: [SubMe](https://github.com/HandsomeWJ/SubMe/blob/master/SubMe.ipynb)\n",
        "\n",
        "Cloud-based machine translation/subtiling for your media files! No more language barriers when you have got SubMe!\n",
        "\n",
        "\n",
        "**Instructions** \n",
        "1. Review 'Project Setup' settings - we recommend to leave values as default if you do not understand what they do.\n",
        "1. Choose your preferred processing method under `processMethod` and then press `Ctrl+F9` to run all. \n",
        "\n",
        "\n",
        "# Types of `processMethod`\n",
        "\n",
        "1. '**Upload from Local**' - Upload your video file from your computer.\n",
        "\n",
        "1. '**Upload from Google Drive**' - you will first have to authenticate your Google Drive account. You will then be prompted to enter the desired number of files to sub. Lastly, you will have to enter the paths for each of these files. The paths can be obtained by right-clicking on the file name from the panel on the left.\n",
        "\n",
        "1. '**Download with magnet**' - you will be prompted to enter a magnet link for every torrent you wish to download. Once you have entered all your magnet links, type `exit` and the downloads will begin.\n",
        "\n",
        "1. '**Download with custom command**' For advanced users only, use custom command to such as curl to download your media files for processing.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i5_78TOBG8Ly"
      },
      "source": [
        "# Project Setup\n",
        "\n",
        "It may take a few minutes to load the program. Please be patient."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gzYGE9j_G8L7",
        "cellView": "form"
      },
      "source": [
        "#@title Project Settings\n",
        "\n",
        "import os\n",
        "\n",
        "#@markdown Repo path to clone from.\n",
        "repo_url = 'https://github.com/BingLingGroup/autosub' #@param {type:\"string\"}\n",
        "repo_name = repo_url.split('/')[-1]\n",
        "\n",
        "#@markdown Repo branch or tag to clone. Leave blank to clone `master`.\n",
        "repo_branch = 'dev' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Root workspace.\n",
        "root_dir = '/content/subme' #@param {type:\"string\"}\n",
        "repo_dir = os.path.join(root_dir, repo_name)\n",
        "\n",
        "#@markdown Attempt to request Google Drive access. Will mount to `/content/drive/My Drive`. \n",
        "#@markdown Will save output to your linked Google Drive account inside `subme` folder. \n",
        "\n",
        "#@markdown Attempt to invoke high memory (25.5 GB) in Colab.\n",
        "high_memory = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown Choose `Upload from Local` if you are uploading a file from your own computer. \n",
        "#@markdown Choose `Upload from Google Drive` if you are uploading a file from Google Drive.\n",
        "#@markdown Choose `Download` to download a torrent file from a magnet.\n",
        "\n",
        "processMethod = \"Upload from Google Drive\" #@param [\"Upload from Local\", \"Upload from Google Drive\", \"Download with magnet\", \"Download with custom command\"]\n",
        "\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gMAAjHqr2GgZ",
        "cellView": "form"
      },
      "source": [
        "#@title Perform Google authentication\n",
        "#@markdown Only necessary for \"Upload from Google Drive\" or \"Download with magnet\"\n",
        "\n",
        "from google.colab import drive\n",
        "\n",
        "if processMethod == \"Upload from Google Drive\" or processMethod == \"Download with magnet\":\n",
        "  drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KY40nikm321k"
      },
      "source": [
        "# Additional Setup\n",
        "\n",
        "Ignore this, for advanced users only.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8ljpIzKi4AmO"
      },
      "source": [
        "#@title System Setup\n",
        "#@markdown Attempt to invoke high memory if enabled in settings.\n",
        "\n",
        "if high_memory:\n",
        "  !pip install gputil\n",
        "\n",
        "  # Import packages\n",
        "  import os,sys,humanize,psutil,GPUtil\n",
        "\n",
        "  def mem_report():\n",
        "    print(\"CPU RAM Free: \" + humanize.naturalsize( psutil.virtual_memory().available ))\n",
        "    \n",
        "    GPUs = GPUtil.getGPUs()\n",
        "    for i, gpu in enumerate(GPUs):\n",
        "      print('GPU {:d} ... Mem Free: {:.0f}MB / {:.0f}MB | Utilization {:3.0f}%'.format(i, gpu.memoryFree, gpu.memoryTotal, gpu.memoryUtil*100))\n",
        "      \n",
        "  mem_report()\n",
        "\n",
        "  if psutil.virtual_memory().available < 13958643712:\n",
        "    print('Attempting to invoke high memory.')\n",
        "    print('This notebook will crash intentionally and Colab should display a prompt to offer you high-RAM.')\n",
        "    print('IF THIS PROMPT DOES NOT SHOW, DISABLE high_memory IN THE SETTINGS!')\n",
        "    d=[]\n",
        "    while(1):\n",
        "      d.append('1')\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "68tEWSj6G8MB"
      },
      "source": [
        "#@title Project Setup\n",
        "#@markdown Setup directories. Clone project from git.\n",
        "\n",
        "output_dir = ('/content/drive/My Drive/subme')\n",
        "\n",
        "os.makedirs(root_dir, exist_ok=True)\n",
        "\n",
        "if not os.path.isdir(repo_dir):\n",
        "  repo_cmd = f'--branch {repo_branch}' if repo_branch else ''\n",
        "  !git clone {repo_url} {repo_cmd} --depth 1 \"{repo_dir}\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LzBNny4yG8MH"
      },
      "source": [
        "#@title Dependency Setup\n",
        "#@markdown Install dependencies from `requirements.txt`.\n",
        "\n",
        "!apt install ffmpeg python3 curl git -y\n",
        "!curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py\n",
        "!python3 get-pip.py\n",
        "!pip install git+{repo_url}@{repo_branch} ffmpeg-normalize langcodes"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1MQIROIBG8MM"
      },
      "source": [
        "# Gathering Videos for Processing"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JHwGXjEzG8MP",
        "cellView": "form"
      },
      "source": [
        "#@#title Upload Videos for Processing\n",
        "#@markdown When uploading from Local, click on 'Choose File' and navigate to the\n",
        "#@markdown the directory on your computer that contains the file you wish to upload.\n",
        "#@markdown You can select multiple files.\n",
        "\n",
        "#@markdown When uploading from Google Drive, you will first have to specify how many\n",
        "#@markdown files you wish to upload. Then you will have to enter the complete file path\n",
        "#@markdown for each of these files. The file path can be obtained by right-clicking on\n",
        "#@markdown the file name from the file browser panel on the left and selecting `Copy Path`.\n",
        "\n",
        "#@markdown When downloading from magnet links, please make sure that only magnet links \n",
        "#@markdown are entered. Magnet links are links which begin with `magnet:`.\n",
        "\n",
        "#@markdown If you see an error in this section, then run this cell again. \n",
        "#@markdown Ensure you have picked the correct processing method from settings above.\n",
        "#@markdown After successful execution, run the next cell using CTRL+F10 or by \n",
        "#@markdown clicking the play button below.\n",
        "\n",
        "from google.colab import files\n",
        "import shutil\n",
        "import re\n",
        "\n",
        "if processMethod == \"Upload from Local\":\n",
        "  mediaList = files.upload()\n",
        "  mediaList = list(mediaList.keys())\n",
        "  for i in range(len(mediaList)):\n",
        "    file_name = mediaList[i]\n",
        "    mediaList[i] = \"/content/\" + file_name\n",
        "  \n",
        "\n",
        "elif processMethod == \"Upload from Google Drive\":\n",
        "  mediaList = []\n",
        "  file_no = int(input(\"No. of files to sub: \"))\n",
        "  for i in range(file_no):\n",
        "    file_name = input(\"Please enter file path: \")\n",
        "    file_name = re.escape(file_name)\n",
        "    mediaList.append(file_name)\n",
        "\n",
        "elif processMethod == \"Download with magnet\":\n",
        "  mediaList = []\n",
        "  !apt install python3-libtorrent\n",
        "  import libtorrent as lt\n",
        "\n",
        "  ses = lt.session()\n",
        "  ses.listen_on(6881, 6891)\n",
        "  downloads = []\n",
        "\n",
        "  params = {\"save_path\": \"/content/drive/My Drive/Torrent\"}\n",
        "\n",
        "  while True:\n",
        "      magnet_link = input(\"Enter Magnet Link(s) Type exit when done: \")\n",
        "      if magnet_link.lower() == \"exit\":\n",
        "          break\n",
        "      downloads.append(\n",
        "          lt.add_magnet_uri(ses, magnet_link, params)\n",
        "      )\n",
        "  \n",
        "  # Start Download \n",
        "\n",
        "  import time\n",
        "  from IPython.display import display\n",
        "  import ipywidgets as widgets\n",
        "\n",
        "  state_str = [\n",
        "      \"queued\",\n",
        "      \"checking\",\n",
        "      \"downloading metadata\",\n",
        "      \"downloading\",\n",
        "      \"finished\",\n",
        "      \"seeding\",\n",
        "      \"allocating\",\n",
        "      \"checking fastresume\",\n",
        "  ]\n",
        "\n",
        "  layout = widgets.Layout(width=\"auto\")\n",
        "  style = {\"description_width\": \"initial\"}\n",
        "  download_bars = [\n",
        "      widgets.FloatSlider(\n",
        "          step=0.01, disabled=True, layout=layout, style=style\n",
        "      )\n",
        "      for _ in downloads\n",
        "  ]\n",
        "  display(*download_bars)\n",
        "\n",
        "  while downloads:\n",
        "      next_shift = 0\n",
        "      for index, download in enumerate(downloads[:]):\n",
        "          bar = download_bars[index + next_shift]\n",
        "          if not download.is_seed():\n",
        "              s = download.status()\n",
        "\n",
        "              bar.description = \" \".join(\n",
        "                  [\n",
        "                      download.name(),\n",
        "                      str(s.download_rate / 1000),\n",
        "                      \"kB/s\",\n",
        "                      state_str[s.state],\n",
        "                  ]\n",
        "              )\n",
        "              bar.value = s.progress * 100\n",
        "          else:\n",
        "              next_shift -= 1\n",
        "              ses.remove_torrent(download)\n",
        "              downloads.remove(download)\n",
        "              bar.close() # Seems to be not working in Colab (see https://github.com/googlecolab/colabtools/issues/726#issue-486731758)\n",
        "              download_bars.remove(bar)\n",
        "              print(download.name(), \"complete\")\n",
        "              mediaList.append(\"/content/drive/My Drive/Torrent/\" + download.name())\n",
        "      time.sleep(5)\n",
        "  \n",
        "elif processMethod == \"Download with custom command\":\n",
        "  mediaList = []\n",
        "\n",
        "  customCommand = input(\"Please enter command you want to execute: \")\n",
        "  !{customCommand}\n",
        "  \n",
        "\n",
        "  #scan for video files in cwd\n",
        "  for file in os.listdir(\"/content\"):\n",
        "    if (file.endswith(\".mp4\") or file.endswith(\".mkv\") or file.endswith(\".webm\") \n",
        "    or file.endswith(\".avi\") or file.endswith(\".flv\") or file.endswith(\".mov\")\n",
        "    or file.endswith(\".wmv\") or file.endswith(\".ogg\")or file.endswith(\".m4v\")):\n",
        "        print(\"Found video: \" + os.path.join(\"/content\", file))\n",
        "        mediaList.append(os.path.join(\"/content\", file))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kMlAiyxcG8MV"
      },
      "source": [
        "# Video Processing"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0kRgJz7qG8MW",
        "cellView": "form"
      },
      "source": [
        "#@markdown Choose media source and destination language.\n",
        "#@markdown Refer to langcode website [here](http://www.lingoes.net/en/translator/langcode.htm).\n",
        "\n",
        "src_language = 'ja' #@param {type:\"string\"}\n",
        "dst_language = 'en' #@param {type:\"string\"}\n",
        "output_format = 'srt'  #@param [\"srt\", \"ass\", \"sub\", \"json\", \"txt\"]\n",
        "\n",
        "for file in mediaList:\n",
        "  print('User mediaList file \"{name}\" '.format(name=file))\n",
        "\n",
        "  if src_language != dst_language:\n",
        "    !autosub -i {file} -S {src_language} -D {dst_language} -F {output_format} -o {file}.{output_format}\n",
        "  else:\n",
        "    !autosub -i {file} -S {src_language} -F {output_format} -o {file}.{output_format}\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "epWpnwASEcU0"
      },
      "source": [
        "# Download your subtitles"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G03Go4fWEhu5",
        "cellView": "form"
      },
      "source": [
        "#@markdown Download your subtitles to your PC.\n",
        "\n",
        "for file in mediaList:\n",
        "  file = re.sub(r'\\\\(.)', r'\\1', file)\n",
        "  files.download(str(file) + \".\" + dst_language + \".\" + output_format)\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}