inference_tf2_colab.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "inference_tf2_colab.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rOvvWAVTkMR7"
      },
      "source": [
        "# Intro to Object Detection Colab\n",
        "\n",
        "Welcome to the object detection colab!  This demo will take you through the steps of running an \"out-of-the-box\" detection model on a collection of images."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fPJESdJcqB3B"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vPs64QA1Zdov"
      },
      "source": [
        "## Imports and Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "LBZ9VWZZFUCT",
        "outputId": "8e7e30c2-1aab-4eb0-e574-c06e8880bebd"
      },
      "source": [
        "!pip install  tensorflow-gpu==2.2.0"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting tensorflow-gpu==2.2.0\n",
            "  Using cached https://files.pythonhosted.org/packages/ef/e6/6a7be21491b528790e733b4bb3d2707f630fef3ee890a1aebaba1639c1e6/tensorflow_gpu-2.2.0-cp37-cp37m-manylinux2010_x86_64.whl\n",
            "Requirement already satisfied: protobuf>=3.8.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==2.2.0) (3.12.4)\n",
            "Requirement already satisfied: absl-py>=0.7.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==2.2.0) (0.12.0)\n",
            "Collecting gast==0.3.3\n",
            "  Downloading https://files.pythonhosted.org/packages/d6/84/759f5dd23fec8ba71952d97bcc7e2c9d7d63bdc582421f3cd4be845f0c98/gast-0.3.3-py2.py3-none-any.whl\n",
            "Collecting tensorboard<2.3.0,>=2.2.0\n",
            "  Using cached https://files.pythonhosted.org/packages/1d/74/0a6fcb206dcc72a6da9a62dd81784bfdbff5fedb099982861dc2219014fb/tensorboard-2.2.2-py3-none-any.whl\n",
            "Requirement already satisfied: astunparse==1.6.3 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==2.2.0) (1.6.3)\n",
            "Requirement already satisfied: wrapt>=1.11.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==2.2.0) (1.12.1)\n",
            "Requirement already satisfied: google-pasta>=0.1.8 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==2.2.0) (0.2.0)\n",
            "Requirement already satisfied: six>=1.12.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==2.2.0) (1.15.0)\n",
            "Requirement already satisfied: wheel>=0.26; python_version >= \"3\" in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==2.2.0) (0.36.2)\n",
            "Collecting tensorflow-estimator<2.3.0,>=2.2.0\n",
            "  Using cached https://files.pythonhosted.org/packages/a4/f5/926ae53d6a226ec0fda5208e0e581cffed895ccc89e36ba76a8e60895b78/tensorflow_estimator-2.2.0-py2.py3-none-any.whl\n",
            "Collecting h5py<2.11.0,>=2.10.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/3f/c0/abde58b837e066bca19a3f7332d9d0493521d7dd6b48248451a9e3fe2214/h5py-2.10.0-cp37-cp37m-manylinux1_x86_64.whl (2.9MB)\n",
            "\u001b[K     |████████████████████████████████| 2.9MB 6.7MB/s \n",
            "\u001b[?25hRequirement already satisfied: keras-preprocessing>=1.1.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==2.2.0) (1.1.2)\n",
            "Requirement already satisfied: numpy<2.0,>=1.16.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==2.2.0) (1.18.5)\n",
            "Requirement already satisfied: scipy==1.4.1; python_version >= \"3\" in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==2.2.0) (1.4.1)\n",
            "Requirement already satisfied: grpcio>=1.8.6 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==2.2.0) (1.32.0)\n",
            "Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==2.2.0) (1.1.0)\n",
            "Requirement already satisfied: opt-einsum>=2.3.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==2.2.0) (3.3.0)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.7/dist-packages (from protobuf>=3.8.0->tensorflow-gpu==2.2.0) (56.1.0)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow-gpu==2.2.0) (3.3.4)\n",
            "Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow-gpu==2.2.0) (2.25.1)\n",
            "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow-gpu==2.2.0) (1.8.0)\n",
            "Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow-gpu==2.2.0) (2.0.0)\n",
            "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow-gpu==2.2.0) (0.4.4)\n",
            "Requirement already satisfied: google-auth<2,>=1.6.3 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow-gpu==2.2.0) (1.30.0)\n",
            "Requirement already satisfied: importlib-metadata; python_version < \"3.8\" in /usr/local/lib/python3.7/dist-packages (from markdown>=2.6.8->tensorboard<2.3.0,>=2.2.0->tensorflow-gpu==2.2.0) (4.0.1)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard<2.3.0,>=2.2.0->tensorflow-gpu==2.2.0) (2020.12.5)\n",
            "Requirement already satisfied: chardet<5,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard<2.3.0,>=2.2.0->tensorflow-gpu==2.2.0) (3.0.4)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard<2.3.0,>=2.2.0->tensorflow-gpu==2.2.0) (1.24.3)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard<2.3.0,>=2.2.0->tensorflow-gpu==2.2.0) (2.10)\n",
            "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.7/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.3.0,>=2.2.0->tensorflow-gpu==2.2.0) (1.3.0)\n",
            "Requirement already satisfied: rsa<5,>=3.1.4; python_version >= \"3.6\" in /usr/local/lib/python3.7/dist-packages (from google-auth<2,>=1.6.3->tensorboard<2.3.0,>=2.2.0->tensorflow-gpu==2.2.0) (4.7.2)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.7/dist-packages (from google-auth<2,>=1.6.3->tensorboard<2.3.0,>=2.2.0->tensorflow-gpu==2.2.0) (0.2.8)\n",
            "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from google-auth<2,>=1.6.3->tensorboard<2.3.0,>=2.2.0->tensorflow-gpu==2.2.0) (4.2.2)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata; python_version < \"3.8\"->markdown>=2.6.8->tensorboard<2.3.0,>=2.2.0->tensorflow-gpu==2.2.0) (3.4.1)\n",
            "Requirement already satisfied: typing-extensions>=3.6.4; python_version < \"3.8\" in /usr/local/lib/python3.7/dist-packages (from importlib-metadata; python_version < \"3.8\"->markdown>=2.6.8->tensorboard<2.3.0,>=2.2.0->tensorflow-gpu==2.2.0) (3.7.4.3)\n",
            "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.7/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.3.0,>=2.2.0->tensorflow-gpu==2.2.0) (3.1.0)\n",
            "Requirement already satisfied: pyasn1>=0.1.3 in /usr/local/lib/python3.7/dist-packages (from rsa<5,>=3.1.4; python_version >= \"3.6\"->google-auth<2,>=1.6.3->tensorboard<2.3.0,>=2.2.0->tensorflow-gpu==2.2.0) (0.4.8)\n",
            "\u001b[31mERROR: tensorflow 2.5.0 has requirement gast==0.4.0, but you'll have gast 0.3.3 which is incompatible.\u001b[0m\n",
            "\u001b[31mERROR: tensorflow 2.5.0 has requirement grpcio~=1.34.0, but you'll have grpcio 1.32.0 which is incompatible.\u001b[0m\n",
            "\u001b[31mERROR: tensorflow 2.5.0 has requirement h5py~=3.1.0, but you'll have h5py 2.10.0 which is incompatible.\u001b[0m\n",
            "\u001b[31mERROR: tensorflow 2.5.0 has requirement numpy~=1.19.2, but you'll have numpy 1.18.5 which is incompatible.\u001b[0m\n",
            "\u001b[31mERROR: tensorflow 2.5.0 has requirement tensorboard~=2.5, but you'll have tensorboard 2.2.2 which is incompatible.\u001b[0m\n",
            "\u001b[31mERROR: tensorflow 2.5.0 has requirement tensorflow-estimator<2.6.0,>=2.5.0rc0, but you'll have tensorflow-estimator 2.2.0 which is incompatible.\u001b[0m\n",
            "Installing collected packages: gast, tensorboard, tensorflow-estimator, h5py, tensorflow-gpu\n",
            "  Found existing installation: gast 0.4.0\n",
            "    Uninstalling gast-0.4.0:\n",
            "      Successfully uninstalled gast-0.4.0\n",
            "  Found existing installation: tensorboard 2.5.0\n",
            "    Uninstalling tensorboard-2.5.0:\n",
            "      Successfully uninstalled tensorboard-2.5.0\n",
            "  Found existing installation: tensorflow-estimator 2.5.0\n",
            "    Uninstalling tensorflow-estimator-2.5.0:\n",
            "      Successfully uninstalled tensorflow-estimator-2.5.0\n",
            "  Found existing installation: h5py 3.1.0\n",
            "    Uninstalling h5py-3.1.0:\n",
            "      Successfully uninstalled h5py-3.1.0\n",
            "  Found existing installation: tensorflow-gpu 2.3.0\n",
            "    Uninstalling tensorflow-gpu-2.3.0:\n",
            "      Successfully uninstalled tensorflow-gpu-2.3.0\n",
            "Successfully installed gast-0.3.3 h5py-2.10.0 tensorboard-2.2.2 tensorflow-estimator-2.2.0 tensorflow-gpu-2.2.0\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.colab-display-data+json": {
              "pip_warning": {
                "packages": [
                  "gast",
                  "h5py",
                  "tensorboard",
                  "tensorflow"
                ]
              }
            }
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 52
        },
        "id": "S6-TjA2eQ2VP",
        "outputId": "86f247c4-40a8-4845-8cb8-42ca809bf785"
      },
      "source": [
        "import tensorflow as tf \n",
        "print(tf.__version__)\n",
        "tf.test.gpu_device_name()"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "2.2.0\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'/device:GPU:0'"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 1
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3Tdsq1p-rhTk"
      },
      "source": [
        "# # import tensorflow as tf \n",
        "# resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')\n",
        "# tf.config.experimental_connect_to_cluster(resolver)\n",
        "# # This is the TPU initialization code that has to be at the beginning.\n",
        "# tf.tpu.experimental.initialize_tpu_system(resolver)\n",
        "# print(\"All devices: \", tf.config.list_logical_devices('TPU'))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oi28cqGGFWnY",
        "outputId": "e3a79516-afb7-4264-b88a-63254609e910"
      },
      "source": [
        "import os\n",
        "import pathlib\n",
        "\n",
        "# Clone the tensorflow models repository if it doesn't already exist\n",
        "if \"models\" in pathlib.Path.cwd().parts:\n",
        "  while \"models\" in pathlib.Path.cwd().parts:\n",
        "    os.chdir('..')\n",
        "elif not pathlib.Path('models').exists():\n",
        "  !git clone --depth 1 https://github.com/tensorflow/models"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Cloning into 'models'...\n",
            "remote: Enumerating objects: 2650, done.\u001b[K\n",
            "remote: Counting objects: 100% (2650/2650), done.\u001b[K\n",
            "remote: Compressing objects: 100% (2203/2203), done.\u001b[K\n",
            "remote: Total 2650 (delta 671), reused 1299 (delta 415), pack-reused 0\u001b[K\n",
            "Receiving objects: 100% (2650/2650), 32.62 MiB | 28.58 MiB/s, done.\n",
            "Resolving deltas: 100% (671/671), done.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NwdsBdGhFanc"
      },
      "source": [
        "# Install the Object Detection API\n",
        "%%bash\n",
        "cd models/research/\n",
        "protoc object_detection/protos/*.proto --python_out=.\n",
        "cp object_detection/packages/tf2/setup.py .\n",
        "python -m pip install ."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yn5_uV1HLvaz"
      },
      "source": [
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import io\n",
        "import scipy.misc\n",
        "import numpy as np\n",
        "from six import BytesIO\n",
        "from PIL import Image, ImageDraw, ImageFont\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "from object_detection.utils import label_map_util\n",
        "from object_detection.utils import config_util\n",
        "from object_detection.utils import visualization_utils as viz_utils\n",
        "from object_detection.builders import model_builder\n",
        "\n",
        "%matplotlib inline"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "PfvrElppq_91",
        "outputId": "3bffdfd5-b6d4-4f13-cb38-0b44a664293f"
      },
      "source": [
        "tf.test.gpu_device_name()"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'/device:GPU:0'"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IogyryF2lFBL"
      },
      "source": [
        "## Utilities"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-y9R0Xllefec"
      },
      "source": [
        "def load_image_into_numpy_array(path):\n",
        "  \"\"\"Load an image from file into a numpy array.\n",
        "\n",
        "  Puts image into numpy array to feed into tensorflow graph.\n",
        "  Note that by convention we put it into a numpy array with shape\n",
        "  (height, width, channels), where channels=3 for RGB.\n",
        "\n",
        "  Args:\n",
        "    path: the file path to the image\n",
        "\n",
        "  Returns:\n",
        "    uint8 numpy array with shape (img_height, img_width, 3)\n",
        "  \"\"\"\n",
        "  img_data = tf.io.gfile.GFile(path, 'rb').read()\n",
        "  image = Image.open(BytesIO(img_data))\n",
        "  (im_width, im_height) = image.size\n",
        "  return np.array(image.getdata()).reshape(\n",
        "      (im_height, im_width, 3)).astype(np.uint8)\n",
        "\n",
        "def get_keypoint_tuples(eval_config):\n",
        "  \"\"\"Return a tuple list of keypoint edges from the eval config.\n",
        "  \n",
        "  Args:\n",
        "    eval_config: an eval config containing the keypoint edges\n",
        "  \n",
        "  Returns:\n",
        "    a list of edge tuples, each in the format (start, end)\n",
        "  \"\"\"\n",
        "  tuple_list = []\n",
        "  kp_list = eval_config.keypoint_edge\n",
        "  for edge in kp_list:\n",
        "    tuple_list.append((edge.start, edge.end))\n",
        "  return tuple_list"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "R4YjnOjME1gy"
      },
      "source": [
        "# @title Choose the model to use, then evaluate the cell.\n",
        "MODELS = {'centernet_with_keypoints': 'centernet_hg104_512x512_kpts_coco17_tpu-32', 'centernet_without_keypoints': 'centernet_hg104_512x512_coco17_tpu-8'}\n",
        "\n",
        "model_display_name = 'centernet_without_keypoints' # @param ['centernet_with_keypoints', 'centernet_without_keypoints']\n",
        "model_name = MODELS[model_display_name]"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6917xnUSlp9x"
      },
      "source": [
        "### Build a detection model and load pre-trained model weights\n",
        "\n",
        "This sometimes takes a little while, please be patient!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ctPavqlyPuU_",
        "outputId": "48ac2f75-ce56-4379-cef3-bf549265e41c"
      },
      "source": [
        "# Download the checkpoint and put it into models/research/object_detection/test_data/\n",
        "\n",
        "if model_display_name == 'centernet_with_keypoints':\n",
        "  !wget http://download.tensorflow.org/models/object_detection/tf2/20200711/centernet_hg104_512x512_kpts_coco17_tpu-32.tar.gz\n",
        "  !tar -xf centernet_hg104_512x512_kpts_coco17_tpu-32.tar.gz\n",
        "  !mv centernet_hg104_512x512_kpts_coco17_tpu-32/checkpoint models/research/object_detection/test_data/\n",
        "else:\n",
        "  !wget http://download.tensorflow.org/models/object_detection/tf2/20200711/centernet_hg104_512x512_coco17_tpu-8.tar.gz\n",
        "  !tar -xf centernet_hg104_512x512_coco17_tpu-8.tar.gz\n",
        "  !mv centernet_hg104_512x512_coco17_tpu-8/checkpoint models/research/object_detection/test_data/"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "--2021-05-24 06:09:13--  http://download.tensorflow.org/models/object_detection/tf2/20200711/centernet_hg104_512x512_coco17_tpu-8.tar.gz\n",
            "Resolving download.tensorflow.org (download.tensorflow.org)... 74.125.197.128, 2607:f8b0:400e:c03::80\n",
            "Connecting to download.tensorflow.org (download.tensorflow.org)|74.125.197.128|:80... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 1426100886 (1.3G) [application/x-tar]\n",
            "Saving to: ‘centernet_hg104_512x512_coco17_tpu-8.tar.gz’\n",
            "\n",
            "centernet_hg104_512 100%[===================>]   1.33G   161MB/s    in 8.5s    \n",
            "\n",
            "2021-05-24 06:09:22 (161 MB/s) - ‘centernet_hg104_512x512_coco17_tpu-8.tar.gz’ saved [1426100886/1426100886]\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4cni4SSocvP_"
      },
      "source": [
        "import os \n",
        "pipeline_config = os.path.join('models/research/object_detection/configs/tf2/',\n",
        "                                'centernet_hourglass104_512x512_coco17_tpu-8' + '.config')\n",
        "model_dir = 'models/research/object_detection/test_data/checkpoint/'\n",
        "\n",
        "# Load pipeline config and build a detection model\n",
        "configs = config_util.get_configs_from_pipeline_file(pipeline_config)\n",
        "model_config = configs['model']\n",
        "detection_model = model_builder.build(\n",
        "      model_config=model_config, is_training=False)\n",
        "\n",
        "# Restore checkpoint\n",
        "ckpt = tf.compat.v2.train.Checkpoint(\n",
        "      model=detection_model)\n",
        "ckpt.restore(os.path.join(model_dir, 'ckpt-0')).expect_partial()\n",
        "\n",
        "def get_model_detection_function(model):\n",
        "  \"\"\"Get a tf.function for detection.\"\"\"\n",
        "\n",
        "  @tf.function\n",
        "  def detect_fn(image):\n",
        "    \"\"\"Detect objects in image.\"\"\"\n",
        "\n",
        "    image, shapes = model.preprocess(image)\n",
        "    prediction_dict = model.predict(image, shapes)\n",
        "    detections = model.postprocess(prediction_dict, shapes)\n",
        "\n",
        "    return detections, prediction_dict, tf.reshape(shapes, [-1])\n",
        "\n",
        "  return detect_fn\n",
        "\n",
        "detect_fn = get_model_detection_function(detection_model)"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NKtD0IeclbL5"
      },
      "source": [
        "# Load label map data (for plotting).\n",
        "\n",
        "Label maps correspond index numbers to category names, so that when our convolution network predicts `5`, we know that this corresponds to `airplane`.  Here we use internal utility functions, but anything that returns a dictionary mapping integers to appropriate string labels would be fine."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5mucYUS6exUJ"
      },
      "source": [
        "label_map_path = '/content/mscoco_label_map.pbtxt'#configs['eval_input_config'].label_map_path\n",
        "label_map = label_map_util.load_labelmap(label_map_path)\n",
        "categories = label_map_util.convert_label_map_to_categories(\n",
        "    label_map,\n",
        "    max_num_classes=label_map_util.get_max_label_map_index(label_map),\n",
        "    use_display_name=True)\n",
        "category_index = label_map_util.create_category_index(categories)\n",
        "label_map_dict = label_map_util.get_label_map_dict(label_map, use_display_name=True)"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RLusV1o-mAx8"
      },
      "source": [
        "### Putting everything together!\n",
        "\n",
        "Run the below code which loads an image, runs it through the detection model and visualizes the detection results, including the keypoints.\n",
        "\n",
        "Note that this will take a long time (several minutes) the first time you run this code due to tf.function's trace-compilation --- on subsequent runs (e.g. on new images), things will be faster.\n",
        "\n",
        "Here are some simple things to try out if you are curious:\n",
        "* Try running inference on your own images (local paths work)\n",
        "* Modify some of the input images and see if detection still works.  Some simple things to try out here (just uncomment the relevant portions of code) include flipping the image horizontally, or converting to grayscale (note that we still expect the input image to have 3 channels).\n",
        "* Print out `detections['detection_boxes']` and try to match the box locations to the boxes in the image.  Notice that coordinates are given in normalized form (i.e., in the interval [0, 1]).\n",
        "* Set min_score_thresh to other values (between 0 and 1) to allow more detections in or to filter out more detections.\n",
        "\n",
        "Note that you can run this cell repeatedly without rerunning earlier cells.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vr_Fux-gfaG9"
      },
      "source": [
        "image_dir = '/content/test_images'\n",
        "# for i in file\n",
        "image_path = os.path.join(image_dir, '1786.jpg')\n",
        "image_np = load_image_into_numpy_array(image_path)\n",
        "\n",
        "# Things to try:\n",
        "# Flip horizontally\n",
        "# image_np = np.fliplr(image_np).copy()\n",
        "\n",
        "# Convert image to grayscale\n",
        "# image_np = np.tile(\n",
        "#     np.mean(image_np, 2, keepdims=True), (1, 1, 3)).astype(np.uint8)\n",
        "\n",
        "input_tensor = tf.convert_to_tensor(\n",
        "    np.expand_dims(image_np, 0), dtype=tf.float32)\n",
        "detections, predictions_dict, shapes = detect_fn(input_tensor)\n",
        "\n",
        "label_id_offset = 1\n",
        "image_np_with_detections = image_np.copy()\n",
        "\n",
        "# Use keypoints if available in detections\n",
        "keypoints, keypoint_scores = None, None\n",
        "if 'detection_keypoints' in detections:\n",
        "  keypoints = detections['detection_keypoints'][0].numpy()\n",
        "  keypoint_scores = detections['detection_keypoint_scores'][0].numpy()\n",
        "\n",
        "viz_utils.visualize_boxes_and_labels_on_image_array(\n",
        "      image_np_with_detections,\n",
        "      detections['detection_boxes'][0].numpy(),\n",
        "      (detections['detection_classes'][0].numpy() + label_id_offset).astype(int),\n",
        "      detections['detection_scores'][0].numpy(),\n",
        "      category_index,\n",
        "      use_normalized_coordinates=True,\n",
        "      max_boxes_to_draw=200,\n",
        "      min_score_thresh=.30,\n",
        "      agnostic_mode=False,\n",
        "      keypoints=keypoints,\n",
        "      keypoint_scores=keypoint_scores,\n",
        "      keypoint_edges=get_keypoint_tuples(configs['eval_config']))\n",
        "\n",
        "plt.figure(figsize=(12,16))\n",
        "plt.imshow(image_np_with_detections)\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-gsqzTQBYDEQ"
      },
      "source": [
        "image_np.shape"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TZykaC6NUBSG"
      },
      "source": [
        "# keypoints = detections['detection_keypoints'][0].numpy()\n",
        "# keypoints\n",
        "def get_results(image_np, frame_no):\n",
        "  input_tensor = tf.convert_to_tensor(\n",
        "  np.expand_dims(image_np, 0), dtype=tf.float32)\n",
        "  detections, predictions_dict, shapes = detect_fn(input_tensor)\n",
        "  im_height,im_width  = image_np.shape[0],image_np.shape[1]\n",
        "  boxes = detections['detection_boxes'][0].numpy()\n",
        "  # detections['detection_classes'][0].numpy()+1\n",
        "  score = detections['detection_scores'][0].numpy()\n",
        "  cls = detections['detection_classes'][0].numpy() + 1\n",
        "  # print(im_width, im_height)\n",
        "  \n",
        "  # img = image_np.copy()\n",
        "  import cv2 \n",
        "  lst =[]\n",
        "  for j,i in enumerate(score):\n",
        "    if i> 0.3:\n",
        "      # if cls[j] not in [37,39]:\n",
        "      #   continue\n",
        "      # print(cls[j], i)\n",
        "      ymin, xmin, ymax, xmax= boxes[j].tolist()\n",
        "      x1,x2,y1,y2  = (xmin * im_width, xmax * im_width,\n",
        "                                    ymin * im_height, ymax * im_height)\n",
        "      x1,y1,x2,y2 = int(x1),int(y1), int(x2), int(y2)\n",
        "      lst.append([x1,y1,x2,y2,cls[j],frame_no])\n",
        "      cv2.rectangle(image_np, (x1,y1), (x2,y2), (0,255,0),2)\n",
        "  if frame_no%100==0:\n",
        "    print(frame_no, len(lst))\n",
        "  \n",
        "      # x,y,\n",
        "      # print(x1,x2,y1,y2)\n",
        "      \n",
        "  return lst, image_np\n",
        "  # cv2.imwrite('output.jpg',img )\n",
        "# plt.figure(figsize=(400,800))\n",
        "# plt.imshow(img)\n"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mKkKkRXv3GRE",
        "outputId": "b508156e-fba6-4dd4-da07-544e09a8a7f7"
      },
      "source": [
        "with open('files.txt','r') as f:\n",
        "  content= f.readlines()\n",
        "print(len(content))\n",
        "content =[int(i[:-5]) for  i in content]\n",
        "content[:10]"
      ],
      "execution_count": 31,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "348\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[1222, 1224, 1226, 1228, 1238, 1242, 1244, 1246, 1306, 1310]"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 31
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NDON9BpkT10D"
      },
      "source": [
        "import cv2\n",
        "\n",
        "video = cv2.VideoCapture(\"/content/drive/MyDrive/cmatch.mp4\");\n",
        "# video.set(cv2.CAP_PROP_FPS, int(20))\n",
        "\n",
        "fourcc = cv2.VideoWriter_fourcc(*'XVID')\n",
        "out = cv2.VideoWriter('results.mp4',fourcc, 20, (1280,720))\n",
        "f=open('tf_output.json', 'a+')\n",
        "\n",
        "\n",
        "# if __name__ == '__main__':\n",
        "    # print(\"Frame rate : {0}\".format(video.get(cv2.CAP_PROP_FPS)))\n",
        "k=0\n",
        "cor =[]\n",
        "while video.isOpened():\n",
        "    ret1, frame2 = video.read()\n",
        "    # cv2.imshow(\"Changed\", frame2)\n",
        "\n",
        "    \n",
        "    if ret1:\n",
        "      \n",
        "        # b = cv2.resize(frame2,(1280,720),fx=0,fy=0, interpolation = cv2.INTER_CUBIC)\n",
        "        # try:\n",
        "        # if k in content:\n",
        "        lst,frame =get_results(frame2,k)\n",
        "        if len(cor)==0 and len(lst)>0:\n",
        "          cor = lst \n",
        "        elif len(lst)>0:\n",
        "          # if len(lst)\n",
        "          cor+=lst\n",
        "        \n",
        "        f.write(str(lst))\n",
        "        print(k, len(lst))\n",
        "        # print(len(lst), k)\n",
        "        k+=1\n",
        "        if k >3000:\n",
        "          break\n",
        "        # cv2.imshow(\"Changed\", b)\n",
        "        # if cv2.waitKey(30) & 0xFF == ord('q') or k>10:  # press q to quit\n",
        "        #     break\n",
        "        b = cv2.resize(frame,(1280,720),fx=0,fy=0, interpolation = cv2.INTER_CUBIC)\n",
        "        out.write(b)\n",
        "    else:\n",
        "        break\n",
        "\n",
        "video.release()\n",
        "out.release()\n",
        "# import pandas as pd\n",
        "# pd.to_csv(cor).to_csv('/content/drive/MyDrive/tf_ou.csv')\n",
        "# cv2.destroyAllWindows()"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3gTgR5GXmphY"
      },
      "source": [
        "import pandas as pd \n",
        "# pd.DataFrame(cor).to_csv('/content/drive/MyDrive/tf_output.csv')"
      ],
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sBX56JVNcYsl",
        "outputId": "c81d1e4b-87ea-4a7c-95f5-50c31014df2a"
      },
      "source": [
        "data = pd.DataFrame(cor)\n",
        "data.values[1]"
      ],
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([None, None, None, None, None, None, None, None, None, None, None],\n",
              "      dtype=object)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CpuvxXyBdc9z"
      },
      "source": [
        "sum((data.values!=None))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-Z3WqcApdtR7"
      },
      "source": [
        "import cv2\n",
        "\n",
        "video = cv2.VideoCapture(\"/content/drive/MyDrive/cmatch.mp4\");\n",
        "# video.set(cv2.CAP_PROP_FPS, int(20))\n",
        "\n",
        "fourcc = cv2.VideoWriter_fourcc(*'XVID')\n",
        "out = cv2.VideoWriter('results.mp4',fourcc, 20, (1280,720))\n",
        "f=open('tf_output.json', 'a+')\n",
        "\n",
        "\n",
        "# if __name__ == '__main__':\n",
        "    # print(\"Frame rate : {0}\".format(video.get(cv2.CAP_PROP_FPS)))\n",
        "k=0\n",
        "cor =[]\n",
        "while video.isOpened():\n",
        "    ret1, frame2 = video.read()\n",
        "    # cv2.imshow(\"Changed\", frame2)\n",
        "\n",
        "    \n",
        "    if ret1:\n",
        "      \n",
        "        # b = cv2.resize(frame2,(1280,720),fx=0,fy=0, interpolation = cv2.INTER_CUBIC)\n",
        "        # try:\n",
        "        # if k in content:\n",
        "        lst,frame =get_results(frame2,k)\n",
        "        if len(cor)==0:\n",
        "          cor = lst \n",
        "        else:\n",
        "          cor+=lst\n",
        "        \n",
        "        f.write(str(lst))\n",
        "        print(k, len(lst))\n",
        "        # print(len(lst), k)\n",
        "        k+=1\n",
        "        if k >3000:\n",
        "          break\n",
        "        # cv2.imshow(\"Changed\", b)\n",
        "        # if cv2.waitKey(30) & 0xFF == ord('q') or k>10:  # press q to quit\n",
        "        #     break\n",
        "        b = cv2.resize(frame,(1280,720),fx=0,fy=0, interpolation = cv2.INTER_CUBIC)\n",
        "        out.write(b)\n",
        "    else:\n",
        "        break\n",
        "\n",
        "video.release()\n",
        "out.release()\n",
        "import pandas as pd\n",
        "pd.to_csv(cor).to_csv('/content/drive/MyDrive/tf_output.csv')\n",
        "# cv2.destroyAllWindows()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LLqD9hrQk2Zq",
        "outputId": "f412bbb4-2c03-42f7-b9d6-3339a798bbf3"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Mounted at /content/drive\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yYG2DxY0D8yM"
      },
      "source": [
        "!mv results.mp4  /content/drive/MyDrive/res_obj/results.mp4"
      ],
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MquUEkm4CVP-"
      },
      "source": [
        "import json \n",
        "with open('/content/drive/MyDrive/cricket_cor.json', 'w') as f:\n",
        "  json.dump(str(cor),f ,indent=2)"
      ],
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lYnOxprty3TD"
      },
      "source": [
        "## Digging into the model's intermediate predictions\n",
        "\n",
        "For this part we will assume that the detection model is a CenterNet model following Zhou et al (https://arxiv.org/abs/1904.07850).  And more specifically, we will assume that `detection_model` is of type `meta_architectures.center_net_meta_arch.CenterNetMetaArch`.\n",
        "\n",
        "As one of its intermediate predictions, CenterNet produces a heatmap of box centers for each class (for example, it will produce a heatmap whose size is proportional to that of the image that lights up at the center of each, e.g., \"zebra\"). In the following, we will visualize these intermediate class center heatmap predictions."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xBgYgSGMhHVi"
      },
      "source": [
        "if detection_model.__class__.__name__ != 'CenterNetMetaArch':\n",
        "  raise AssertionError('The meta-architecture for this section '\n",
        "  'is assumed to be CenterNetMetaArch!')\n",
        "\n",
        "def get_heatmap(predictions_dict, class_name):\n",
        "  \"\"\"Grabs class center logits and apply inverse logit transform.\n",
        "\n",
        "  Args:\n",
        "    predictions_dict: dictionary of tensors containing a `object_center`\n",
        "      field of shape [1, heatmap_width, heatmap_height, num_classes]\n",
        "    class_name: string name of category (e.g., `horse`)\n",
        "\n",
        "  Returns:\n",
        "    heatmap: 2d Tensor heatmap representing heatmap of centers for a given class\n",
        "      (For CenterNet, this is 128x128 or 256x256) with values in [0,1]\n",
        "  \"\"\"\n",
        "  class_index = label_map_dict[class_name]\n",
        "  class_center_logits = predictions_dict['object_center'][0]\n",
        "  class_center_logits = class_center_logits[0][\n",
        "    :, :, class_index - label_id_offset]\n",
        "  heatmap = tf.exp(class_center_logits) / (tf.exp(class_center_logits) + 1)\n",
        "  return heatmap\n",
        "\n",
        "def unpad_heatmap(heatmap, image_np):\n",
        "  \"\"\"Reshapes/unpads heatmap appropriately.\n",
        "\n",
        "  Reshapes/unpads heatmap appropriately to match image_np.\n",
        "\n",
        "  Args:\n",
        "    heatmap: Output of `get_heatmap`, a 2d Tensor\n",
        "    image_np: uint8 numpy array with shape (img_height, img_width, 3).  Note\n",
        "      that due to padding, the relationship between img_height and img_width\n",
        "      might not be a simple scaling.\n",
        "\n",
        "  Returns:\n",
        "    resized_heatmap_unpadded: a resized heatmap (2d Tensor) that is the same\n",
        "      size as `image_np`\n",
        "  \"\"\"\n",
        "  heatmap = tf.tile(tf.expand_dims(heatmap, 2), [1, 1, 3]) * 255\n",
        "  pre_strided_size = detection_model._stride * heatmap.shape[0]\n",
        "  resized_heatmap = tf.image.resize(\n",
        "      heatmap, [pre_strided_size, pre_strided_size],\n",
        "      method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)\n",
        "  resized_heatmap_unpadded = tf.slice(resized_heatmap, begin=[0,0,0], size=shapes)\n",
        "  return tf.image.resize(\n",
        "      resized_heatmap_unpadded,\n",
        "      [image_np.shape[0], image_np.shape[1]],\n",
        "      method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)[:,:,0]\n",
        "\n",
        "\n",
        "class_name = 'kite'\n",
        "heatmap = get_heatmap(predictions_dict, class_name)\n",
        "resized_heatmap_unpadded = unpad_heatmap(heatmap, image_np)\n",
        "plt.figure(figsize=(12,16))\n",
        "plt.imshow(image_np_with_detections)\n",
        "plt.imshow(resized_heatmap_unpadded, alpha=0.7,vmin=0, vmax=160, cmap='viridis')\n",
        "plt.title('Object center heatmap (class: ' + class_name + ')')\n",
        "plt.show()\n",
        "\n",
        "class_name = 'person'\n",
        "heatmap = get_heatmap(predictions_dict, class_name)\n",
        "resized_heatmap_unpadded = unpad_heatmap(heatmap, image_np)\n",
        "plt.figure(figsize=(12,16))\n",
        "plt.imshow(image_np_with_detections)\n",
        "plt.imshow(resized_heatmap_unpadded, alpha=0.7,vmin=0, vmax=160, cmap='viridis')\n",
        "plt.title('Object center heatmap (class: ' + class_name + ')')\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}