spatial_saliency_colab.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oMdtOGQRttwD",
        "outputId": "13ff9454-015d-4c7c-b82d-5ea4b60473ba"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: tensorflow-gpu==1.15.0 in /usr/local/lib/python3.7/dist-packages (1.15.0)\n",
            "Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==1.15.0) (1.1.0)\n",
            "Requirement already satisfied: numpy<2.0,>=1.16.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==1.15.0) (1.21.5)\n",
            "Requirement already satisfied: keras-preprocessing>=1.0.5 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==1.15.0) (1.1.2)\n",
            "Requirement already satisfied: google-pasta>=0.1.6 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==1.15.0) (0.2.0)\n",
            "Requirement already satisfied: six>=1.10.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==1.15.0) (1.15.0)\n",
            "Requirement already satisfied: absl-py>=0.7.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==1.15.0) (1.0.0)\n",
            "Requirement already satisfied: keras-applications>=1.0.8 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==1.15.0) (1.0.8)\n",
            "Requirement already satisfied: opt-einsum>=2.3.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==1.15.0) (3.3.0)\n",
            "Requirement already satisfied: tensorboard<1.16.0,>=1.15.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==1.15.0) (1.15.0)\n",
            "Requirement already satisfied: tensorflow-estimator==1.15.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==1.15.0) (1.15.1)\n",
            "Requirement already satisfied: wrapt>=1.11.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==1.15.0) (1.14.0)\n",
            "Requirement already satisfied: grpcio>=1.8.6 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==1.15.0) (1.44.0)\n",
            "Requirement already satisfied: gast==0.2.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==1.15.0) (0.2.2)\n",
            "Requirement already satisfied: protobuf>=3.6.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==1.15.0) (3.17.3)\n",
            "Requirement already satisfied: astor>=0.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==1.15.0) (0.8.1)\n",
            "Requirement already satisfied: wheel>=0.26 in /usr/local/lib/python3.7/dist-packages (from tensorflow-gpu==1.15.0) (0.37.1)\n",
            "Requirement already satisfied: h5py in /usr/local/lib/python3.7/dist-packages (from keras-applications>=1.0.8->tensorflow-gpu==1.15.0) (3.1.0)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.7/dist-packages (from tensorboard<1.16.0,>=1.15.0->tensorflow-gpu==1.15.0) (3.3.6)\n",
            "Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.7/dist-packages (from tensorboard<1.16.0,>=1.15.0->tensorflow-gpu==1.15.0) (1.0.1)\n",
            "Requirement already satisfied: setuptools>=41.0.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard<1.16.0,>=1.15.0->tensorflow-gpu==1.15.0) (57.4.0)\n",
            "Requirement already satisfied: importlib-metadata>=4.4 in /usr/local/lib/python3.7/dist-packages (from markdown>=2.6.8->tensorboard<1.16.0,>=1.15.0->tensorflow-gpu==1.15.0) (4.11.3)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata>=4.4->markdown>=2.6.8->tensorboard<1.16.0,>=1.15.0->tensorflow-gpu==1.15.0) (3.8.0)\n",
            "Requirement already satisfied: typing-extensions>=3.6.4 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata>=4.4->markdown>=2.6.8->tensorboard<1.16.0,>=1.15.0->tensorflow-gpu==1.15.0) (4.1.1)\n",
            "Requirement already satisfied: cached-property in /usr/local/lib/python3.7/dist-packages (from h5py->keras-applications>=1.0.8->tensorflow-gpu==1.15.0) (1.5.2)\n"
          ]
        }
      ],
      "source": [
        "!pip install tensorflow-gpu==1.15.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tW5e_Rfk1T01",
        "outputId": "68f161bb-bb5c-4d28-c90f-f095430795f7"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting wget\n",
            "  Downloading wget-3.2.zip (10 kB)\n",
            "Building wheels for collected packages: wget\n",
            "  Building wheel for wget (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for wget: filename=wget-3.2-py3-none-any.whl size=9675 sha256=3179dacf782c8a4b2823659248777149452daa950ae609686d6b070990a4d6d3\n",
            "  Stored in directory: /root/.cache/pip/wheels/a1/b6/7c/0e63e34eb06634181c63adacca38b79ff8f35c37e3c13e3c02\n",
            "Successfully built wget\n",
            "Installing collected packages: wget\n",
            "Successfully installed wget-3.2\n",
            "Collecting pysaliency\n",
            "  Downloading pysaliency-0.2.21.tar.gz (156 kB)\n",
            "\u001b[K     |████████████████████████████████| 156 kB 2.4 MB/s \n",
            "\u001b[?25hCollecting boltons\n",
            "  Downloading boltons-21.0.0-py2.py3-none-any.whl (193 kB)\n",
            "\u001b[K     |████████████████████████████████| 193 kB 32.8 MB/s \n",
            "\u001b[?25hCollecting deprecation\n",
            "  Downloading deprecation-2.1.0-py2.py3-none-any.whl (11 kB)\n",
            "Requirement already satisfied: dill in /usr/local/lib/python3.7/dist-packages (from pysaliency) (0.3.4)\n",
            "Requirement already satisfied: imageio in /usr/local/lib/python3.7/dist-packages (from pysaliency) (2.4.1)\n",
            "Requirement already satisfied: natsort in /usr/local/lib/python3.7/dist-packages (from pysaliency) (5.5.0)\n",
            "Requirement already satisfied: numba in /usr/local/lib/python3.7/dist-packages (from pysaliency) (0.51.2)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from pysaliency) (1.21.5)\n",
            "Collecting piexif\n",
            "  Downloading piexif-1.1.3-py2.py3-none-any.whl (20 kB)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from pysaliency) (2.23.0)\n",
            "Collecting schema\n",
            "  Downloading schema-0.7.5-py2.py3-none-any.whl (17 kB)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from pysaliency) (1.4.1)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.7/dist-packages (from pysaliency) (57.4.0)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from pysaliency) (4.64.0)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.7/dist-packages (from deprecation->pysaliency) (21.3)\n",
            "Requirement already satisfied: pillow in /usr/local/lib/python3.7/dist-packages (from imageio->pysaliency) (7.1.2)\n",
            "Requirement already satisfied: llvmlite<0.35,>=0.34.0.dev0 in /usr/local/lib/python3.7/dist-packages (from numba->pysaliency) (0.34.0)\n",
            "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging->deprecation->pysaliency) (3.0.8)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->pysaliency) (3.0.4)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->pysaliency) (2021.10.8)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests->pysaliency) (1.24.3)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->pysaliency) (2.10)\n",
            "Requirement already satisfied: contextlib2>=0.5.5 in /usr/local/lib/python3.7/dist-packages (from schema->pysaliency) (0.5.5)\n",
            "Building wheels for collected packages: pysaliency\n",
            "  Building wheel for pysaliency (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for pysaliency: filename=pysaliency-0.2.21-cp37-cp37m-linux_x86_64.whl size=332170 sha256=6bed8372bb0c283073e89db97bd95521032d143e9eb8a966fc17223c0a5f7568\n",
            "  Stored in directory: /root/.cache/pip/wheels/d4/ad/d7/4f7f72d119a3e9cb85dc95b5cc88e1111fab0705b3d2b5feeb\n",
            "Successfully built pysaliency\n",
            "Installing collected packages: schema, piexif, deprecation, boltons, pysaliency\n",
            "Successfully installed boltons-21.0.0 deprecation-2.1.0 piexif-1.1.3 pysaliency-0.2.21 schema-0.7.5\n"
          ]
        }
      ],
      "source": [
        "!pip install wget\n",
        "!pip install pysaliency"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "q32jlGyb1Cqk",
        "outputId": "7f77e59e-6f9e-4d0d-d471-7e14ed6b02e1"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            ">> Downloading COCOSearch18 dataset..."
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Downloading...\n",
            "From: https://drive.google.com/uc?export=download&id=1vEzgF54LPK2adlI7DdlXWGkYV76L-jjK\n",
            "To: /content/targets.zip\n",
            "100%|██████████| 14.7M/14.7M [00:00<00:00, 42.7MB/s]\n",
            "Downloading...\n",
            "From: https://drive.google.com/uc?id=1OkpX_Md-lFwCo5TB_cq0Qxoe4oEB8eKG\n",
            "To: /content/bbox_annos.npy\n",
            "100%|██████████| 127k/127k [00:00<00:00, 24.5MB/s]\n",
            "Downloading...\n",
            "From: https://drive.google.com/u/0/uc?export=download&confirm=ATmP&id=1ff0va472Xs1bvidCwRlW3Ctf7Hbyyn7p\n",
            "To: /content/weights/vgg16_hybrid.zip\n",
            "100%|██████████| 519M/519M [00:08<00:00, 64.5MB/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "done!\n"
          ]
        }
      ],
      "source": [
        "'''This file downloads COCO-Search18 dataset, target images,\n",
        " and VGG16 pretrained weights on ImageNet.'''\n",
        "\n",
        "import os\n",
        "import wget\n",
        "import zipfile\n",
        "import gdown\n",
        "import argparse\n",
        "\n",
        "def unzip(zip_path, extract_path):\n",
        "    \"\"\"extracts the files in a zip file\n",
        "       in the specified directory\n",
        "    args:\n",
        "        zip_path (str): the path to the zip file\n",
        "        extract_path (str): the path to save the extracted files \n",
        "    \"\"\"\n",
        "\n",
        "    with zipfile.ZipFile(zip_path, \"r\") as zip_ref:\n",
        "        for file in zip_ref.namelist():\n",
        "            zip_ref.extract(file, extract_path)\n",
        "\n",
        "    os.remove(zip_path)\n",
        "\n",
        "    return\n",
        "\n",
        "\n",
        "def download_cocosearch(data_path):\n",
        "    \"\"\"Downloads the COCOSearch18 dataset. The dataset\n",
        "       contains the image stimuli and label/annotation files.\n",
        "    Args:\n",
        "        data_path (str): Defines the path where the dataset will be\n",
        "                         downloaded and extracted to.\n",
        "    \"\"\"\n",
        "\n",
        "    print(\">> Downloading COCOSearch18 dataset...\", end=\"\", flush=True)\n",
        "\n",
        "    os.makedirs(data_path, exist_ok=True)\n",
        "\n",
        "    urls = ['http://vision.cs.stonybrook.edu/~cvlab_download/COCOSearch18-images-TP.zip',\n",
        "            'https://saliency.tuebingen.ai/data/coco_search18_TP.zip']\n",
        "\n",
        "    for url in urls:\n",
        "        filename = wget.download(url, data_path)\n",
        "        unzip(filename, data_path)\n",
        "\n",
        "    url = \"https://drive.google.com/uc?export=download&id=1vEzgF54LPK2adlI7DdlXWGkYV76L-jjK\"\n",
        "\n",
        "    gdown.download(url, data_path, quiet=False)\n",
        "    unzip(os.path.join(data_path , 'targets.zip'), data_path)\n",
        "\n",
        "    # Downloading target object bounding box annotation\n",
        "    url = \"https://drive.google.com/uc?id=1OkpX_Md-lFwCo5TB_cq0Qxoe4oEB8eKG\"\n",
        "    output_path = os.path.join(data_path , 'bbox_annos.npy')\n",
        "    gdown.download(url, output_path, quiet=False)\n",
        "\n",
        "    url = \"https://drive.google.com/u/0/uc?export=download&confirm=ATmP&id=1ff0va472Xs1bvidCwRlW3Ctf7Hbyyn7p\"\n",
        "    weights_path = os.path.join(data_path , 'weights')\n",
        "    os.makedirs(weights_path, exist_ok=True)\n",
        "    gdown.download(url, os.path.join(weights_path , 'vgg16_hybrid.zip'), quiet=False)\n",
        "    unzip(os.path.join(weights_path , 'vgg16_hybrid.zip'), weights_path)\n",
        "\n",
        "    print(\"done!\", flush=True)\n",
        "    return\n",
        "\n",
        "if __name__ == '__main__':\n",
        "\n",
        "    download_cocosearch('/content/')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sv-KNlv81Ylv",
        "outputId": "8601d20d-9aad-4b20-e7d1-4c1d9eab68ad"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "3652\n",
            "324\n",
            "324\n"
          ]
        }
      ],
      "source": [
        "import numpy as np\n",
        "from os.path import join\n",
        "from itertools import groupby\n",
        "import argparse\n",
        "import os\n",
        "import json\n",
        "import random\n",
        "import cv2\n",
        "\n",
        "\n",
        "def GaussianMask(sizex, sizey, sigma=11, center=None, fix=1):\n",
        "    \"\"\"Blurs each fixation point by convolving it\n",
        "        with a Gaussian kernel. This function is adopted from\n",
        "        https://github.com/takyamamoto/Fixation-Densitymap repository.\n",
        "    args:\n",
        "        sizex (int): mask width\n",
        "        sizey (int): mask height\n",
        "        sigma (int): gaussian std\n",
        "        center (tuple): gaussian mean\n",
        "        fix (int or float): gaussian max\n",
        "    returns:\n",
        "        gaussian mask\n",
        "    \"\"\"\n",
        "\n",
        "    x = np.arange(0, sizex, 1, float)\n",
        "    y = np.arange(0, sizey, 1, float)\n",
        "    x, y = np.meshgrid(x, y)\n",
        "\n",
        "    if center is None:\n",
        "        x0 = sizex // 2\n",
        "        y0 = sizey // 2\n",
        "    else:\n",
        "        if np.isnan(center[0]) == False and np.isnan(center[1]) == False:\n",
        "            x0 = center[0]\n",
        "            y0 = center[1]\n",
        "        else:\n",
        "            return np.zeros((sizey, sizex))\n",
        "\n",
        "    return fix * np.exp(-0.5 * ((x - x0) ** 2 + (y - y0) ** 2) / sigma ** 2)\n",
        "\n",
        "\n",
        "def preprocess_fixations(phase,\n",
        "                         task_img_pair,\n",
        "                         trajs,\n",
        "                         im_h,\n",
        "                         im_w,\n",
        "                         bbox,\n",
        "                         sigma,\n",
        "                         dldir,\n",
        "                         datadir,\n",
        "                         truncate_num=-1):\n",
        "    \"\"\"Processes fixation data and creates \n",
        "        fixation maps. Resizes all search and target images \n",
        "        and save them in the corresponding directories. \n",
        "        Splits data into train-validation-test sets. \n",
        "        Augments the training data.\n",
        "        saves unblurred fixation maps for saliency metric computation.\n",
        "        saves target bbox overlayed test images for the purpose of results visualization.\n",
        "    Args:\n",
        "        phase (str): train or valid set (test set is separated from train set)\n",
        "        task_img_pair (list): a list of task-image pairs \n",
        "        trajs (list): a list of all trials \n",
        "        im_h (int): resize search images to this height\n",
        "        im_w (int): resize search images to this width\n",
        "        bbox  (dict): target object bbox for each task-image pair\n",
        "        sigma (int): sigma for Gaussian blurring function\n",
        "        dldir (str): directory of downloaded data\n",
        "        datadir (str): directory to save the preprocessed train/val/test sets\n",
        "        truncate_num (int): maximum number of fixations to be processed from each trial\n",
        "    \"\"\"\n",
        "    fix_labels = []\n",
        "    stimuli = []\n",
        "    heat_maps_list = []\n",
        "    min_fix_x = 100000\n",
        "    max_fix_x = -100000\n",
        "    min_fix_y = 100000\n",
        "    max_fix_y = -100000\n",
        "    flat_test_task_img_pair = []\n",
        "\n",
        "    if phase == 'train':\n",
        "        test_task_img_pair = []\n",
        "        for key, group in groupby(task_img_pair, lambda x: x.split('_')[0]):\n",
        "            key_and_group = {key: random.sample(list(group), 18)}\n",
        "            test_task_img_pair.append(key_and_group[key])\n",
        "\n",
        "        flat_test_task_img_pair = [item for sublist in test_task_img_pair for item in sublist]\n",
        "\n",
        "    for traj in trajs:\n",
        "        for i in range(len(traj['X'])):\n",
        "            if traj['X'][i] < 0 or traj['Y'][i] < 0 or traj['X'][i] > 1680 or traj['Y'][i] > 1050:\n",
        "                continue\n",
        "\n",
        "            if traj['X'][i] < min_fix_x:\n",
        "                min_fix_x = traj['X'][i]\n",
        "\n",
        "            if traj['X'][i] > max_fix_x:\n",
        "                max_fix_x = traj['X'][i]\n",
        "\n",
        "            if traj['Y'][i] < min_fix_y:\n",
        "                min_fix_y = traj['Y'][i]\n",
        "\n",
        "            if traj['Y'][i] > max_fix_y:\n",
        "                max_fix_y = traj['Y'][i]\n",
        "\n",
        "    for task_img in task_img_pair:\n",
        "\n",
        "        heatmap = np.zeros((im_h, im_w), np.float32)\n",
        "        heatmap_unblurred = np.zeros((im_h, im_w), np.float32)\n",
        "\n",
        "        x1 = bbox[task_img][0]\n",
        "        y1 = bbox[task_img][1]\n",
        "        w_image = bbox[task_img][2]\n",
        "        h_image = bbox[task_img][3]\n",
        "\n",
        "        for traj in trajs:\n",
        "\n",
        "            if (traj['task'] + '_' + traj['name']) == task_img:\n",
        "\n",
        "                # first fixations are fixed at the screen center\n",
        "                traj['X'][0], traj['Y'][0] = im_w / 2, im_h / 2\n",
        "                if truncate_num < 1:\n",
        "                    traj_len = len(traj['X'])\n",
        "                else:\n",
        "                    traj_len = min(truncate_num, len(traj['X']))\n",
        "\n",
        "                for i in range(1, traj_len):\n",
        "                    # remove out of boundary fixations\n",
        "                    if traj['X'][i] < 0 or traj['Y'][i] < 0 or traj['X'][i] > 1680 or traj['Y'][i] > 1050:\n",
        "                        continue\n",
        "                    fix = (\n",
        "                        ((traj['X'][i] - min_fix_x) / max_fix_x) * (512),\n",
        "                        ((traj['Y'][i] - min_fix_y) / max_fix_y) * (320))\n",
        "                    \n",
        "                    # masking the target, uncomment if you want to mask the target\n",
        "                    '''\n",
        "                    if (x1<=fix[0]<=x1+w_image and y1<=fix[1]<=y1+h_image):\n",
        "                        continue\n",
        "                    else:\n",
        "                    '''\n",
        "                    heatmap += GaussianMask(im_w, im_h, sigma, (fix[0], fix[1]))\n",
        "                    heatmap_unblurred[int(fix[1]), int(fix[0])] = 1\n",
        "\n",
        "        # Normalization\n",
        "        heatmap = heatmap / np.amax(heatmap)\n",
        "        heatmap_np = heatmap * 255\n",
        "        heatmap = heatmap_np.astype(\"uint8\")\n",
        "\n",
        "        heatmap_unblurred = heatmap_unblurred / np.amax(heatmap_unblurred)\n",
        "        heatmap_unblurred_np = heatmap_unblurred * 255\n",
        "        heatmap_unblurred = heatmap_unblurred_np.astype(\"uint8\")\n",
        "\n",
        "        source = os.path.join(dldir , 'images' , str(task_img.split('_')[0]) , str(task_img.split('_')[1]))\n",
        "        heatmap_flip = cv2.flip(heatmap, 1)\n",
        "        img = cv2.imread(source)\n",
        "        img_resized = cv2.resize(img, (im_w, im_h), interpolation=cv2.INTER_AREA)\n",
        "        # bbox = [top left x position, top left y position, width, height].\n",
        "        img_resized_flip = cv2.flip(img_resized, 1)\n",
        "\n",
        "        target_0 = cv2.imread(os.path.join(dldir , 'targets' , (task_img.split('_')[\n",
        "            0] + '_0.png')))  # img_resized[y1:y1+h_image , x1:x1+w_image, :]\n",
        "\n",
        "        target_1 = cv2.imread(os.path.join(dldir , 'targets' , (task_img.split('_')[\n",
        "            0] + '_1.png'))) \n",
        "\n",
        "        target_2 = cv2.imread(os.path.join(dldir , 'targets' , (task_img.split('_')[\n",
        "            0] + '_2.png'))) \n",
        "\n",
        "        target_3 = cv2.imread(os.path.join(dldir , 'targets' , (task_img.split('_')[\n",
        "            0] + '_3.png'))) \n",
        "\n",
        "        target_4 = cv2.imread(os.path.join(dldir , 'targets' ,(task_img.split('_')[\n",
        "            0] + '_4.png'))) \n",
        "\n",
        "        target_0 = cv2.resize(target_0, (64, 64), interpolation=cv2.INTER_AREA)\n",
        "        target_flip_0 = cv2.flip(target_0, 1)\n",
        "\n",
        "        target_1 = cv2.resize(target_1, (64, 64), interpolation=cv2.INTER_AREA)\n",
        "        target_flip_1 = cv2.flip(target_1, 1)\n",
        "\n",
        "        target_2 = cv2.resize(target_2, (64, 64), interpolation=cv2.INTER_AREA)\n",
        "        target_flip_2 = cv2.flip(target_2, 1)\n",
        "\n",
        "        target_3 = cv2.resize(target_3, (64, 64), interpolation=cv2.INTER_AREA)\n",
        "        target_flip_3 = cv2.flip(target_3, 1)\n",
        "\n",
        "        target_4 = cv2.resize(target_4, (64, 64), interpolation=cv2.INTER_AREA)\n",
        "        target_flip_4 = cv2.flip(target_4, 1)\n",
        "\n",
        "        img_target_frame=cv2.rectangle(img_resized.copy(),(x1,y1),(x1+w_image,y1+h_image),(0,255,0),2)\n",
        "        \n",
        "        unblur = False\n",
        "        flip_f = False\n",
        "\n",
        "        if phase == 'train':\n",
        "\n",
        "            if task_img in flat_test_task_img_pair:\n",
        " \n",
        "                unblur = True\n",
        "\n",
        "                out_name = os.path.join(datadir , 'saliencymap/test' , str(task_img))\n",
        "                out_name_np = os.path.join(datadir , 'saliencymap/test' , (os.path.splitext(str(task_img))[0]+'.npy'))\n",
        "                \n",
        "                with open(out_name_np, \"wb\") as file:\n",
        "                    np.save(file, heatmap_np )\n",
        "\n",
        "                destination = os.path.join(datadir , 'stimuli/test' , str(task_img))\n",
        "                target_path_0 = os.path.join(datadir , 'target_0/test' , str(task_img))\n",
        "                target_path_1 = os.path.join(datadir , 'target_1/test' , str(task_img))\n",
        "                target_path_2 = os.path.join(datadir , 'target_2/test' , str(task_img))\n",
        "                target_path_3 = os.path.join(datadir , 'target_3/test' , str(task_img))\n",
        "                target_path_4 = os.path.join(datadir , 'target_4/test' , str(task_img))\n",
        "\n",
        "                img_target_rect_path = os.path.join(datadir , 'stimuli/test_targ_bbox' , str(task_img))\n",
        "                cv2.imwrite(img_target_rect_path, img_target_frame)\n",
        "                \n",
        "                out_name_unblur = os.path.join(datadir , 'saliencymap/test_unblur' , str(task_img))\n",
        "                out_name_unblur_npy = os.path.join(datadir , 'saliencymap/test_unblur' , (os.path.splitext(str(task_img))[0]+'.npy'))\n",
        "\n",
        "            else:\n",
        "\n",
        "                flip_f = True\n",
        "\n",
        "                out_name = os.path.join(datadir , 'saliencymap/train' , str(task_img))\n",
        "                destination = os.path.join(datadir , 'stimuli/train' , str(task_img))\n",
        "\n",
        "                target_path_0 = os.path.join(datadir , 'target_0/train' , str(task_img))\n",
        "                target_path_1 = os.path.join(datadir , 'target_1/train' , str(task_img))\n",
        "                target_path_2 = os.path.join(datadir , 'target_2/train' , str(task_img))\n",
        "                target_path_3 = os.path.join(datadir , 'target_3/train' , str(task_img))\n",
        "                target_path_4 = os.path.join(datadir , 'target_4/train' , str(task_img))\n",
        "\n",
        "                sal_out_flip = os.path.join(datadir , 'saliencymap/train' , (str(task_img.split('.')[0]) + '_flip.' + str(\n",
        "                    task_img.split('.')[1])))\n",
        "                stim_out_flip = os.path.join(datadir , 'stimuli/train' , (str(task_img.split('.')[0]) + '_flip.' + str(\n",
        "                    task_img.split('.')[1])))\n",
        "\n",
        "                tar_out_flip_0 = os.path.join(datadir , 'target_0/train' , (str(task_img.split('.')[0]) + '_flip.' + str(\n",
        "                    task_img.split('.')[1])))\n",
        "                tar_out_flip_1 = os.path.join(datadir , 'target_1/train' , (str(task_img.split('.')[0]) + '_flip.' + str(\n",
        "                    task_img.split('.')[1])))\n",
        "                tar_out_flip_2 = os.path.join(datadir , 'target_2/train' , (str(task_img.split('.')[0]) + '_flip.' + str(\n",
        "                    task_img.split('.')[1])))\n",
        "                tar_out_flip_3 = os.path.join(datadir , 'target_3/train' , (str(task_img.split('.')[0]) + '_flip.' + str(\n",
        "                    task_img.split('.')[1])))\n",
        "                tar_out_flip_4 = os.path.join(datadir , 'target_4/train' , (str(task_img.split('.')[0]) + '_flip.' + str(\n",
        "                    task_img.split('.')[1])))\n",
        "        else:\n",
        "            \n",
        "            out_name = os.path.join(datadir , 'saliencymap/valid' , str(task_img))\n",
        "            destination = os.path.join(datadir , 'stimuli/valid' , str(task_img))\n",
        "\n",
        "            target_path_0 = os.path.join(datadir , 'target_0/valid' , str(task_img))\n",
        "            target_path_1 = os.path.join(datadir , 'target_1/valid' , str(task_img))\n",
        "            target_path_2 = os.path.join(datadir , 'target_2/valid' , str(task_img))\n",
        "            target_path_3 = os.path.join(datadir , 'target_3/valid' , str(task_img))\n",
        "            target_path_4 = os.path.join(datadir , 'target_4/valid' , str(task_img))\n",
        "\n",
        "        #uncomment this part to save colorful heatmap of fixations overlayed on images\n",
        "        '''#create groundtruth heatmaps\n",
        "        heatmap_color = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)\n",
        "\n",
        "        # Create mask\n",
        "        threshold = 30 #10\n",
        "\n",
        "        mask = np.where(heatmap <= threshold, 1, 0)\n",
        "        mask = np.reshape(mask, (im_h, im_w, 1))\n",
        "        mask = np.repeat(mask, 3, axis=2)\n",
        "\n",
        "        # Marge images\n",
        "\n",
        "        marge = img_resized * mask + heatmap_color * (1 - mask)\n",
        "        marge = marge.astype(\"uint8\")\n",
        "        alpha = 0.5\n",
        "\n",
        "        marge = cv2.addWeighted(img_resized, 1 - alpha, marge, alpha, 0)'''\n",
        "\n",
        "        cv2.imwrite(destination, img_resized)\n",
        "        cv2.imwrite(out_name, heatmap)\n",
        "        #cv2.imwrite(out_name, marge)\n",
        "        cv2.imwrite(target_path_0, target_0)\n",
        "        cv2.imwrite(target_path_1, target_1)\n",
        "        cv2.imwrite(target_path_2, target_2)\n",
        "        cv2.imwrite(target_path_3, target_3)\n",
        "        cv2.imwrite(target_path_4, target_4)\n",
        "\n",
        "        if flip_f:\n",
        "            cv2.imwrite(stim_out_flip, img_resized_flip)\n",
        "            cv2.imwrite(sal_out_flip, heatmap_flip)\n",
        "            cv2.imwrite(tar_out_flip_0, target_flip_0)\n",
        "            cv2.imwrite(tar_out_flip_1, target_flip_1)\n",
        "            cv2.imwrite(tar_out_flip_2, target_flip_2)\n",
        "            cv2.imwrite(tar_out_flip_3, target_flip_3)\n",
        "            cv2.imwrite(tar_out_flip_4, target_flip_4)\n",
        "\n",
        "        if unblur:\n",
        "            cv2.imwrite(out_name_unblur, heatmap_unblurred)\n",
        "            with open(out_name_unblur_npy, \"wb\") as file:\n",
        "                    np.save(file, heatmap_unblurred_np)\n",
        "\n",
        "    return \n",
        "\n",
        "def process_data(trajs_train,\n",
        "                 trajs_valid,\n",
        "                 target_annos,\n",
        "                 sigma,\n",
        "                 dldir,\n",
        "                 datadir):\n",
        "    \"\"\"creates task-image pairs for training and validation sets\n",
        "        then calls preprocess_fixations func to \n",
        "        create fixation maps and train-test-valid split.\n",
        "    args:\n",
        "        trajs_train (list): a list of all trials in the original dataset training split \n",
        "        trajs_valid (list): a list of all trials in the original dataset validation split \n",
        "        target_annos (dict):  contains target object bbox for each task-image pair\n",
        "        sigma (int): sigma for Gaussian blurring function\n",
        "        dldir (str): directory of downloaded data\n",
        "        datadir (str): directory to save the preprocessed train/val/test sets\n",
        "    \"\"\"\n",
        "\n",
        "    im_w = 512\n",
        "    im_h = 320\n",
        "    #max_traj_length = 6\n",
        "\n",
        "    target_init_fixs = {}\n",
        "    cat_names = list(np.unique([x['task'] for x in trajs_train]))\n",
        "    catIds = dict(zip(cat_names, list(range(len(cat_names)))))\n",
        "\n",
        "    # training fixation data\n",
        "    train_task_img_pair = np.unique(\n",
        "        [traj['task'] + '_' + traj['name'] for traj in trajs_train])\n",
        "\n",
        "    # uncomment this part to process train data for only a single category\n",
        "    '''train_task_img_pair = []\n",
        "    for traj in trajs_train:\n",
        "      if traj['task'] =='tv':\n",
        "        train_task_img_pair.append(traj['task'] + '_' + traj['name'])\n",
        "    train_task_img_pair = np.unique(np.asarray(train_task_img_pair))'''\n",
        "\n",
        "    preprocess_fixations(\n",
        "        'train',\n",
        "        train_task_img_pair,\n",
        "        trajs_train,\n",
        "        im_h,\n",
        "        im_w,\n",
        "        target_annos,\n",
        "        sigma,\n",
        "        dldir,\n",
        "        datadir,\n",
        "        truncate_num=-1)\n",
        "\n",
        "    # validation fixation data\n",
        "    valid_task_img_pair = np.unique(\n",
        "        [traj['task'] + '_' + traj['name'] for traj in trajs_valid])\n",
        "\n",
        "    # uncomment this part to process valid data for only a single category\n",
        "    '''valid_task_img_pair = []\n",
        "    for traj in trajs_valid:\n",
        "      if traj['task'] =='tv':\n",
        "        valid_task_img_pair.append(traj['task'] + '_' + traj['name'])\n",
        "    valid_task_img_pair = np.unique(np.array(valid_task_img_pair))'''\n",
        "\n",
        "    preprocess_fixations(\n",
        "        'valid',\n",
        "        valid_task_img_pair,\n",
        "        trajs_valid,\n",
        "        im_h,\n",
        "        im_w,\n",
        "        target_annos,\n",
        "        sigma,\n",
        "        dldir,\n",
        "        datadir,\n",
        "        truncate_num=-1)\n",
        "\n",
        "    return \n",
        "\n",
        "\n",
        "if __name__ == '__main__':\n",
        "\n",
        "    #The directory to save images along with their fixation maps.\n",
        "    datadir = os.path.join('.' , 'cocosearch/')\n",
        "\n",
        "    sl_map = os.path.join(datadir, 'saliencymap')\n",
        "    tr_sl_map = os.path.join(sl_map, 'train')\n",
        "    v_sl_map = os.path.join(sl_map, 'valid')\n",
        "    te_sl_map = os.path.join(sl_map, 'test')\n",
        "    te_unbur_sl_map = os.path.join(sl_map, 'test_unblur')\n",
        "\n",
        "    stimuli = os.path.join(datadir, 'stimuli')\n",
        "    tr_stimuli = os.path.join(stimuli, 'train')\n",
        "    v_stimuli = os.path.join(stimuli, 'valid')\n",
        "    te_stimuli = os.path.join(stimuli, 'test')\n",
        "    targ_t_stimuli = os.path.join(stimuli, 'test_targ_bbox')\n",
        "\n",
        "    target_0 = os.path.join(datadir, 'target_0')\n",
        "    tr_target_0 = os.path.join(target_0, 'train')\n",
        "    v_target_0 = os.path.join(target_0, 'valid')\n",
        "    te_target_0 = os.path.join(target_0, 'test')\n",
        "\n",
        "    target_1 = os.path.join(datadir, 'target_1')\n",
        "    tr_target_1 = os.path.join(target_1, 'train')\n",
        "    v_target_1 = os.path.join(target_1, 'valid')\n",
        "    te_target_1 = os.path.join(target_1, 'test')\n",
        "\n",
        "    target_2 = os.path.join(datadir, 'target_2')\n",
        "    tr_target_2 = os.path.join(target_2, 'train')\n",
        "    v_target_2 = os.path.join(target_2, 'valid')\n",
        "    te_target_2 = os.path.join(target_2, 'test')\n",
        "\n",
        "    target_3 = os.path.join(datadir, 'target_3')\n",
        "    tr_target_3 = os.path.join(target_3, 'train')\n",
        "    v_target_3 = os.path.join(target_3, 'valid')\n",
        "    te_target_3 = os.path.join(target_3, 'test')\n",
        "\n",
        "    target_4 = os.path.join(datadir, 'target_4')\n",
        "    tr_target_4 = os.path.join(target_4, 'train')\n",
        "    v_target_4 = os.path.join(target_4, 'valid')\n",
        "    te_target_4 = os.path.join(target_4, 'test')\n",
        "\n",
        "    os.makedirs('./', exist_ok=True)\n",
        "    os.makedirs(datadir, exist_ok=True)\n",
        "\n",
        "    os.makedirs(sl_map, exist_ok=True)\n",
        "    os.makedirs(tr_sl_map, exist_ok=True)\n",
        "    os.makedirs(v_sl_map, exist_ok=True)\n",
        "    os.makedirs(te_sl_map, exist_ok=True)\n",
        "    os.makedirs(te_unbur_sl_map, exist_ok=True)\n",
        "\n",
        "    os.makedirs(stimuli, exist_ok=True)\n",
        "    os.makedirs(tr_stimuli, exist_ok=True)\n",
        "    os.makedirs(v_stimuli, exist_ok=True)\n",
        "    os.makedirs(te_stimuli, exist_ok=True)\n",
        "    os.makedirs( targ_t_stimuli, exist_ok=True)\n",
        "\n",
        "    os.makedirs(target_0, exist_ok=True)\n",
        "    os.makedirs(tr_target_0, exist_ok=True)\n",
        "    os.makedirs(v_target_0, exist_ok=True)\n",
        "    os.makedirs(te_target_0, exist_ok=True)\n",
        "\n",
        "    os.makedirs(target_1, exist_ok=True)\n",
        "    os.makedirs(tr_target_1, exist_ok=True)\n",
        "    os.makedirs(v_target_1, exist_ok=True)\n",
        "    os.makedirs(te_target_1, exist_ok=True)\n",
        "\n",
        "    os.makedirs(target_2, exist_ok=True)\n",
        "    os.makedirs(tr_target_2, exist_ok=True)\n",
        "    os.makedirs(v_target_2, exist_ok=True)\n",
        "    os.makedirs(te_target_2, exist_ok=True)\n",
        "\n",
        "    os.makedirs(target_3, exist_ok=True)\n",
        "    os.makedirs(tr_target_3, exist_ok=True)\n",
        "    os.makedirs(v_target_3, exist_ok=True)\n",
        "    os.makedirs(te_target_3, exist_ok=True)\n",
        "\n",
        "    os.makedirs(target_4, exist_ok=True)\n",
        "    os.makedirs(tr_target_4, exist_ok=True)\n",
        "    os.makedirs(v_target_4, exist_ok=True)\n",
        "    os.makedirs(te_target_4, exist_ok=True)\n",
        "\n",
        "    dataset_root = '/content'\n",
        "\n",
        "    # bounding box of the target object (for search efficiency evaluation)\n",
        "    bbox_annos = np.load(join(dataset_root, 'bbox_annos.npy'),\n",
        "                         allow_pickle=True).item()\n",
        "\n",
        "    # load ground-truth human scanpaths\n",
        "\n",
        "    with open(join(dataset_root,\n",
        "                   'coco_search18_fixations_TP_train_split1.json')) as json_file:\n",
        "        human_scanpaths_train = json.load(json_file)\n",
        "\n",
        "    with open(join(dataset_root,\n",
        "                   'coco_search18_fixations_TP_validation_split1.json')) as json_file:\n",
        "        human_scanpaths_valid = json.load(json_file)\n",
        "\n",
        "    # exclude incorrect scanpaths\n",
        "    human_scanpaths_train = list(\n",
        "        filter(lambda x: x['correct'] == 1, human_scanpaths_train))\n",
        "    human_scanpaths_valid = list(\n",
        "        filter(lambda x: x['correct'] == 1, human_scanpaths_valid))\n",
        "\n",
        "    sigma = 11\n",
        "    # process fixation data\n",
        "    process_data(human_scanpaths_train, human_scanpaths_valid, bbox_annos,\n",
        "                           sigma, dataset_root, datadir)\n",
        "\n",
        "    train = next(os.walk(tr_stimuli))[2] \n",
        "    print(len(train))\n",
        "\n",
        "    valid = next(os.walk(v_stimuli))[2] \n",
        "    print(len(valid))\n",
        "\n",
        "    test = next(os.walk(te_stimuli))[2] \n",
        "    print(len(test))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "yUBSLQwkLwgM"
      },
      "outputs": [],
      "source": [
        "!rm -r './images/'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "TvySNcIfvC1R"
      },
      "outputs": [],
      "source": [
        "class config:\n",
        "  \"\"\"General training parameters that define the maximum number of\n",
        "  training epochs, the batch size, and learning rate for the ADAM\n",
        "  optimization method. To reproduce the results from the paper,\n",
        "  these values should not be changed. The device can be either\n",
        "  \"cpu\" or \"gpu\", which then optimizes the model accordingly after\n",
        "  training or uses the correct version for inference when testing.\n",
        "  \"\"\"\n",
        "\n",
        "  PARAMS = {\n",
        "      \"n_epochs\": 7,\n",
        "      \"batch_size\": 1,\n",
        "      \"n_training_steps\": 10000,\n",
        "      \"learning_rate\": 1e-5,\n",
        "      \"learning_power\": 0.5,\n",
        "      \"momentum\": 0.9,\n",
        "      \"device\": \"gpu\"\n",
        "  }\n",
        "\n",
        "  \"\"\"The predefined input image sizes for the search and target images.\n",
        "  To reproduce the results from the paper, these values should\n",
        "  not be changed. They must be divisible by 8 due to the model's\n",
        "  downsampling operations.\n",
        "  \"\"\" \n",
        "\n",
        "  DIMS = {\n",
        "      \"image_size_cocosearch\": (320, 512),\n",
        "      \"image_target_size_cocosearch\": (64, 64)\n",
        "  }"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "ch74rvDZwfJi"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import sys\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import random \n",
        "\n",
        "class COCOSEARCH:\n",
        "\n",
        "    \"\"\"This class represents the COCO-Search18 dataset. It consists of 3101\n",
        "       target-present images. All stimuli are of size 1680x1050 pixels\n",
        "       and are resized to 512x320 (height by width). It also randomly chooses\n",
        "       one of the five available sample target images for image stimuli. \n",
        "       Thus at each epoch when the dataset iterator is re-initialized, \n",
        "       the network randomly chooses a different set of sample targets.       \n",
        "\n",
        "    Attributes:\n",
        "        n_train: Number of training instances as defined in the dataset.\n",
        "        n_valid: Number of validation instances as defined in the dataset.\n",
        "\n",
        "    Returns:\n",
        "        tuple: A tuple that consists of dataset objects holding the training\n",
        "               and validation set instances respectively.\n",
        "    \"\"\"\n",
        "    n_train = 0\n",
        "    n_valid = 0\n",
        "\n",
        "    def __init__(self, data_path):\n",
        "\n",
        "        type(self).n_train = len(next(os.walk(data_path + \"cocosearch/stimuli/train\"))[2])\n",
        "        type(self).n_valid = len(next(os.walk(data_path + \"cocosearch/stimuli/valid\"))[2])\n",
        "        \n",
        "        self._stimuli_size = config.DIMS[\"image_size_cocosearch\"]\n",
        "        self._target_size = config.DIMS[\"image_target_size_cocosearch\"]\n",
        "\n",
        "        self._dir_stimuli_train = data_path + \"cocosearch/stimuli/train\"\n",
        "        self._dir_stimuli_valid = data_path + \"cocosearch/stimuli/valid\"\n",
        "\n",
        "        self._dir_saliency_train = data_path + \"cocosearch/saliencymap/train\"\n",
        "        self._dir_saliency_valid = data_path + \"cocosearch/saliencymap/valid\"\n",
        "\n",
        "        targ_ind_train = str(random.randint(0, 4))\n",
        "        targ_ind_valid = str(random.randint(0, 4))\n",
        "\n",
        "        self._dir_target_train = data_path + \"cocosearch/target_\" + targ_ind_train + \"/train\"\n",
        "        self._dir_target_valid = data_path + \"cocosearch/target_\" + targ_ind_valid + \"/valid\"\n",
        "\n",
        "\n",
        "    def load_data(self):\n",
        "\n",
        "        train_list_x = _get_file_list(self._dir_stimuli_train)\n",
        "        train_list_y = _get_file_list(self._dir_saliency_train)\n",
        "        train_list_z = _get_file_list(self._dir_target_train)\n",
        "\n",
        "        _check_consistency(zip(train_list_x, train_list_y, train_list_z), self.n_train)\n",
        "\n",
        "        train_set = _fetch_dataset((train_list_x, train_list_y, train_list_z),\n",
        "                                    self._stimuli_size, self._target_size, True)\n",
        "\n",
        "        valid_list_x = _get_file_list(self._dir_stimuli_valid)\n",
        "        valid_list_y = _get_file_list(self._dir_saliency_valid)\n",
        "        valid_list_z = _get_file_list(self._dir_target_valid)\n",
        "\n",
        "        _check_consistency(zip(valid_list_x, valid_list_y, valid_list_z), self.n_valid)\n",
        "\n",
        "        valid_set = _fetch_dataset((valid_list_x, valid_list_y, valid_list_z),\n",
        "                                    self._stimuli_size, self._target_size, False)\n",
        "\n",
        "        return (train_set, valid_set)\n",
        "\n",
        "\n",
        "class TEST:\n",
        "\n",
        "    \"\"\"This class represents test set instances used for inference through\n",
        "       a trained network. All stimuli are resized to the preferred spatial\n",
        "       dimensions of the chosen model. This can, however, lead to cases of\n",
        "       excessive image padding.\n",
        "    Returns:\n",
        "        object: A dataset object that holds all test set instances\n",
        "                specified under the path variable.\n",
        "    \"\"\"\n",
        "    n_test = 0\n",
        "    def __init__(self, dataset, data_path):\n",
        "        \n",
        "        type(self).n_test = len(next(os.walk(data_path + \"cocosearch/stimuli/test\"))[2])\n",
        "\n",
        "        self._stimuli_size = config.DIMS[\"image_size_cocosearch\"]\n",
        "        self._target_size = config.DIMS[\"image_target_size_cocosearch\"]\n",
        "\n",
        "        targ_ind_test = str(random.randint(0, 4))\n",
        "\n",
        "        self._dir_stimuli_test = data_path + \"cocosearch/stimuli/test\"\n",
        "        self._dir_target_test = data_path + \"cocosearch/target_\"+ targ_ind_test + \"/test\"\n",
        "        self._dir_saliency_test = data_path + \"cocosearch/saliencymap/test\"\n",
        " \n",
        "\n",
        "    def load_data(self):\n",
        "\n",
        "        test_list_x = _get_file_list(self._dir_stimuli_test)\n",
        "        test_list_y = _get_file_list(self._dir_saliency_test)\n",
        "        test_list_z = _get_file_list(self._dir_target_test)\n",
        "\n",
        "        _check_consistency(zip(test_list_x, test_list_y, test_list_z), self.n_test)\n",
        "\n",
        "        test_set = _fetch_dataset((test_list_x, test_list_y, test_list_z), self._stimuli_size, self._target_size,\n",
        "                                  False, online=True)\n",
        "\n",
        "        return test_set\n",
        "\n",
        "\n",
        "def get_dataset_iterator(phase, dataset, data_path):\n",
        "\n",
        "    \"\"\"\n",
        "    Entry point to make an initializable dataset iterator for either\n",
        "       training or testing a model by calling the respective dataset class.\n",
        "    Args:\n",
        "        phase (str): Holds the current phase, which can be \"train\" or \"test\".\n",
        "        dataset (str): Denotes the dataset to be used during training or the\n",
        "                       suitable resizing procedure when testing a model.\n",
        "        data_path (str): Points to the directory where training or testing\n",
        "                         data instances are stored.\n",
        "    Returns:\n",
        "        iterator: An initializable dataset iterator holding the relevant data.\n",
        "        initializer: An operation required to initialize the correct iterator.\n",
        "    \"\"\"\n",
        "\n",
        "    if phase == \"train\":\n",
        "\n",
        "        current_module = sys.modules[__name__]\n",
        "        class_name = \"%s\" % dataset.upper()\n",
        "\n",
        "        dataset_class = getattr(current_module, class_name)(data_path)\n",
        "        train_set, valid_set = dataset_class.load_data()\n",
        "\n",
        "        iterator = tf.data.Iterator.from_structure(train_set.output_types,\n",
        "                                                   train_set.output_shapes)\n",
        "        next_element = iterator.get_next()\n",
        "\n",
        "        train_init_op = iterator.make_initializer(train_set)\n",
        "        valid_init_op = iterator.make_initializer(valid_set)\n",
        "\n",
        "        return next_element, train_init_op, valid_init_op\n",
        "\n",
        "    if phase == \"test\":\n",
        "\n",
        "        test_class = TEST(dataset, data_path)\n",
        "        test_set = test_class.load_data()\n",
        "\n",
        "        iterator = tf.data.Iterator.from_structure(test_set.output_types,\n",
        "                                                   test_set.output_shapes)\n",
        "        next_element = iterator.get_next()\n",
        "\n",
        "        init_op = iterator.make_initializer(test_set)\n",
        "\n",
        "        return next_element, init_op\n",
        "\n",
        "\n",
        "def postprocess_saliency_map(saliency_map, target_size):\n",
        "    \"\"\"This function resizes and crops a single saliency map to the original\n",
        "       dimensions of the input image. The output is then encoded as a jpeg\n",
        "       file suitable for saving to disk.\n",
        "    Args:\n",
        "        saliency_map (tensor, float32): 3D tensor that holds the values of a\n",
        "                                        saliency map in the range from 0 to 1.\n",
        "        target_size (tensor, int32): 1D tensor that specifies the size to which\n",
        "                                     the saliency map is resized and cropped.\n",
        "    Returns:\n",
        "        tensor, str: A tensor of the saliency map encoded as a jpeg file.\n",
        "    \"\"\"\n",
        "\n",
        "    saliency_map_np = saliency_map * 255.0\n",
        "    saliency_map = _resize_image(saliency_map_np , target_size, True)\n",
        "\n",
        "\n",
        "    saliency_map = tf.round(saliency_map)\n",
        "    saliency_map = tf.cast(saliency_map, tf.uint8)\n",
        "\n",
        "    saliency_map_jpg = tf.image.encode_jpeg(saliency_map, \"grayscale\", 100)\n",
        "\n",
        "    return saliency_map_jpg, saliency_map_np \n",
        "\n",
        "\n",
        "def _fetch_dataset(files, stimuli_size, target_size, shuffle, online=False):\n",
        "\n",
        "    \"\"\"Here the list of file directories is shuffled (only when training),\n",
        "       loaded, batched, and prefetched to ensure high GPU utilization.\n",
        "    Args:\n",
        "        files (list, str): A list that holds the paths to all file instances.\n",
        "        stimuli_size (tuple, int): A tuple that specifies the size to which\n",
        "                                  the search images and saliency maps will be reshaped.\n",
        "        target_size (tuple, int): A tuple that specifies the size to which\n",
        "                                  the target image will be reshaped.\n",
        "        shuffle (bool): Determines whether the dataset will be shuffled or not.\n",
        "        online (bool, optional): Flag that decides whether the batch size must\n",
        "                                 be 1 or can take any value. Defaults to False.\n",
        "    Returns:\n",
        "        object: A dataset object that contains the batched and prefetched data\n",
        "                instances along with their shapes and file paths.\n",
        "    \"\"\"\n",
        "\n",
        "    dataset = tf.data.Dataset.from_tensor_slices(files)\n",
        "\n",
        "    if shuffle:\n",
        "        dataset = dataset.shuffle(len(files[0]))\n",
        "\n",
        "    dataset = dataset.map(lambda *files: _parse_function(files, stimuli_size, target_size),\n",
        "                          num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "\n",
        "    batch_size = 1 if online else config.PARAMS[\"batch_size\"]\n",
        "\n",
        "    dataset = dataset.batch(batch_size)\n",
        "    dataset = dataset.prefetch(5)\n",
        "\n",
        "    return dataset\n",
        "\n",
        "\n",
        "def _parse_function(files, stimuli_size, target_size):\n",
        "    \"\"\"This function reads image data dependent on the image type and\n",
        "       whether it constitutes a stimulus or saliency map. All instances\n",
        "       are then reshaped and padded to yield the target dimensionality.\n",
        "    Args:\n",
        "        files (tuple, str): A tuple with the paths to all file instances.\n",
        "                            The first element contains the stimuli and, if\n",
        "                            present, the second one the ground truth maps.\n",
        "        stimuli_size (tuple, int): A tuple that specifies the size to which\n",
        "                                  the stimuli/saliency map will be reshaped.\n",
        "        target_size (tuple, int): A tuple that specifies the size to which\n",
        "                                  the target will be reshaped.\n",
        "    Returns:\n",
        "        list: A list that holds the image instances along with their\n",
        "              shapes and file paths.\n",
        "    \"\"\"\n",
        "\n",
        "    image_list = []\n",
        "\n",
        "    for count, filename in enumerate(files):\n",
        "\n",
        "        image_str = tf.read_file(filename)\n",
        "        channels = 3 if (count == 0 or count == 2) else 1\n",
        "        image = tf.cond(tf.image.is_jpeg(image_str),\n",
        "                        lambda: tf.image.decode_jpeg(image_str,\n",
        "                                                     channels=channels),\n",
        "                        lambda: tf.image.decode_png(image_str,\n",
        "                                                    channels=channels))\n",
        "        original_size = tf.shape(image)[:2]\n",
        "\n",
        "        if count == 2: #target images\n",
        "            image = _resize_image(image, target_size)\n",
        "        \n",
        "        elif count == 0 or count == 1: #saliency maps and stimuli\n",
        "\n",
        "            image = _resize_image(image, stimuli_size)\n",
        "\n",
        "        image_list.append(image)\n",
        "\n",
        "\n",
        "    image_list.append(original_size)\n",
        "    image_list.append(files)\n",
        "\n",
        "    return image_list\n",
        "\n",
        "\n",
        "def _resize_image(image, target_size, overfull=False):\n",
        "    \"\"\"This resizing procedure preserves the original aspect ratio and might be\n",
        "       followed by padding or cropping. Depending on whether the target size is\n",
        "       smaller or larger than the current image size, the area or bicubic\n",
        "       interpolation method will be utilized.\n",
        "    Args:\n",
        "        image (tensor, uint8): A tensor with the values of an image instance.\n",
        "        target_size (tuple, int): A tuple that specifies the size to which\n",
        "                                  the data will be resized.\n",
        "        overfull (bool, optional): Denotes whether the resulting image will be\n",
        "                                   larger or equal to the specified target\n",
        "                                   size. This is crucial for the following\n",
        "                                   padding or cropping. Defaults to False.\n",
        "    Returns:\n",
        "        tensor, float32: 4D tensor that holds the values of the resized image.\n",
        "    .. seealso:: The reasoning for using either area or bicubic interpolation\n",
        "                 methods is based on the OpenCV documentation recommendations.\n",
        "                 [https://bit.ly/2XAavw0]\n",
        "    \"\"\"\n",
        "\n",
        "    current_size = tf.shape(image)[:2]\n",
        "\n",
        "    height_ratio = target_size[0] / current_size[0]\n",
        "    width_ratio = target_size[1] / current_size[1]\n",
        "\n",
        "    if overfull:\n",
        "        target_ratio = tf.maximum(height_ratio, width_ratio)\n",
        "    else:\n",
        "        target_ratio = tf.minimum(height_ratio, width_ratio)\n",
        "\n",
        "    target_size = tf.cast(current_size, tf.float64) * target_ratio\n",
        "    target_size = tf.cast(tf.round(target_size), tf.int32)\n",
        "\n",
        "\n",
        "    shrinking = tf.cond(tf.logical_or(current_size[0] > target_size[0],\n",
        "                                      current_size[1] > target_size[1]),\n",
        "                        lambda: tf.constant(True),\n",
        "                        lambda: tf.constant(False))\n",
        "\n",
        "    image = tf.expand_dims(image, 0)\n",
        "\n",
        "    image = tf.cond(shrinking,\n",
        "                    lambda: tf.image.resize_area(image, target_size,\n",
        "                                                 align_corners=True),\n",
        "                    lambda: tf.image.resize_bicubic(image, target_size,\n",
        "                                                    align_corners=True))\n",
        "\n",
        "    image = tf.clip_by_value(image[0], 0.0, 255.0)\n",
        "\n",
        "    return image\n",
        "\n",
        "\n",
        "def _get_file_list(data_path):\n",
        "    \"\"\"This function detects all image files within the specified parent\n",
        "       directory for either training or testing. The path content cannot\n",
        "       be empty, otherwise an error occurs.\n",
        "    Args:\n",
        "        data_path (str): Points to the directory where training or testing\n",
        "                         data instances are stored.\n",
        "    Returns:\n",
        "        list, str: A sorted list that holds the paths to all file instances.\n",
        "    \"\"\"\n",
        "\n",
        "    data_list = []\n",
        "\n",
        "    if os.path.isfile(data_path):\n",
        "        data_list.append(data_path)\n",
        "    else:\n",
        "        for subdir, dirs, files in os.walk(data_path):\n",
        "            for file in files:\n",
        "                if file.lower().endswith((\".png\", \".jpg\", \".jpeg\")):\n",
        "                    data_list.append(os.path.join(subdir, file))\n",
        "\n",
        "    data_list.sort()\n",
        "\n",
        "    if not data_list:\n",
        "        raise FileNotFoundError(\"No data was found\")\n",
        "\n",
        "    return data_list\n",
        "\n",
        "\n",
        "def _check_consistency(zipped_file_lists, n_total_files):\n",
        "    \"\"\"A consistency check that makes sure all files could successfully be\n",
        "       found and stimuli names correspond to the ones of ground truth maps.\n",
        "    Args:\n",
        "        zipped_file_lists (tuple, str): A tuple of train and valid path names.\n",
        "        n_total_files (int): The total number of files expected in the list.\n",
        "    \"\"\"\n",
        "\n",
        "    assert len(list(zipped_file_lists)) == n_total_files, \"Files are missing\"\n",
        "\n",
        "    for file_tuple in zipped_file_lists:\n",
        "        file_names = [os.path.basename(entry) for entry in list(file_tuple)]\n",
        "        file_names = [os.path.splitext(entry)[0] for entry in file_names]\n",
        "        file_names = [entry.replace(\"_fixMap\", \"\") for entry in file_names]\n",
        "        file_names = [entry.replace(\"_fixPts\", \"\") for entry in file_names]\n",
        "\n",
        "        assert len(set(file_names)) == 1, \"File name mismatch\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "1lmiCOR5-Bg6"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "class loss:\n",
        "  def kld(y_true, y_pred, eps=1e-7):\n",
        "      \"\"\"This function computes the Kullback-Leibler divergence between ground\n",
        "        truth saliency maps and their predictions. Values are first divided by\n",
        "        their sum for each image to yield a distribution that adds to 1.\n",
        "      Args:\n",
        "          y_true (tensor, float32): A 4d tensor that holds the ground truth\n",
        "                                    saliency maps with values between 0 and 255.\n",
        "          y_pred (tensor, float32): A 4d tensor that holds the predicted saliency\n",
        "                                    maps with values between 0 and 1.\n",
        "          eps (scalar, float, optional): A small factor to avoid numerical\n",
        "                                        instabilities. Defaults to 1e-7.\n",
        "      Returns:\n",
        "          tensor, float32: A 0D tensor that holds the averaged error.\n",
        "      \"\"\"\n",
        "\n",
        "      sum_per_image = tf.reduce_sum(y_true, axis=(1, 2, 3), keep_dims=True)\n",
        "      y_true /= eps + sum_per_image\n",
        "\n",
        "      sum_per_image = tf.reduce_sum(y_pred, axis=(1, 2, 3), keep_dims=True)\n",
        "      y_pred /= eps + sum_per_image\n",
        "\n",
        "      loss = y_true * tf.log(eps + y_true / (eps + y_pred))\n",
        "      loss = tf.reduce_sum(loss, axis=(1, 2, 3))\n",
        "\n",
        "      return loss\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "OBYhAzkhwu5z"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import tensorflow as tf\n",
        "from tensorflow.python.tools import freeze_graph\n",
        "from tensorflow.tools.graph_transforms import TransformGraph\n",
        "\n",
        "\n",
        "class MSINET:\n",
        "    \"\"\"The class representing the MSI-Net based on the VGG16 model. It\n",
        "       implements a definition of the computational graph, as well as\n",
        "       functions related to network training.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self):\n",
        "        self._output = None\n",
        "        self._mapping = {}\n",
        "\n",
        "        if config.PARAMS[\"device\"] == \"gpu\":\n",
        "            self._data_format = \"channels_first\"\n",
        "            self._channel_axis = 1\n",
        "            self._dims_axis = (2, 3)\n",
        "        elif config.PARAMS[\"device\"] == \"cpu\":\n",
        "            self._data_format = \"channels_last\"\n",
        "            self._channel_axis = 3\n",
        "            self._dims_axis = (1, 2)\n",
        "\n",
        "    def _encoder(self, images):\n",
        "        \"\"\"The encoder of the model consists of a pretrained VGG16 architecture\n",
        "           with 13 convolutional layers. All dense layers are discarded and the\n",
        "           last 3 layers are dilated at a rate of 2 to account for the omitted\n",
        "           downsampling. Finally, the activations from 3 layers are combined.\n",
        "        Args:\n",
        "            images (tensor, float32): A 4D tensor that holds the RGB image\n",
        "                                      batches used as input to the network.\n",
        "        \"\"\"\n",
        "\n",
        "        imagenet_mean = tf.constant([103.939, 116.779, 123.68])\n",
        "        imagenet_mean = tf.reshape(imagenet_mean, [1, 1, 1, 3])\n",
        "\n",
        "        images -= imagenet_mean\n",
        "\n",
        "        if self._data_format == \"channels_first\":\n",
        "            images = tf.transpose(images, (0, 3, 1, 2))\n",
        "\n",
        "        layer01 = tf.layers.conv2d(images, 64, 3,\n",
        "                                   padding=\"same\",\n",
        "                                   activation=tf.nn.relu,\n",
        "                                   data_format=self._data_format,\n",
        "                                   name=\"conv1/conv1_1\", reuse=tf.AUTO_REUSE)\n",
        "\n",
        "        layer02 = tf.layers.conv2d(layer01, 64, 3,\n",
        "                                   padding=\"same\",\n",
        "                                   activation=tf.nn.relu,\n",
        "                                   data_format=self._data_format,\n",
        "                                   name=\"conv1/conv1_2\", reuse=tf.AUTO_REUSE)\n",
        "\n",
        "        layer03 = tf.layers.max_pooling2d(layer02, 2, 2,\n",
        "                                          data_format=self._data_format)\n",
        "\n",
        "        layer04 = tf.layers.conv2d(layer03, 128, 3,\n",
        "                                   padding=\"same\",\n",
        "                                   activation=tf.nn.relu,\n",
        "                                   data_format=self._data_format,\n",
        "                                   name=\"conv2/conv2_1\", reuse=tf.AUTO_REUSE)\n",
        "\n",
        "        layer05 = tf.layers.conv2d(layer04, 128, 3,\n",
        "                                   padding=\"same\",\n",
        "                                   activation=tf.nn.relu,\n",
        "                                   data_format=self._data_format,\n",
        "                                   name=\"conv2/conv2_2\", reuse=tf.AUTO_REUSE)\n",
        "\n",
        "        layer06 = tf.layers.max_pooling2d(layer05, 2, 2,\n",
        "                                          data_format=self._data_format)\n",
        "\n",
        "        layer07 = tf.layers.conv2d(layer06, 256, 3,\n",
        "                                   padding=\"same\",\n",
        "                                   activation=tf.nn.relu,\n",
        "                                   data_format=self._data_format,\n",
        "                                   name=\"conv3/conv3_1\", reuse=tf.AUTO_REUSE)\n",
        "\n",
        "        layer08 = tf.layers.conv2d(layer07, 256, 3,\n",
        "                                   padding=\"same\",\n",
        "                                   activation=tf.nn.relu,\n",
        "                                   data_format=self._data_format,\n",
        "                                   name=\"conv3/conv3_2\", reuse=tf.AUTO_REUSE)\n",
        "\n",
        "        layer09 = tf.layers.conv2d(layer08, 256, 3,\n",
        "                                   padding=\"same\",\n",
        "                                   activation=tf.nn.relu,\n",
        "                                   data_format=self._data_format,\n",
        "                                   name=\"conv3/conv3_3\", reuse=tf.AUTO_REUSE)\n",
        "\n",
        "        #layer09 = tf.layers.dropout(layer09, rate=0.5)\n",
        "\n",
        "        layer10 = tf.layers.max_pooling2d(layer09, 2, 2,\n",
        "                                          data_format=self._data_format)\n",
        "\n",
        "        layer11 = tf.layers.conv2d(layer10, 512, 3,\n",
        "                                   padding=\"same\",\n",
        "                                   activation=tf.nn.relu,\n",
        "                                   data_format=self._data_format,\n",
        "                                   name=\"conv4/conv4_1\", reuse=tf.AUTO_REUSE)\n",
        "\n",
        "        layer12 = tf.layers.conv2d(layer11, 512, 3,\n",
        "                                   padding=\"same\",\n",
        "                                   activation=tf.nn.relu,\n",
        "                                   data_format=self._data_format,\n",
        "                                   name=\"conv4/conv4_2\", reuse=tf.AUTO_REUSE)\n",
        "\n",
        "        layer13 = tf.layers.conv2d(layer12, 512, 3,\n",
        "                                   padding=\"same\",\n",
        "                                   activation=tf.nn.relu,\n",
        "                                   data_format=self._data_format,\n",
        "                                   name=\"conv4/conv4_3\", reuse=tf.AUTO_REUSE)\n",
        "\n",
        "        #layer13 = tf.layers.dropout(layer13, rate=0.5)\n",
        "\n",
        "        layer14 = tf.layers.max_pooling2d(layer13, 2, 1,\n",
        "                                          padding=\"same\",\n",
        "                                          data_format=self._data_format)\n",
        "\n",
        "        layer15 = tf.layers.conv2d(layer14, 512, 3,\n",
        "                                   padding=\"same\",\n",
        "                                   activation=tf.nn.relu,\n",
        "                                   dilation_rate=2,\n",
        "                                   data_format=self._data_format,\n",
        "                                   name=\"conv5/conv5_1\", reuse=tf.AUTO_REUSE)\n",
        "\n",
        "        layer16 = tf.layers.conv2d(layer15, 512, 3,\n",
        "                                   padding=\"same\",\n",
        "                                   activation=tf.nn.relu,\n",
        "                                   dilation_rate=2,\n",
        "                                   data_format=self._data_format,\n",
        "                                   name=\"conv5/conv5_2\", reuse=tf.AUTO_REUSE)\n",
        "\n",
        "        layer17 = tf.layers.conv2d(layer16, 512, 3,\n",
        "                                   padding=\"same\",\n",
        "                                   activation=tf.nn.relu,\n",
        "                                   dilation_rate=2,\n",
        "                                   data_format=self._data_format,\n",
        "                                   name=\"conv5/conv5_3\", reuse=tf.AUTO_REUSE)\n",
        "\n",
        "        #layer17 = tf.layers.dropout(layer17, rate=0.5)\n",
        "\n",
        "        layer18 = tf.layers.max_pooling2d(layer17, 2, 1,\n",
        "                                          padding=\"same\",\n",
        "                                          data_format=self._data_format)\n",
        "\n",
        "        encoder_output = tf.concat([layer10, layer14, layer18],\n",
        "                                   axis=self._channel_axis, name='concat')\n",
        "\n",
        "        self._output = encoder_output\n",
        "\n",
        "    def _aspp(self, features):\n",
        "        \"\"\"The ASPP module samples information at multiple spatial scales in\n",
        "           parallel via convolutional layers with different dilation factors.\n",
        "           The activations are then combined with global scene context and\n",
        "           represented as a common tensor.\n",
        "        Args:\n",
        "            features (tensor, float32): A 4D tensor that holds the features\n",
        "                                        produced by the encoder module.\n",
        "        \"\"\"\n",
        "\n",
        "        branch1 = tf.layers.conv2d(features, 256, 1,\n",
        "                                   padding=\"same\",\n",
        "                                   activation=tf.nn.relu,\n",
        "                                   data_format=self._data_format,\n",
        "                                   name=\"aspp/conv1_1\", reuse=tf.AUTO_REUSE)\n",
        "\n",
        "        branch2 = tf.layers.conv2d(features, 256, 3,\n",
        "                                   padding=\"same\",\n",
        "                                   activation=tf.nn.relu,\n",
        "                                   dilation_rate=4,\n",
        "                                   data_format=self._data_format,\n",
        "                                   name=\"aspp/conv1_2\", reuse=tf.AUTO_REUSE)\n",
        "\n",
        "        branch3 = tf.layers.conv2d(features, 256, 3,\n",
        "                                   padding=\"same\",\n",
        "                                   activation=tf.nn.relu,\n",
        "                                   dilation_rate=8,\n",
        "                                   data_format=self._data_format,\n",
        "                                   name=\"aspp/conv1_3\", reuse=tf.AUTO_REUSE)\n",
        "\n",
        "        branch4 = tf.layers.conv2d(features, 256, 3,\n",
        "                                   padding=\"same\",\n",
        "                                   activation=tf.nn.relu,\n",
        "                                   dilation_rate=12,\n",
        "                                   data_format=self._data_format,\n",
        "                                   name=\"aspp/conv1_4\", reuse=tf.AUTO_REUSE)\n",
        "\n",
        "        branch5 = tf.reduce_mean(features,\n",
        "                                 axis=self._dims_axis,\n",
        "                                 keepdims=True)\n",
        "\n",
        "        branch5 = tf.layers.conv2d(branch5, 256, 1,\n",
        "                                   padding=\"valid\",\n",
        "                                   activation=tf.nn.relu,\n",
        "                                   data_format=self._data_format,\n",
        "                                   name=\"aspp/conv1_5\", reuse=tf.AUTO_REUSE)\n",
        "\n",
        "        shape = tf.shape(features)\n",
        "\n",
        "        branch5 = self._upsample(branch5, shape, 1)\n",
        "\n",
        "        context = tf.concat([branch1, branch2, branch3, branch4, branch5],\n",
        "                            axis=self._channel_axis)\n",
        "\n",
        "        aspp_output = tf.layers.conv2d(context, 256, 1,\n",
        "                                       padding=\"same\",\n",
        "                                       activation=tf.nn.relu,\n",
        "                                       data_format=self._data_format,\n",
        "                                       name=\"aspp/conv2\", reuse=tf.AUTO_REUSE)\n",
        "        self._output = aspp_output\n",
        "\n",
        "    def _decoder(self, features):\n",
        "        \"\"\"The decoder model applies a series of 3 upsampling blocks that each\n",
        "           performs bilinear upsampling followed by a 3x3 convolution to avoid\n",
        "           checkerboard artifacts in the image space. Unlike all other layers,\n",
        "           the output of the model is not modified by a ReLU.\n",
        "        Args:\n",
        "            features (tensor, float32): A 4D tensor that holds the features\n",
        "                                        produced by the ASPP module.\n",
        "        \"\"\"\n",
        "\n",
        "        shape = tf.shape(features)\n",
        "\n",
        "        layer1 = self._upsample(features, shape, 2)\n",
        "\n",
        "        layer2 = tf.layers.conv2d(layer1, 128, 3,\n",
        "                                  padding=\"same\",\n",
        "                                  activation=tf.nn.relu,\n",
        "                                  data_format=self._data_format,\n",
        "                                  name=\"decoder/conv1\", reuse=tf.AUTO_REUSE)\n",
        "\n",
        "        shape = tf.shape(layer2)\n",
        "\n",
        "        layer3 = self._upsample(layer2, shape, 2)\n",
        "\n",
        "        layer4 = tf.layers.conv2d(layer3, 64, 3,\n",
        "                                  padding=\"same\",\n",
        "                                  activation=tf.nn.relu,\n",
        "                                  data_format=self._data_format,\n",
        "                                  name=\"decoder/conv2\", reuse=tf.AUTO_REUSE)\n",
        "\n",
        "        shape = tf.shape(layer4)\n",
        "\n",
        "        layer5 = self._upsample(layer4, shape, 2)\n",
        "\n",
        "        layer6 = tf.layers.conv2d(layer5, 32, 3,\n",
        "                                  padding=\"same\",\n",
        "                                  activation=tf.nn.relu,\n",
        "                                  data_format=self._data_format,\n",
        "                                  name=\"decoder/conv3\", reuse=tf.AUTO_REUSE)\n",
        "\n",
        "        decoder_output = tf.layers.conv2d(layer6, 1, 3,\n",
        "                                          padding=\"same\",\n",
        "                                          data_format=self._data_format,\n",
        "                                          name=\"decoder/conv4\")\n",
        "\n",
        "        if self._data_format == \"channels_first\":\n",
        "            decoder_output = tf.transpose(decoder_output, (0, 2, 3, 1))\n",
        "\n",
        "        self._output = decoder_output\n",
        "\n",
        "    def _overlay(self, feature1, feature2):\n",
        "        \"\"\"This function convolves the features extracted from stimuli\n",
        "            and target streams. It uses a convolutional layer \n",
        "            and uses features of the target image as convolution filters\n",
        "            for stimulus features.\n",
        "        Args:\n",
        "            feature1 (tensor, float32): A 4D tensor that contains the features of stimuli.\n",
        "            feature2 (tensor, int32): A 4D tensor that contains the features of target.\n",
        "        Returns:\n",
        "            tensor, float32: A 4D tensor that holds the output of convolving the two streams.\n",
        "        \"\"\"\n",
        "\n",
        "        shape_channel = (feature2.get_shape())[1]\n",
        "        feature2 = tf.squeeze(feature2, axis=0)\n",
        "        feature2 = tf.transpose(feature2, [1, 2, 0])\n",
        "        feature2 = tf.expand_dims(feature2, axis=3)\n",
        "        feature2 = tf.concat(shape_channel * [feature2], axis=3)\n",
        "        feature2 = tf.cast(feature2, dtype=tf.float32)\n",
        "\n",
        "        overlay_output = tf.nn.conv2d(feature1, feature2, strides=[1, 1, 1, 1], padding='SAME', data_format='NCHW',\n",
        "                                      name=\"overlay\")\n",
        "\n",
        "        self._output = overlay_output\n",
        "\n",
        "    def _upsample(self, stack, shape, factor):\n",
        "        \"\"\"This function resizes the input to a desired shape via the\n",
        "           bilinear upsampling method.\n",
        "        Args:\n",
        "            stack (tensor, float32): A 4D tensor with the function input.\n",
        "            shape (tensor, int32): A 1D tensor with the reference shape.\n",
        "            factor (scalar, int): An integer denoting the upsampling factor.\n",
        "        Returns:\n",
        "            tensor, float32: A 4D tensor that holds the activations after\n",
        "                             bilinear upsampling of the input.\n",
        "        \"\"\"\n",
        "\n",
        "        if self._data_format == \"channels_first\":\n",
        "            stack = tf.transpose(stack, (0, 2, 3, 1))\n",
        "\n",
        "        stack = tf.image.resize_bilinear(stack, (shape[self._dims_axis[0]] * factor,\n",
        "                                                 shape[self._dims_axis[1]] * factor))\n",
        "\n",
        "        if self._data_format == \"channels_first\":\n",
        "            stack = tf.transpose(stack, (0, 3, 1, 2))\n",
        "\n",
        "        return stack\n",
        "\n",
        "    def _normalize(self, maps, eps=1e-7):\n",
        "        \"\"\"This function normalizes the output values to a range\n",
        "           between 0 and 1 per saliency map.\n",
        "        Args:\n",
        "            maps (tensor, float32): A 4D tensor that holds the model output.\n",
        "            eps (scalar, float, optional): A small factor to avoid numerical\n",
        "                                           instabilities. Defaults to 1e-7.\n",
        "        \"\"\"\n",
        "\n",
        "        min_per_image = tf.reduce_min(maps, axis=(1, 2, 3), keep_dims=True)\n",
        "        maps -= min_per_image\n",
        "\n",
        "        max_per_image = tf.reduce_max(maps, axis=(1, 2, 3), keep_dims=True)\n",
        "        maps = tf.divide(maps, eps + max_per_image, name=\"output\")\n",
        "\n",
        "        self._output = maps\n",
        "\n",
        "    def _pretraining(self):\n",
        "        \"\"\"The first 26 variables of the model here are based on the VGG16\n",
        "           network. Therefore, their names are matched to the ones of the\n",
        "           pretrained VGG16 checkpoint for correct initialization.\n",
        "        \"\"\"\n",
        "\n",
        "        for var in tf.global_variables()[:26]:\n",
        "            key = var.name.split(\"/\", 1)[1]\n",
        "            key = key.replace(\"kernel:0\", \"weights\")\n",
        "            key = key.replace(\"bias:0\", \"biases\")\n",
        "            self._mapping[key] = var\n",
        "\n",
        "    def forward(self, stimuli):\n",
        "        \"\"\"Public method to forward RGB images through the feature \n",
        "            extraction parts of the network.\n",
        "        Args:\n",
        "            images (tensor, float32): A 4D tensor that holds the values of the\n",
        "                                      raw input images.\n",
        "        Returns:\n",
        "            tensor, float32: A 4D tensor that holds the values of the\n",
        "                             extracted features from ASPP module.\n",
        "        \"\"\"\n",
        "\n",
        "        self._encoder(stimuli)\n",
        "        self._aspp(self._output)\n",
        "\n",
        "        return self._output\n",
        "\n",
        "    def one_stream(self, stimuli_features):\n",
        "        \"\"\"creates the output of the one stream network, which is the\n",
        "            predicted fixation density map.\n",
        "        Args:\n",
        "            stimuli_features (tensor, float32): A 4D tensor containing the features of stimuli.\n",
        "        Returns:\n",
        "            tensor, float32: A 4D tensor that holds the values of the\n",
        "                             predicted saliency maps.\n",
        "        \"\"\"\n",
        "        self._decoder(self._output)\n",
        "        self._normalize(self._output)\n",
        "\n",
        "        return self._output\n",
        "\n",
        "    def output_stream(self, stimuli_features, target_features):\n",
        "        \"\"\"creates the output of the two stream network, which is the\n",
        "            predicted fixation density map.\n",
        "        Args:\n",
        "            stimuli_features (tensor, float32): A 4D tensor containing the features of stimuli.\n",
        "            target_features (tensor, float32): A 4D tensor containing the features of target.\n",
        "        Returns:\n",
        "            tensor, float32: A 4D tensor that holds the values of the\n",
        "                             predicted saliency maps.\n",
        "        \"\"\"\n",
        "        self._overlay(stimuli_features, target_features)\n",
        "        self._decoder(self._output)\n",
        "        self._normalize(self._output)\n",
        "\n",
        "        return self._output\n",
        "\n",
        "    def train(self, ground_truth, predicted_maps, learning_rate):\n",
        "        \"\"\"Public method to define the loss function and optimization\n",
        "           algorithm for training the model.\n",
        "        Args:\n",
        "            ground_truth (tensor, float32): A 4D tensor with the ground truth.\n",
        "            predicted_maps (tensor, float32): A 4D tensor with the predictions.\n",
        "            learning_rate (scalar, float): Defines the learning rate.\n",
        "        Returns:\n",
        "            object: The optimizer element used to train the model.\n",
        "            tensor, float32: A 0D tensor that holds the averaged error.\n",
        "        \"\"\"\n",
        "\n",
        "        error = loss.kld(ground_truth, predicted_maps)\n",
        "\n",
        "        global_step = tf.train.get_or_create_global_step()\n",
        "\n",
        "        adjusted_global_step = global_step\n",
        "\n",
        "        base_learning_rate = learning_rate\n",
        "\n",
        "        '''learning_rate = tf.train.polynomial_decay(\n",
        "            base_learning_rate,\n",
        "            adjusted_global_step,\n",
        "            config.PARAMS[\"n_training_steps\"],\n",
        "            end_learning_rate=0,\n",
        "            power=config.PARAMS[\"learning_power\"])'''\n",
        "\n",
        "        optimizer = tf.train.AdamOptimizer(learning_rate)              \n",
        "        #optimizer = tf.train.MomentumOptimizer(learning_rate, config.PARAMS[\"momentum\"])\n",
        "        optimizer = optimizer.minimize(error)  \n",
        "        return optimizer, error\n",
        "\n",
        "    def save(self, saver, sess, dataset, path, device):\n",
        "        \"\"\"This saves a model checkpoint to disk and creates\n",
        "           the folder if it doesn't exist yet.\n",
        "        Args:\n",
        "            saver (object): An object for saving the model.\n",
        "            sess (object): The current TF training session.\n",
        "            path (str): The path used for saving the model.\n",
        "            device (str): Represents either \"cpu\" or \"gpu\".\n",
        "        \"\"\"\n",
        "\n",
        "        os.makedirs(path, exist_ok=True)\n",
        "\n",
        "        saver.save(sess, path + \"model_%s_%s.ckpt\" % (dataset, device),\n",
        "                   write_meta_graph=False, write_state=False)\n",
        "\n",
        "    def restore(self, sess, dataset, paths, device):\n",
        "        \"\"\"This function allows continued training from a prior checkpoint and\n",
        "           training from scratch with the pretrained VGG16 weights. In case the\n",
        "           dataset is either CAT2000 or MIT1003, a prior checkpoint based on\n",
        "           the SALICON dataset is required.\n",
        "        Args:\n",
        "            sess (object): The current TF training session.\n",
        "            dataset ([type]): The dataset used for training.\n",
        "            paths (dict, str): A dictionary with all path elements.\n",
        "            device (str): Represents either \"cpu\" or \"gpu\".\n",
        "        Returns:\n",
        "            object: A saver object for saving the model.\n",
        "        \"\"\"\n",
        "\n",
        "        model_name = \"model_%s_%s\" % (dataset, device)\n",
        "        salicon_name = \"model_salicon_%s\" % device\n",
        "        vgg16_name = \"vgg16_hybrid\"\n",
        "\n",
        "        ext1 = \".ckpt.data-00000-of-00001\"\n",
        "        ext2 = \".ckpt.index\"\n",
        "\n",
        "        saver = tf.train.Saver()\n",
        "\n",
        "        if os.path.isfile(paths[\"latest\"] + model_name + ext1) and \\\n",
        "                os.path.isfile(paths[\"latest\"] + model_name + ext2):\n",
        "            saver.restore(sess, paths[\"latest\"] + model_name + \".ckpt\")\n",
        "        else:\n",
        "            self._pretraining()\n",
        "            loader = tf.train.Saver(self._mapping)\n",
        "            loader.restore(sess, paths[\"weights\"] + vgg16_name + \".ckpt\")\n",
        "\n",
        "        return saver\n",
        "\n",
        "    def optimize(self, sess, dataset, path, device):\n",
        "        \"\"\"The best performing model is frozen, optimized for inference\n",
        "           by removing unneeded training operations, and written to disk.\n",
        "        Args:\n",
        "            sess (object): The current TF training session.\n",
        "            path (str): The path used for saving the model.\n",
        "            device (str): Represents either \"cpu\" or \"gpu\".\n",
        "        .. seealso:: https://bit.ly/2VBBdqQ and https://bit.ly/2W7YqBa\n",
        "        \"\"\"\n",
        "\n",
        "        model_name = \"model_%s_%s\" % (dataset, device)\n",
        "        model_path = path + model_name\n",
        "\n",
        "        tf.train.write_graph(sess.graph.as_graph_def(),\n",
        "                             path, model_name + \".pbtxt\")\n",
        "\n",
        "        freeze_graph.freeze_graph(model_path + \".pbtxt\", \"\", False,\n",
        "                                  model_path + \".ckpt\", \"output\",\n",
        "                                  \"save/restore_all\", \"save/Const:0\",\n",
        "                                  model_path + \".pb\", True, \"\")\n",
        "\n",
        "        os.remove(model_path + \".pbtxt\")\n",
        "\n",
        "        graph_def = tf.GraphDef()\n",
        "\n",
        "        with tf.gfile.Open(model_path + \".pb\", \"rb\") as file:\n",
        "            graph_def.ParseFromString(file.read())\n",
        "\n",
        "        transforms = [\"remove_nodes(op=Identity)\",\n",
        "                      \"merge_duplicate_nodes\",\n",
        "                      \"strip_unused_nodes\",\n",
        "                      \"fold_constants(ignore_errors=true)\"]\n",
        "\n",
        "        optimized_graph_def = TransformGraph(graph_def,\n",
        "                                             [\"input\"],\n",
        "                                             [\"output\"],\n",
        "                                             transforms)\n",
        "\n",
        "        tf.train.write_graph(optimized_graph_def,\n",
        "                             logdir=path,\n",
        "                             as_text=False,\n",
        "                             name=model_name + \".pb\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "uXqZZHnpw7WT"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import time\n",
        "from datetime import timedelta\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.ticker as plticker\n",
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "\n",
        "class History:\n",
        "    \"\"\"This class represents the training history of a model. It can load the\n",
        "       prior history when training continues, keeps track of the training and\n",
        "       validation error, and finally plots them as a curve after each epoch.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, n_train_batches, n_valid_batches,\n",
        "                 dataset, path, device):\n",
        "        self.train_history = []\n",
        "        self.valid_history = []\n",
        "\n",
        "        self._prior_epochs = 0\n",
        "\n",
        "        self._train_error = 0\n",
        "        self._valid_error = 0\n",
        "\n",
        "        self._n_train_batches = n_train_batches\n",
        "        self._n_valid_batches = n_valid_batches\n",
        "\n",
        "        self._path = path\n",
        "        self._id = (dataset, device)\n",
        "\n",
        "        self._get_prior_history()\n",
        "\n",
        "    def _get_prior_history(self):\n",
        "        if os.path.isfile(self._path + \"train_%s_%s.txt\" % self._id):\n",
        "            with open(self._path + \"train_%s_%s.txt\" % self._id, \"r\") as file:\n",
        "                for line in file.readlines():\n",
        "                    self.train_history.append(float(line))\n",
        "\n",
        "        if os.path.isfile(self._path + \"valid_%s_%s.txt\" % self._id):\n",
        "            with open(self._path + \"valid_%s_%s.txt\" % self._id, \"r\") as file:\n",
        "                for line in file.readlines():\n",
        "                    self.valid_history.append(float(line))\n",
        "\n",
        "        self.prior_epochs = len(self.train_history)\n",
        "\n",
        "    def update_train_step(self, train_error):\n",
        "        self._train_error += train_error\n",
        "\n",
        "    def update_valid_step(self, valid_error):\n",
        "        self._valid_error += valid_error\n",
        "\n",
        "    def get_mean_train_error(self, reset=True):\n",
        "        mean_train_error = self._train_error / self._n_train_batches\n",
        "\n",
        "        if reset:\n",
        "            self._train_error = 0\n",
        "\n",
        "        return mean_train_error\n",
        "\n",
        "    def get_mean_valid_error(self, reset=True):\n",
        "        mean_valid_error = self._valid_error / self._n_valid_batches\n",
        "\n",
        "        if reset:\n",
        "            self._valid_error = 0\n",
        "\n",
        "        return mean_valid_error\n",
        "\n",
        "    def save_history(self):\n",
        "        mean_train_loss = self.get_mean_train_error(False)\n",
        "        mean_valid_loss = self.get_mean_valid_error(False)\n",
        "\n",
        "        self.train_history.append(mean_train_loss)\n",
        "        self.valid_history.append(mean_valid_loss)\n",
        "\n",
        "        os.makedirs(self._path, exist_ok=True)\n",
        "\n",
        "        with open(self._path + \"train_%s_%s.txt\" % self._id, \"a\") as file:\n",
        "            file.write(\"%f\\n\" % self.train_history[-1])\n",
        "\n",
        "        with open(self._path + \"valid_%s_%s.txt\" % self._id, \"a\") as file:\n",
        "            file.write(\"%f\\n\" % self.valid_history[-1])\n",
        "\n",
        "        if len(self.train_history) > 1:\n",
        "            axes = plt.figure().gca()\n",
        "\n",
        "            x_range = np.arange(1, len(self.train_history) + 1)\n",
        "\n",
        "            plt.plot(x_range, self.train_history, label=\"train\", linewidth=2)\n",
        "            plt.plot(x_range, self.valid_history, label=\"valid\", linewidth=2)\n",
        "\n",
        "            plt.legend()\n",
        "            plt.xlabel(\"epochs\")\n",
        "            plt.ylabel(\"error\")\n",
        "\n",
        "            locations = plticker.MultipleLocator(base=1.0)\n",
        "            axes.xaxis.set_major_locator(locations)\n",
        "\n",
        "            plt.savefig(self._path + \"curve_%s_%s.png\" % self._id)\n",
        "            plt.close()\n",
        "\n",
        "\n",
        "class Progbar:\n",
        "    \"\"\"This class represents a progress bar for the terminal that visualizes\n",
        "       the training progress for each epoch, estimated time of accomplishment,\n",
        "       and then summarizes the training and validation loss together with the\n",
        "       elapsed time.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, n_train_data, n_train_batches,\n",
        "                 batch_size, n_epochs, prior_epochs):\n",
        "        self._train_time = 0\n",
        "        self._valid_time = 0\n",
        "\n",
        "        self._start_time = time.time()\n",
        "\n",
        "        self._batch_size = batch_size\n",
        "\n",
        "        self._n_train_data = n_train_data\n",
        "        self._n_train_batches = n_train_batches\n",
        "\n",
        "        self._target_epoch = str(n_epochs + prior_epochs).zfill(2)\n",
        "        self._current_epoch = str(prior_epochs + 1).zfill(2)\n",
        "\n",
        "    def _flush(self):\n",
        "        self._train_time = 0\n",
        "        self._valid_time = 0\n",
        "\n",
        "        self._start_time = time.time()\n",
        "\n",
        "        current_epoch_int = int(self._current_epoch) + 1\n",
        "        self._current_epoch = str(current_epoch_int).zfill(2)\n",
        "\n",
        "    def update_train_step(self, current_batch):\n",
        "        current_batch += 1\n",
        "\n",
        "        self._train_time = time.time() - self._start_time\n",
        "        batch_train_time = self._train_time / current_batch\n",
        "\n",
        "        eta = (self._n_train_batches - current_batch) * batch_train_time\n",
        "        eta = str(timedelta(seconds=np.ceil(eta)))\n",
        "\n",
        "        progress_line = \"=\" * (20 * current_batch // self._n_train_batches)\n",
        "\n",
        "        current_instance = current_batch * self._batch_size\n",
        "        current_instance = np.clip(current_instance, 0, self._n_train_data)\n",
        "\n",
        "        progress_frac = \"%i/%i\" % (current_instance, self._n_train_data)\n",
        "\n",
        "        information = (self._current_epoch, self._target_epoch,\n",
        "                       progress_line, progress_frac, eta)\n",
        "\n",
        "        progbar_output = \"Epoch %s/%s [%-20s] %s (ETA: %s)\" % information\n",
        "\n",
        "        print(progbar_output, end=\"\\r\", flush=True)\n",
        "\n",
        "    def update_valid_step(self):\n",
        "        self._valid_time = time.time() - self._start_time - self._train_time\n",
        "\n",
        "    def write_summary(self, mean_train_loss, mean_valid_loss):\n",
        "        train_time = str(timedelta(seconds=np.ceil(self._train_time)))\n",
        "        valid_time = str(timedelta(seconds=np.ceil(self._valid_time)))\n",
        "\n",
        "        train_information = (mean_train_loss, train_time)\n",
        "        valid_information = (mean_valid_loss, valid_time)\n",
        "\n",
        "        train_output = \"\\n\\tTrain loss: %.6f (%s)\" % train_information\n",
        "        valid_output = \"\\tValid loss: %.6f (%s)\" % valid_information\n",
        "\n",
        "        print(train_output, flush=True)\n",
        "        print(valid_output, flush=True)\n",
        "\n",
        "        self._flush()\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "l_QUJXNu-Tzo",
        "outputId": "a7c1470b-e541-4612-d9af-1d1ec3cb4835"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            ">> Start testing with COCOSEARCH gpu model...\n"
          ]
        }
      ],
      "source": [
        "import argparse\n",
        "import os\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import random\n",
        "import sys\n",
        "import cv2\n",
        "\n",
        "seed_value = 32\n",
        "os.environ['PYTHONHASHSEED'] = str(seed_value)\n",
        "random.seed(seed_value)\n",
        "np.random.seed(seed_value)\n",
        "tf.random.set_random_seed(seed_value)\n",
        "\n",
        "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"2\"\n",
        "tf.logging.set_verbosity(tf.logging.ERROR)\n",
        "\n",
        "def define_paths(current_path):\n",
        "    \"\"\"A helper function to define all relevant path elements for the\n",
        "       locations of data, weights, and the results from either training\n",
        "       or testing a model.\n",
        "    Args:\n",
        "        current_path (str): The path string of this script.\n",
        "        args (object): A namespace object with values from command line.\n",
        "    Returns:\n",
        "        dict: A dictionary with all path elements.\n",
        "    \"\"\"\n",
        "\n",
        "    data_path = current_path \n",
        "    results_path = current_path + \"results/\"\n",
        "    weights_path = current_path + \"weights/\"\n",
        "\n",
        "    history_path = results_path + \"history/\"\n",
        "    images_path = results_path + \"images/\"\n",
        "    ckpts_path = results_path + \"ckpts/\"\n",
        "\n",
        "    best_path = ckpts_path + \"best/\"\n",
        "    latest_path = ckpts_path + \"latest/\"\n",
        "\n",
        "    paths = {\n",
        "        \"data\": data_path,\n",
        "        \"history\": history_path,\n",
        "        \"images\": images_path,\n",
        "        \"best\": best_path,\n",
        "        \"latest\": latest_path,\n",
        "        \"weights\": weights_path\n",
        "    }\n",
        "\n",
        "    return paths\n",
        "\n",
        "\n",
        "def train_model(dataset, paths, device):\n",
        "    \"\"\"The main function for executing network training. It loads the specified\n",
        "       dataset iterator, saliency model, and helper classes. Training is then\n",
        "       performed in a new session by iterating over all batches for a number of\n",
        "       epochs. After validation on an independent set, the model is saved and\n",
        "       the training history is updated.\n",
        "    Args:\n",
        "        dataset (str): Denotes the dataset to be used during training.\n",
        "        paths (dict, str): A dictionary with all path elements.\n",
        "        device (str): Represents either \"cpu\" or \"gpu\".\n",
        "    \"\"\"\n",
        "\n",
        "    iterator = get_dataset_iterator(\"train\", dataset, paths[\"data\"])\n",
        "\n",
        "    next_element, train_init_op, valid_init_op = iterator\n",
        "\n",
        "    input_images, ground_truths, input_targets  = next_element[:3]\n",
        "    ground_truths = tf.divide(ground_truths, 255)\n",
        "\n",
        "    input_plhd = tf.placeholder_with_default(input_images,\n",
        "                                             (None, None, None, 3),\n",
        "                                             name=\"input\")\n",
        "\n",
        "    input_target_img = tf.placeholder_with_default(input_targets,\n",
        "                                                    (None, None, None, 3),\n",
        "                                                    name=\"input_2\")\n",
        "\n",
        "    msinet = MSINET()\n",
        "\n",
        "    feature_map_stimuli = msinet.forward(input_plhd)\n",
        "\n",
        "    feature_map_target = msinet.forward(input_target_img)\n",
        "\n",
        "    predicted_maps = msinet.output_stream(feature_map_stimuli, feature_map_target)\n",
        "\n",
        "    # uncomment if you want to test with one stream network\n",
        "    # predicted_maps = msinet.one_stream(feature_map_stimuli)\n",
        "\n",
        "    optimizer, loss = msinet.train(ground_truths, predicted_maps,\n",
        "                                   config.PARAMS[\"learning_rate\"])\n",
        "\n",
        "    curr_module = sys.modules[__name__]\n",
        "\n",
        "    n_train_data = getattr(curr_module, dataset.upper()).n_train\n",
        "    n_valid_data = getattr(curr_module, dataset.upper()).n_valid\n",
        "\n",
        "    print(n_train_data)\n",
        "    print(n_valid_data)\n",
        "\n",
        "    n_train_batches = int(np.ceil(n_train_data / config.PARAMS[\"batch_size\"]))\n",
        "    n_valid_batches = int(np.ceil(n_valid_data / config.PARAMS[\"batch_size\"]))\n",
        "\n",
        "    history = History(n_train_batches,\n",
        "                            n_valid_batches,\n",
        "                            dataset,\n",
        "                            paths[\"history\"],\n",
        "                            device)\n",
        "\n",
        "    progbar = Progbar(n_train_data,\n",
        "                            n_train_batches,\n",
        "                            config.PARAMS[\"batch_size\"],\n",
        "                            config.PARAMS[\"n_epochs\"],\n",
        "                            history.prior_epochs)\n",
        "\n",
        "\n",
        "    with tf.Session() as sess:\n",
        "        \n",
        "        sess.run(tf.global_variables_initializer())\n",
        "        saver = msinet.restore(sess, dataset, paths, device)\n",
        "        writer = tf.summary.FileWriter('./tflogs', sess.graph)\n",
        "        print(\">> Start training on %s...\" % dataset.upper())\n",
        "\n",
        "        for epoch in range(config.PARAMS[\"n_epochs\"]):\n",
        "            sess.run(train_init_op)\n",
        "\n",
        "            for batch in range(n_train_batches):\n",
        "\n",
        "                _ , error = sess.run([optimizer, loss])\n",
        "\n",
        "                history.update_train_step(error)\n",
        "                progbar.update_train_step(batch)\n",
        "\n",
        "            sess.run(valid_init_op)\n",
        "\n",
        "            for batch in range(n_valid_batches):\n",
        "                error = sess.run(loss)\n",
        "\n",
        "                history.update_valid_step(error)\n",
        "                progbar.update_valid_step()\n",
        "\n",
        "            msinet.save(saver, sess, dataset, paths[\"latest\"], device)\n",
        "\n",
        "            history.save_history()\n",
        "\n",
        "            progbar.write_summary(history.get_mean_train_error(),\n",
        "                                history.get_mean_valid_error())\n",
        "\n",
        "            if history.valid_history[-1] == min(history.valid_history):\n",
        "                msinet.save(saver, sess, dataset, paths[\"best\"], device)\n",
        "                msinet.optimize(sess, dataset, paths[\"best\"], device)\n",
        "\n",
        "                print(\"\\tBest model!\", flush=True)\n",
        "\n",
        "\n",
        "def test_model(dataset, paths, device):\n",
        "    \"\"\"The main function for executing network testing. It loads the specified\n",
        "       dataset iterator and optimized saliency model. By default, when no model\n",
        "       checkpoint is found locally, the pretrained weights will be downloaded.\n",
        "       Testing only works for models trained on the same device as specified in\n",
        "       the config file.\n",
        "    Args:\n",
        "        dataset (str): Denotes the dataset that was used during training.\n",
        "        paths (dict, str): A dictionary with all path elements.\n",
        "        device (str): Represents either \"cpu\" or \"gpu\".\n",
        "    \"\"\"\n",
        "\n",
        "    iterator = get_dataset_iterator(\"test\", dataset, paths[\"data\"])\n",
        "\n",
        "    next_element, init_op = iterator\n",
        "\n",
        "    input_images, ground_truths, input_targets, original, file_path = next_element\n",
        "\n",
        "    original_shape = (320, 512)\n",
        "\n",
        "    graph_def = tf.GraphDef()\n",
        "\n",
        "    model_name = \"model_%s_%s.pb\" % (dataset, device)\n",
        "\n",
        "    if os.path.isfile(paths[\"best\"] + model_name):\n",
        "        with tf.gfile.Open(paths[\"best\"] + model_name, \"rb\") as file:\n",
        "            graph_def.ParseFromString(file.read())\n",
        "\n",
        "    predicted_maps = tf.import_graph_def(graph_def,\n",
        "                                         input_map={\"input\": input_images, \"input_2\": input_targets},\n",
        "                                         return_elements=[\"output:0\"])\n",
        "\n",
        "    predicted_maps = tf.squeeze(predicted_maps, axis=0)\n",
        "    input_images = tf.squeeze(input_images, axis=0)\n",
        "    jpeg , npy = postprocess_saliency_map(predicted_maps[0], original_shape)\n",
        "    \n",
        "    curr_module = sys.modules[__name__]\n",
        "    n_test_data = getattr(curr_module , 'TEST').n_test\n",
        "\n",
        "    print(\">> Start testing with %s %s model...\" % (dataset.upper(), device))\n",
        "\n",
        "    with tf.Session() as sess:\n",
        "        sess.run(init_op)\n",
        "\n",
        "        while True:\n",
        "            try:\n",
        "\n",
        "                output_file_jpeg, output_file_npy, path = sess.run(\n",
        "                    [jpeg, npy, file_path])\n",
        "\n",
        "            except tf.errors.OutOfRangeError:\n",
        "                break\n",
        "\n",
        "            path = path[0][0].decode(\"utf-8\")\n",
        "\n",
        "            filename = os.path.basename(path)\n",
        "            filename = os.path.splitext(filename)[0]\n",
        "            filename_jpeg = filename + \".jpg\"\n",
        "            filename_npy = filename + \".npy\"\n",
        "\n",
        "            os.makedirs(paths[\"images\"], exist_ok=True)\n",
        "\n",
        "            with open(paths[\"images\"] + filename_jpeg, \"wb\") as file:\n",
        "                file.write(output_file_jpeg)\n",
        "\n",
        "            with open(paths[\"images\"] + filename_npy, \"wb\") as file:\n",
        "                np.save(file, output_file_npy)\n",
        "               \n",
        "\n",
        "def jet_map(paths, threshold=30, alpha=0.5):\n",
        "\n",
        "    \"\"\"creates the jet map of predicted fixation density maps\n",
        "       for the test data.\n",
        "    Args:\n",
        "        paths (str): paths to the test search stimuli and results folder\n",
        "        threshold (int): threshold used to generate a mask of predicted density maps.\n",
        "        alpha (float): weight for overlaying the get color map and the original image\n",
        "    \"\"\"\n",
        "\n",
        "    test_data_path = paths['data'] + 'cocosearch/stimuli/test_targ_bbox/'\n",
        "    prediction_path = paths['images']\n",
        "    output_dir = paths['images'] + 'images_jet/'\n",
        "    gnd_dir = paths['data'] + 'cocosearch/saliencymap/test/'\n",
        "    gnd_output_dir = paths['images'] + 'groundtruth_jet/'\n",
        "\n",
        "    if not os.path.exists(gnd_output_dir):\n",
        "        os.makedirs(gnd_output_dir)\n",
        "\n",
        "    if not os.path.exists(output_dir):\n",
        "        os.makedirs(output_dir)\n",
        "\n",
        "    for subdir, dirs, files in os.walk(test_data_path):\n",
        "        for file in files:\n",
        "            if file.lower().endswith((\".png\", \".jpg\", \".jpeg\")):\n",
        "\n",
        "              ##predicted Saliency\n",
        "              img = cv2.imread(test_data_path + file)\n",
        "              heatmap = cv2.imread(prediction_path + file)\n",
        "              heatmap_color = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)\n",
        "              \n",
        "              # Create mask\n",
        "              mask = np.where(heatmap<=threshold, 1, 0)\n",
        "              mask = np.reshape(mask, (img.shape[0] , img.shape[1], 3))\n",
        "\n",
        "              # Marge images\n",
        "              marge = img*mask + heatmap_color*(1-mask)\n",
        "              marge = marge.astype(\"uint8\")\n",
        "\n",
        "              marge = cv2.addWeighted(img, 1-alpha, marge, alpha,0)\n",
        "              cv2.imwrite( output_dir + file ,marge)\n",
        "\n",
        "              ##Groundtruth saliency\n",
        "              heatmap_gnd = cv2.imread(gnd_dir + file)\n",
        "              heatmap_gnd_color = cv2.applyColorMap(heatmap_gnd, cv2.COLORMAP_JET)\n",
        "              \n",
        "              # Create mask\n",
        "              mask_gnd = np.where(heatmap_gnd<=threshold, 1, 0)\n",
        "              mask_gnd = np.reshape(mask_gnd, (img.shape[0] , img.shape[1], 3))\n",
        "\n",
        "              # Marge images\n",
        "              marge_gnd = img*mask_gnd + heatmap_gnd_color*(1-mask_gnd)\n",
        "              marge_gnd = marge_gnd.astype(\"uint8\")\n",
        "\n",
        "              marge_gnd = cv2.addWeighted(img, 1-alpha, marge_gnd, alpha,0)\n",
        "              cv2.imwrite(gnd_output_dir + file , marge_gnd)\n",
        "\n",
        "def main():\n",
        "    \n",
        "    \"\"\"The main function reads the command line arguments, invokes the\n",
        "       creation of appropriate path variables, and starts the training\n",
        "       or testing procedure for a model.\n",
        "    \"\"\"\n",
        "\n",
        "    phases_list = [\"train\", \"test\"]\n",
        "\n",
        "    dataset = 'cocosearch'\n",
        "\n",
        "    path = '/content/'\n",
        "    paths = define_paths(path)\n",
        "\n",
        "    tf.reset_default_graph()\n",
        "\n",
        "    train_model(dataset, paths, config.PARAMS[\"device\"])\n",
        "\n",
        "    test_model(dataset, paths, config.PARAMS[\"device\"])\n",
        "\n",
        "    threshold = 30 # threshold to generate jet maps\n",
        "\n",
        "    jet_map(paths, threshold, alpha=0.5)\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    main()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "cYbBmb-Urpm_"
      },
      "outputs": [],
      "source": [
        "\"\"\"This script caluclates the saliency metrics. \n",
        "The code is borrowed from \n",
        "https://github.com/tarunsharma1/saliency_metrics repository with\n",
        "slight modfications. \"\"\"\n",
        "\n",
        "import numpy as np\n",
        "import random\n",
        "import math\n",
        "\n",
        "def normalize_map(s_map):\n",
        "    # normalize the salience map (as done in MIT code)\n",
        "    norm_s_map = (s_map - np.min(s_map))/((np.max(s_map)-np.min(s_map))*1.0)\n",
        "    return norm_s_map\n",
        "\n",
        "def discretize_gt(gt):\n",
        "\timport warnings\n",
        "\twarnings.warn('can improve the way GT is discretized')\n",
        "\treturn gt/255\n",
        "\n",
        "def auc_judd(s_map,gt):\n",
        "\t# ground truth is discrete, s_map is continous and normalized\n",
        "\tgt = discretize_gt(gt)\n",
        "\t# thresholds are calculated from the salience map, only at places where fixations are present\n",
        "\tthresholds = []\n",
        "\tfor i in range(0,gt.shape[0]):\n",
        "\t\tfor k in range(0,gt.shape[1]):\n",
        "\t\t\tif gt[i][k]>0:\n",
        "\t\t\t\tthresholds.append(s_map[i][k])\n",
        "\n",
        "\tnum_fixations = np.sum(gt)\n",
        "\t# num fixations is no. of salience map values at gt >0\n",
        "\n",
        "\tthresholds = sorted(set(thresholds))\t\n",
        "\t#fp_list = []\n",
        "\t#tp_list = []\n",
        "\tarea = []\n",
        "\tarea.append((0.0,0.0))\n",
        "\tfor thresh in thresholds:\n",
        "\t\t# in the salience map, keep only those pixels with values above threshold\n",
        "\t\ttemp = np.zeros(s_map.shape)\n",
        "\t\ttemp[s_map>=thresh] = 1.0\n",
        "\t\tassert np.max(gt)==1.0, 'something is wrong with ground truth..not discretized properly max value > 1.0'\n",
        "\t\tassert np.max(s_map)==1.0, 'something is wrong with salience map..not normalized properly max value > 1.0'\n",
        "\t\tnum_overlap = np.where(np.add(temp,gt)==2)[0].shape[0]\n",
        "\t\ttp = num_overlap/(num_fixations*1.0)\n",
        "\t\t\n",
        "\t\t# total number of pixels > threshold - number of pixels that overlap with gt / total number of non fixated pixels\n",
        "\t\t# this becomes nan when gt is full of fixations..this won't happen\n",
        "\t\tfp = (np.sum(temp) - num_overlap)/((np.shape(gt)[0] * np.shape(gt)[1]) - num_fixations)\n",
        "\t\t\n",
        "\t\tarea.append((round(tp,4),round(fp,4)))\n",
        "\t\t#tp_list.append(tp)\n",
        "\t\t#fp_list.append(fp)\n",
        "\n",
        "\t#tp_list.reverse()\n",
        "\t#fp_list.reverse()\n",
        "\tarea.append((1.0,1.0))\n",
        "\t#tp_list.append(1.0)\n",
        "\t#fp_list.append(1.0)\n",
        "\t#print tp_list\n",
        "\tarea.sort(key = lambda x:x[0])\n",
        "\ttp_list =  [x[0] for x in area]\n",
        "\tfp_list =  [x[1] for x in area]\n",
        "\treturn np.trapz(np.array(tp_list), np.array(fp_list))\n",
        "\n",
        "def auc_borji(s_map, gt, splits=100, stepsize=0.1):\n",
        "\tgt = discretize_gt(gt)\n",
        "\tnum_fixations = np.sum(gt)\n",
        "\n",
        "\tnum_pixels = s_map.shape[0]*s_map.shape[1]\n",
        "\trandom_numbers = []\n",
        "\tfor i in range(0,splits):\n",
        "\t\ttemp_list = []\n",
        "\t\tfor k in range(0, int(num_fixations)):\n",
        "\t\t\ttemp_list.append(np.random.randint(num_pixels))\n",
        "\t\trandom_numbers.append(temp_list)\n",
        "\n",
        "\taucs = []\n",
        "\t# for each split, calculate auc\n",
        "\tfor i in random_numbers:\n",
        "\t\tr_sal_map = []\n",
        "\t\tfor k in i:\n",
        "\t\t\tr_sal_map.append(s_map[k%s_map.shape[0]-1, int(k/s_map.shape[0])])\n",
        "\t\t# in these values, we need to find thresholds and calculate auc\n",
        "\t\tthresholds = [0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9]\n",
        "\n",
        "\t\tr_sal_map = np.array(r_sal_map)\n",
        "\n",
        "\t\t# once threshs are got\n",
        "\t\tthresholds = sorted(set(thresholds))\n",
        "\t\tarea = []\n",
        "\t\tarea.append((0.0,0.0))\n",
        "\t\tfor thresh in thresholds:\n",
        "\t\t\t# in the salience map, keep only those pixels with values above threshold\n",
        "\t\t\ttemp = np.zeros(s_map.shape)\n",
        "\t\t\ttemp[s_map>=thresh] = 1.0\n",
        "\t\t\tnum_overlap = np.where(np.add(temp,gt)==2)[0].shape[0]\n",
        "\t\t\ttp = num_overlap/(num_fixations*1.0)\n",
        "\t\t\t\n",
        "\t\t\t#fp = (np.sum(temp) - num_overlap)/((np.shape(gt)[0] * np.shape(gt)[1]) - num_fixations)\n",
        "\t\t\t# number of values in r_sal_map, above the threshold, divided by num of random locations = num of fixations\n",
        "\t\t\tfp = len(np.where(r_sal_map>thresh)[0])/(num_fixations*1.0)\n",
        "\n",
        "\t\t\tarea.append((round(tp,4),round(fp,4)))\n",
        "\t\t\n",
        "\t\tarea.append((1.0,1.0))\n",
        "\t\tarea.sort(key = lambda x:x[0])\n",
        "\t\ttp_list =  [x[0] for x in area]\n",
        "\t\tfp_list =  [x[1] for x in area]\n",
        "\n",
        "\t\taucs.append(np.trapz(np.array(tp_list),np.array(fp_list)))\n",
        "\t\n",
        "\treturn np.mean(aucs)\n",
        "\n",
        "def auc_shuff(s_map,gt,other_map,splits=100,stepsize=0.1):\n",
        "\tgt = discretize_gt(gt)\n",
        "\t#print(np.max(other_map))\n",
        "\t#other_map = discretize_gt(other_map)\n",
        "\t#print(np.max(other_map))\n",
        "\tnum_fixations = np.sum(gt)\n",
        "\t\n",
        "\tx,y = np.where(other_map==1.0)\n",
        "\n",
        "\tother_map_fixs = []\n",
        "\tfor j in zip(x,y):\n",
        "\t\tother_map_fixs.append(j[0]*other_map.shape[0] + j[1])\n",
        "\tind = len(other_map_fixs)\n",
        "\tassert ind==np.sum(other_map), 'something is wrong in auc shuffle'\n",
        "\n",
        "\n",
        "\tnum_fixations_other = min(ind,num_fixations)\n",
        "\n",
        "\tnum_pixels = s_map.shape[0]*s_map.shape[1]\n",
        "\trandom_numbers = []\n",
        "\tfor i in range(0,splits):\n",
        "\t\ttemp_list = []\n",
        "\t\tt1 = np.random.permutation(ind)\n",
        "\t\tfor k in t1:\n",
        "\t\t\ttemp_list.append(other_map_fixs[k])\n",
        "\t\trandom_numbers.append(temp_list)\t\n",
        "\n",
        "\taucs = []\n",
        "\t# for each split, calculate auc\n",
        "\tfor i in random_numbers:\n",
        "\t\tr_sal_map = []\n",
        "\t\tfor k in i:\n",
        "\n",
        "\t\t\tr_sal_map.append(s_map[k%s_map.shape[0]-1, int(k/s_map.shape[0])])\n",
        "\t\t# in these values, we need to find thresholds and calculate auc\n",
        "\t\tthresholds = [0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9]\n",
        "\n",
        "\t\tr_sal_map = np.array(r_sal_map)\n",
        "\n",
        "\t\t# once threshs are got\n",
        "\t\tthresholds = sorted(set(thresholds))\n",
        "\t\tarea = []\n",
        "\t\tarea.append((0.0,0.0))\n",
        "\t\tfor thresh in thresholds:\n",
        "\t\t\t# in the salience map, keep only those pixels with values above threshold\n",
        "\t\t\ttemp = np.zeros(s_map.shape)\n",
        "\t\t\ttemp[s_map>=thresh] = 1.0\n",
        "\t\t\tnum_overlap = np.where(np.add(temp,gt)==2)[0].shape[0]\n",
        "\t\t\ttp = num_overlap/(num_fixations*1.0)\n",
        "\t\t\t\n",
        "\t\t\t#fp = (np.sum(temp) - num_overlap)/((np.shape(gt)[0] * np.shape(gt)[1]) - num_fixations)\n",
        "\t\t\t# number of values in r_sal_map, above the threshold, divided by num of random locations = num of fixations\n",
        "\t\t\tfp = len(np.where(r_sal_map>thresh)[0])/(num_fixations*1.0)\n",
        "\n",
        "\t\t\tarea.append((round(tp,4),round(fp,4)))\n",
        "\t\t\n",
        "\t\tarea.append((1.0,1.0))\n",
        "\t\tarea.sort(key = lambda x:x[0])\n",
        "\t\ttp_list =  [x[0] for x in area]\n",
        "\t\tfp_list =  [x[1] for x in area]\n",
        "\n",
        "\t\taucs.append(np.trapz(np.array(tp_list),np.array(fp_list)))\n",
        "\t\n",
        "\treturn np.mean(aucs)\n",
        "\n",
        "\n",
        "def nss(s_map,gt):\n",
        "\n",
        "\tgt = discretize_gt(gt)\n",
        "\ts_map_norm = (s_map - np.mean(s_map))/np.std(s_map)\n",
        "\n",
        "\tx,y = np.where(gt==1)\n",
        "\ttemp = []\n",
        "\tfor i in zip(x,y):\n",
        "\t\ttemp.append(s_map_norm[i[0],i[1]])\n",
        "\treturn np.mean(temp)\n",
        "\n",
        "\n",
        "def infogain(s_map,gt,baseline_map):\n",
        "\t\n",
        "\tgt = discretize_gt(gt)\n",
        "\t# assuming s_map and baseline_map are normalized\n",
        "\teps = 2.2204e-16\n",
        "\n",
        "\ts_map = s_map/(np.sum(s_map)*1.0)\n",
        "\tbaseline_map = baseline_map/(np.sum(baseline_map)*1.0)\n",
        "\n",
        "\t# for all places where gt=1, calculate info gain\n",
        "\ttemp = []\n",
        "\tx,y = np.where(gt==1.0)\n",
        "\t#print(x,y)\n",
        "\tfor i in zip(x,y):\n",
        "\t\ttemp.append(np.log2(eps + s_map[i[0],i[1]]) - np.log2(eps + baseline_map[i[0],i[1]]))\n",
        "\n",
        "\treturn np.mean(temp)\n",
        "\n",
        "\n",
        "def similarity(s_map,gt):\n",
        "\t# here gt is not discretized nor normalized\n",
        "\ts_map = s_map/(np.sum(s_map)*1.0)\n",
        "\tgt = gt/(np.sum(gt)*1.0)\n",
        "\t#print(s_map.shape)\n",
        "\t#print(gt.shape)\n",
        "\tx,y = np.where(gt>0)\n",
        "\tsim = 0.0\n",
        "\tfor i in zip(x,y):\n",
        "\t\tsim = sim + min(gt[i[0],i[1]],s_map[i[0],i[1]])\n",
        "\treturn sim\n",
        "\n",
        "\n",
        "def cc(s_map,gt):\n",
        "\ts_map_norm = (s_map - np.mean(s_map))/np.std(s_map)\n",
        "\tgt_norm = (gt - np.mean(gt))/np.std(gt)\n",
        "\ta = s_map_norm\n",
        "\tb= gt_norm\n",
        "\tr = (a*b).sum() / math.sqrt((a*a).sum() * (b*b).sum());\n",
        "\treturn r\n",
        "\n",
        "\n",
        "def kld(s_map,gt):\n",
        "\ts_map = s_map/(np.sum(s_map)*1.0)\n",
        "\tgt = gt/(np.sum(gt)*1.0)\n",
        "\teps = 2.2204e-16\n",
        "\treturn np.sum(gt * np.log(eps + gt/(s_map + eps)))\n",
        "\n",
        "\n",
        "def calculate_metrics(y_pred , y_true, y_true_binary, bl_fixmap, bl_salmap):\n",
        "\n",
        "\ty_pred = np.squeeze(y_pred)\n",
        "\ty_true = np.squeeze(y_true)\n",
        "\ty_true_binary = np.squeeze(y_true_binary)\n",
        "\tbl_fixmap = np.squeeze(bl_fixmap)\n",
        "\tbl_salmap = np.squeeze(bl_salmap)\n",
        "\tcc_error = cc(y_pred, y_true)\n",
        "\ty_pred_n = normalize_map(y_pred)\n",
        "\ty_true_n = normalize_map(y_true)\n",
        "\tbl_salmap_n = normalize_map(bl_salmap)\n",
        "\tkld_error = kld(y_pred, y_true)#\n",
        "\tinfog_error = infogain(y_pred_n, y_true_binary, bl_salmap_n)#?\n",
        "\tsim_error = similarity(y_pred_n, y_true_n)\n",
        "\tnss_error = nss(y_pred, y_true_binary)\n",
        "\tauc_error = auc_judd(y_pred_n, y_true_binary)\n",
        "\tauc_borji_error = auc_borji(y_pred_n, y_true_binary)\n",
        "\tsauc_error = auc_shuff(y_pred_n, y_true_binary, bl_fixmap)\n",
        "\n",
        "\treturn kld_error , cc_error , sim_error, nss_error, auc_error, infog_error, sauc_error, auc_borji_error"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "-huYfdLrrqUp",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f811ab30-1d72-46a1-e912-91fca93cacbe"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/ipykernel_launcher.py:17: UserWarning: can improve the way GT is discretized\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "locally computed:\n",
            "AUC: 0.934979192314815\n",
            "AUC Borji: 0.8620678619958333\n",
            "SAUC: 0.8502377083796303\n",
            "NSS: 4.203472264455976\n",
            "image_KLD: 1.0647463152805965\n",
            "CC: 0.6563009473182845\n",
            "SIM: 0.4650988454010617\n",
            "INFO GAIN: 2.311629964751863\n"
          ]
        }
      ],
      "source": [
        "import pysaliency\n",
        "from pysaliency.baseline_utils import BaselineModel, GoldModel\n",
        "import argparse\n",
        "import os\n",
        "import json\n",
        "import numpy as np\n",
        "import cv2\n",
        "import tensorflow as tf\n",
        "import random\n",
        "import glob\n",
        "import shutil\n",
        "from pandas import DataFrame\n",
        "#tf.compat.v1.enable_eager_execution()\n",
        "\n",
        "def compute_saliency_metrics(data_path, use_pysaliency, csv_path):\n",
        "    \"\"\"computes the saliency metrics for test set.\n",
        "    Args:\n",
        "        data_path(str): The path to where the dataset is stored.\n",
        "        use_pysaliency(bool): Whether to use pysaliency library to measure\n",
        "                              saliency metrics.\n",
        "        csv_path (str): The path to where a csv file containing\n",
        "                        the computed saliency metrics should be stored.\n",
        "    \"\"\"\n",
        "    test_list = []\n",
        "    category = ['bottle', 'bowl' , 'cup', 'car', 'chair', \n",
        "                   'clock', 'fork', 'keyboard', 'knife', \n",
        "                   'laptop', 'microwave', 'mouse', 'oven',\n",
        "                   'potted plant', 'sink', 'stop sign', \n",
        "                   'toilet', 'tv']\n",
        "    test_stimuli_path = data_path + \"cocosearch/stimuli/test\"\n",
        "\n",
        "    for subdir, dirs, files in os.walk(test_stimuli_path):\n",
        "        for file in files:\n",
        "            if file.lower().endswith((\".png\", \".jpg\", \".jpeg\")):\n",
        "                test_list.append(os.path.join(subdir, file))\n",
        "\n",
        "    dir_saliency_test = data_path + \"cocosearch/saliencymap/test\"\n",
        "    dir_saliency_test_unblur = data_path + \"cocosearch/saliencymap/test_unblur\"\n",
        "    dir_stimuli_test = pysaliency.FileStimuli(test_list)\n",
        "    dir_saliency_test_img = os.path.join(\"cocosearch/saliencymap\" , 'pysaliency_test_sal_img')\n",
        "    \n",
        "    if use_pysaliency:\n",
        "\n",
        "        if os.path.exists(dir_saliency_test_img):\n",
        "            shutil.rmtree(dir_saliency_test_img)\n",
        "            os.makedirs(dir_saliency_test_img)\n",
        "        else:\n",
        "            os.makedirs(dir_saliency_test_img)\n",
        "        \n",
        "        dir_results_test_img = os.path.join(data_path + 'results/images' , 'pysaliency_result_img')\n",
        "\n",
        "        if os.path.exists(dir_results_test_img):\n",
        "            shutil.rmtree(dir_results_test_img)\n",
        "            os.makedirs(dir_results_test_img)\n",
        "        else:\n",
        "            os.makedirs(dir_results_test_img)\n",
        "\n",
        "        for item in glob.glob(os.path.join(dir_saliency_test , '*.jpg')):\n",
        "            os.symlink(item , os.path.join(dir_saliency_test_img , os.path.basename(item)))\n",
        "\n",
        "        for item in glob.glob(os.path.join(data_path + 'results/images' , '*.jpg')):\n",
        "            os.symlink(item, os.path.join(dir_results_test_img , os.path.basename(item)))\n",
        "\n",
        "\n",
        "    with open(data_path + 'coco_search18_fixations_TP_train_split1.json') as json_file:\n",
        "        human_scanpaths_train = json.load(json_file)\n",
        "\n",
        "    xs = []\n",
        "    ys = []\n",
        "    ts = []\n",
        "    ns = []\n",
        "    train_subjects = []\n",
        "    min_fix_x = 100000\n",
        "    max_fix_x = -100000\n",
        "    min_fix_y = 100000\n",
        "    max_fix_y = -100000\n",
        "\n",
        "    for traj in human_scanpaths_train:\n",
        "        for i in range(len(traj['X'])):\n",
        "            if traj['X'][i] < 0 or traj['Y'][i] < 0 or traj['X'][i] > 1680 or traj['Y'][i] > 1050:\n",
        "                continue\n",
        "\n",
        "            if traj['X'][i] < min_fix_x:\n",
        "                min_fix_x = traj['X'][i]\n",
        "\n",
        "            if traj['X'][i] > max_fix_x:\n",
        "                max_fix_x = traj['X'][i]\n",
        "\n",
        "            if traj['Y'][i] < min_fix_y:\n",
        "                min_fix_y = traj['Y'][i]\n",
        "\n",
        "            if traj['Y'][i] > max_fix_y:\n",
        "                max_fix_y = traj['Y'][i]\n",
        "\n",
        "    #for cat in category:\n",
        "    #print(str(cat))\n",
        "    m_kld_error, m_cc_error, m_sim_error, m_nss_error, m_auc_error, m_infog_error, m_sauc_error, m_auc_b_error = [],[],[],[],[],[],[],[]\n",
        "    stimuli_names = []\n",
        "    #count = len(glob.glob(os.path.join(test_stimuli_path , str(cat)+ '*.jpg')))#\n",
        "    count =  len(os.listdir(test_stimuli_path))\n",
        "    for subdir, dirs, files in os.walk(test_stimuli_path):\n",
        "\n",
        "        \n",
        "        for n, stimulus in enumerate(files):\n",
        "\n",
        "            #if stimulus.startswith(str(cat)):   \n",
        "            base_line_salmap = np.zeros((320 , 512))\n",
        "            for j in files:\n",
        "                if j!= stimulus: #and j.startswith(str(cat)):\n",
        "                    base_line_salmap = base_line_salmap + np.load(os.path.join(dir_saliency_test , os.path.splitext(j)[0]+'.npy'), allow_pickle=True)\n",
        "\n",
        "            base_line_salmap /= np.max(base_line_salmap)\n",
        "            M = 1\n",
        "            random_ind = random.sample(range(0, count), M)\n",
        "            base_line_fixmap = np.zeros((320 , 512))\n",
        "            for i in random_ind:\n",
        "\n",
        "                while(files[i]==stimulus):\n",
        "                    i = random.randint(0, count-1)\n",
        "\n",
        "                if files[i]!=stimulus: # and files[i].startswith(str(cat)):\n",
        "                    base_line_fixmap = base_line_fixmap + (np.load(os.path.join(dir_saliency_test_unblur , os.path.splitext(files[i])[0]+'.npy'), allow_pickle=True))/255\n",
        "                    \n",
        "            base_line_fixmap[np.where(base_line_fixmap>1.0)] = 1.0\n",
        "\n",
        "            ##gnd_map = tf.image.decode_jpeg(tf.read_file(os.path.join(dir_saliency_test , stimulus)), channels=1).numpy()\n",
        "            ##gnd_bin_map = tf.image.decode_jpeg(tf.read_file(os.path.join(dir_saliency_test_unblur , stimulus)), channels=1).numpy()\n",
        "            ##pred_map = tf.image.decode_jpeg(tf.read_file(os.path.join(data_path + 'results/images' , os.path.splitext(stimulus)[0]+'.png')), channels=1).numpy()\n",
        "\n",
        "            gnd_map  = np.load(os.path.join(dir_saliency_test , os.path.splitext(stimulus)[0]+'.npy'), allow_pickle=True)      \n",
        "            gnd_bin_map  = np.load(os.path.join(dir_saliency_test_unblur , os.path.splitext(stimulus)[0]+'.npy'), allow_pickle=True)\n",
        "            pred_map = np.load(os.path.join(data_path + 'results/images' , os.path.splitext(stimulus)[0]+'.npy'), allow_pickle=True)\n",
        "\n",
        "            kl_test_error, cc_test_error, sim_test_error, nss_test_error, auc_test_error, infog_test_error, sauc_test_error, auc_b_test_error = calculate_metrics(pred_map, gnd_map, gnd_bin_map, base_line_fixmap, base_line_salmap)\n",
        "\n",
        "            stimuli_names.append(stimulus)\n",
        "            m_kld_error.append(kl_test_error)\n",
        "            m_cc_error.append(cc_test_error)\n",
        "            m_sim_error.append(sim_test_error)\n",
        "            m_nss_error.append(nss_test_error)\n",
        "            m_auc_error.append(auc_test_error)\n",
        "            m_infog_error.append(infog_test_error)\n",
        "            m_sauc_error.append(sauc_test_error)\n",
        "            m_auc_b_error.append(auc_b_test_error)\n",
        "            \n",
        "            if use_pysaliency:\n",
        "\n",
        "                stimulus_size = dir_stimuli_test.sizes[n]\n",
        "                height, width = stimulus_size\n",
        "\n",
        "                for traj in human_scanpaths_train:\n",
        "\n",
        "                    \n",
        "                    if ((traj['task'] + '_' + traj['name'])==stimulus and traj['correct'] == 1):\n",
        "                        subject_name = traj['subject']\n",
        "\n",
        "                        xs_ = []\n",
        "                        ys_ = []\n",
        "                        durations_ = []\n",
        "\n",
        "                        for ind in range( len(traj['X'])):\n",
        "                            \n",
        "                            if 1680>=traj['X'][ind]>=0 and 1050>=traj['Y'][ind]>=0:\n",
        "                                if not (traj['X'][ind] in xs_ and traj['Y'][ind] in ys_):\n",
        "                                    xs_.append(((traj['X'][ind] - min_fix_x) / max_fix_x)*(512))\n",
        "                                    ys_.append(((traj['Y'][ind] - min_fix_y) / max_fix_y)*(320))\n",
        "                                    durations_.append(traj['T'][ind])\n",
        "\n",
        "\n",
        "                        time_l_ = durations_ [:-1]\n",
        "                        time_l_.insert(0,0)\n",
        "                        time_l_array_ = np.array(time_l_)\n",
        "                        ts_ = [np.sum(time_l_array_[0:ind]) for ind in range(1, len(time_l_array_)+1)]\n",
        "\n",
        "                        xs.append(xs_)\n",
        "                        ys.append(ys_)\n",
        "                        ts.append(ts_)\n",
        "                        ns.append(n)\n",
        "                        train_subjects.append(subject_name )\n",
        "\n",
        "    if use_pysaliency:\n",
        "\n",
        "        fixations = pysaliency.FixationTrains.from_fixation_trains(xs, ys, ts, ns, train_subjects, attributes=False, scanpath_attributes=None)\n",
        "\n",
        "        my_model = pysaliency.SaliencyMapModelFromDirectory(dir_stimuli_test, dir_results_test_img)\n",
        "\n",
        "        ground_truth = pysaliency.SaliencyMapModelFromDirectory(dir_stimuli_test,  dir_saliency_test_img)\n",
        "\n",
        "        auc = my_model.AUC(dir_stimuli_test , fixations)\n",
        "        sauc = my_model.sAUC(dir_stimuli_test , fixations)\n",
        "        nss = my_model.NSS(dir_stimuli_test , fixations)\n",
        "        cc = my_model.CC(dir_stimuli_test , ground_truth)\n",
        "        sim = my_model.SIM(dir_stimuli_test , ground_truth)\n",
        "\n",
        "        img_kld = my_model.image_based_kl_divergence(dir_stimuli_test , ground_truth)\n",
        "\n",
        "        print('metrics computed by pysaliency:')\n",
        "        print('AUC:', auc )\n",
        "        print('sAUC:', sauc )\n",
        "        print('NSS:', nss )\n",
        "        print('image_KLD:', img_kld) \n",
        "        print('CC:', cc )\n",
        "        print('SIM:', sim )\n",
        "\n",
        "    kl_mean_test_error = sum(m_kld_error) / len(m_kld_error)#m_kld_error / count\n",
        "    cc_mean_test_error = sum(m_cc_error) / len(m_cc_error)#m_cc_error / count\n",
        "    sim_mean_test_error = sum(m_sim_error) / len(m_sim_error)#m_sim_error / count \n",
        "    nss_mean_test_error = sum(m_nss_error) / len(m_nss_error)#m_nss_error / count\n",
        "    auc_mean_test_error = sum(m_auc_error) / len(m_auc_error)#m_auc_error / count\n",
        "    infog_mean_test_error = sum(m_infog_error) / len(m_infog_error)#m_infog_error / count\n",
        "    sauc_mean_test_error = sum(m_sauc_error) / len(m_sauc_error)#m_sauc_error / count\n",
        "    auc_b_mean_test_error = sum(m_auc_b_error) / len(m_auc_b_error)#m_auc_b_error / count\n",
        "\n",
        "    print('locally computed:')\n",
        "    print('AUC:',  auc_mean_test_error)\n",
        "    print('AUC Borji:',  auc_b_mean_test_error)\n",
        "    print('SAUC:',  sauc_mean_test_error)\n",
        "    print('NSS:',  nss_mean_test_error )\n",
        "    print('image_KLD:', kl_mean_test_error) # this should be same as our loss function\n",
        "    print('CC:', cc_mean_test_error )\n",
        "    print('SIM:',sim_mean_test_error )\n",
        "    print('INFO GAIN:', infog_mean_test_error )\n",
        "\n",
        "    sample = {}\n",
        "\n",
        "    sample['kl'] = sorted(zip(stimuli_names, m_kld_error), key=lambda x: x[1])\n",
        "    sample['cc'] = sorted(zip(stimuli_names, m_cc_error), key=lambda x: x[1])\n",
        "    sample['sim'] = sorted(zip(stimuli_names, m_sim_error), key=lambda x: x[1])\n",
        "    sample['nss'] = sorted(zip(stimuli_names, m_nss_error), key=lambda x: x[1])\n",
        "    sample['auc'] = sorted(zip(stimuli_names, m_auc_error), key=lambda x: x[1])\n",
        "    sample['infog'] = sorted(zip(stimuli_names, m_infog_error), key=lambda x: x[1])\n",
        "    sample['sauc'] = sorted(zip(stimuli_names, m_sauc_error), key=lambda x: x[1])\n",
        "    sample['auc_borji'] = sorted(zip(stimuli_names, m_auc_b_error), key=lambda x: x[1])\n",
        "\n",
        "    with open(csv_path + 'ordering_test_images_based_on_score.json', 'w') as fp:\n",
        "        json.dump(str(sample), fp)\n",
        "\n",
        "    df = DataFrame({'AUC': [auc_mean_test_error], 'AUC Borji':[auc_b_mean_test_error] , 'SAUC':[sauc_mean_test_error] , 'NSS':[nss_mean_test_error] , \\\n",
        "                        'image_KLD':[kl_mean_test_error], 'CC':[cc_mean_test_error], 'SIM':[sim_mean_test_error], 'INFO GAIN':[infog_mean_test_error]})\n",
        "\n",
        "    output_path= csv_path + 'sal_metrics.xlsx'\n",
        "    #output_path= csv_path + str(cat) + '_sal_metrics.xlsx'\n",
        "    df.to_csv(output_path, mode='a', header=not os.path.exists(output_path), index=False)\n",
        "\n",
        "def str2bool(v):\n",
        "    if isinstance(v, bool):\n",
        "        return v\n",
        "    if v.lower() in ('yes', 'true', 't', 'y', '1'):\n",
        "        return True\n",
        "    elif v.lower() in ('no', 'false', 'f', 'n', '0'):\n",
        "        return False\n",
        "    else:\n",
        "        raise argparse.ArgumentTypeError('Boolean value expected.')\n",
        "\n",
        "def main():\n",
        "    \n",
        "    \"\"\"The main function reads the command line arguments, invokes the\n",
        "       creation of appropriate path variables, and starts the training\n",
        "       or testing procedure for a model.\n",
        "    \"\"\"\n",
        "\n",
        "    path = './'\n",
        "    csv_path = './results/'\n",
        "    use_pysaliency = False\n",
        "\n",
        "    compute_saliency_metrics(path, use_pysaliency, csv_path)\n",
        "        \n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    main()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5rQyQ4oUc-W2"
      },
      "outputs": [],
      "source": [
        "from google.colab import files\n",
        "\n",
        "!zip -r './results_th30.zip' './results/images/images_jet'\n",
        "files.download('./results_th30.zip') "
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "spatial_saliency_colab.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.5"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}