colab_worker.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/cyrildiagne/instagram-3d-photo/blob/master/colab_worker.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DWBsYMM4vH3N",
        "colab_type": "text"
      },
      "source": [
        "Restart the Runtime after running this cell."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sLSEMdPrhlSw",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install google-cloud-pubsub google-cloud-storage"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "72_LnuZbu9Jg",
        "colab_type": "text"
      },
      "source": [
        "Upload the insta3d-colab-key.json that was generated by the script `script/create_colab_key.sh` before running this cell."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wnpniHuXyDwp",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "06c1fb84-4999-4a50-bdc5-eb1104278810"
      },
      "source": [
        "# TODO: Upload the colab serviceaccount key file.\n",
        "%env GOOGLE_APPLICATION_CREDENTIALS=/content/insta3d-colab-key.json\n",
        "\n",
        "PROJECT_ID = '' # TODO: Set your GCP Project ID.\n",
        "BUCKET_NAME = '' # TODO: Set your GCP Bucket Name."
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "env: GOOGLE_APPLICATION_CREDENTIALS=/content/insta3d-colab-key.json\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M_zTQBX8BgcK",
        "colab_type": "text"
      },
      "source": [
        "# Initialize 3d-photo-inpainting"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AywhpgqvBdi0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Install dependencies\n",
        "!pip3 install torch==1.4.0+cu100 torchvision==0.5.0+cu100 -f https://download.pytorch.org/whl/torch_stable.html\n",
        "!pip3 install opencv-python==4.2.0.32\n",
        "!pip3 install vispy==0.6.4\n",
        "!pip3 install moviepy==1.0.2\n",
        "!pip3 install transforms3d==0.3.1\n",
        "!pip3 install networkx==2.3\n",
        "!sudo apt install sed\n",
        "\n",
        "# Get code from custom fork & retrieve checkpoints.\n",
        "%cd /content/\n",
        "!git clone https://github.com/cyrildiagne/3d-photo-inpainting.git\n",
        "%cd 3d-photo-inpainting\n",
        "!sh download.sh"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TzUDswfpBeTB",
        "colab_type": "text"
      },
      "source": [
        "# Start Worker"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6XApU3MyNV4n",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "bea33b01-2011-476d-99cf-97605b01be76"
      },
      "source": [
        "%cd /content/3d-photo-inpainting\n",
        "\n",
        "import os\n",
        "import time\n",
        "import json\n",
        "import requests\n",
        "from google.cloud import pubsub\n",
        "from google.cloud import storage\n",
        "from IPython.display import clear_output\n",
        "\n",
        "subscriber = pubsub.SubscriberClient()\n",
        "subscription_path = subscriber.subscription_path(PROJECT_ID, 'insta3d-sub')\n",
        "\n",
        "gcp = storage.Client()\n",
        "bucket = gcp.bucket(BUCKET_NAME)\n",
        "\n",
        "output_folder = '/content/3d-photo-inpainting'\n",
        "image_file = output_folder + '/image/img.jpg'\n",
        "depth_file = output_folder + '/depth/img.npy'\n",
        "mesh_file = output_folder + '/mesh/img.ply'\n",
        "video_file = output_folder + '/video/img_circle.mp4'\n",
        "\n",
        "save_extra = False\n",
        "\n",
        "if os.path.exists(output_folder + '/image/moon.jpg'):\n",
        "  os.remove(output_folder + '/image/moon.jpg')\n",
        "\n",
        "def process(img_url, output_path):\n",
        "  # Downlaod the image.\n",
        "  r = requests.get(img_url)\n",
        "  with open(image_file, 'wb') as f:\n",
        "    f.write(r.content)\n",
        "  # Process the image.\n",
        "  !python main.py --config argument.yml\n",
        "  # Upload results to storage.\n",
        "  bucket.blob(output_path).upload_from_filename(video_file)\n",
        "  if save_extra is True:\n",
        "    bucket.blob(output_path[:-4] + '.jpg').upload_from_filename(image_file)\n",
        "    bucket.blob(output_path[:-4] + '.npy').upload_from_filename(depth_file)\n",
        "    bucket.blob(output_path[:-4] + '.ply').upload_from_filename(mesh_file)\n",
        "  # Cleanup tmp files\n",
        "  os.remove(depth_file)\n",
        "  os.remove(image_file)\n",
        "  os.remove(video_file)\n",
        "  os.remove(mesh_file)\n",
        "\n",
        "while True:\n",
        "  response = subscriber.pull(subscription_path, max_messages=1)\n",
        "  if len(response.received_messages):\n",
        "    for msg in response.received_messages:\n",
        "      print('Received message.')\n",
        "      # Load info.\n",
        "      msg_data = json.loads(msg.message.data.decode('utf-8'))\n",
        "      # Check if it's not already being processed.\n",
        "      data_file = 'insta3d/renders/' + msg_data['id'] + '/info.json'\n",
        "      data = json.loads(bucket.blob(data_file).download_as_string())\n",
        "      print(msg_data['url'])\n",
        "      # Acknowledge message.\n",
        "      subscriber.acknowledge(subscription_path, [msg.ack_id])\n",
        "      try:\n",
        "        # Mark as processing in the status file.\n",
        "        status = { 'id': data['id'], 'status': 'processing' }\n",
        "        b = bucket.blob(data_file)\n",
        "        b.cache_control = 'no-cache, max-age=0'\n",
        "        b.upload_from_string(json.dumps(status))\n",
        "        # Process.\n",
        "        img_url = msg_data['url']\n",
        "        output_file = msg_data['output_file']\n",
        "        process(img_url, output_file)\n",
        "        # Mark as ready in the status file.\n",
        "        status = {\n",
        "            'id': data['id'],\n",
        "            'status': 'ready',\n",
        "            'url': f'https://storage.googleapis.com/{BUCKET_NAME}/{output_file}'\n",
        "        }\n",
        "        b = bucket.blob(data_file)\n",
        "        b.cache_control = 'no-cache, max-age=0'\n",
        "        b.upload_from_string(json.dumps(status))\n",
        "        print('Done')\n",
        "      except Exception as e:\n",
        "        print(e)\n",
        "        # Mark as error.\n",
        "        status = { 'id': data['id'], 'status': 'error', 'error': str(e) }\n",
        "        b = bucket.blob(data_file)\n",
        "        b.cache_control = 'no-cache, max-age=0'\n",
        "        b.upload_from_string(json.dumps(status))\n",
        "  else:\n",
        "    clear_output()\n",
        "    print('Queue is empty. Waiting for new message...')\n",
        "    time.sleep(1) # wait 1s before checking again"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Queue is empty. Waiting for new message...\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "Insta3D Worker",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "authorship_tag": "ABX9TyPrkN2IAaP/FtSuf3BmehgY",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}