generate_anchor_image_dict.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import os.path\n",
    "import pickle\n",
    "import shutil\n",
    "import pandas as pd\n",
    "import random\n",
    "import math\n",
    "import matplotlib.pyplot as plt\n",
    "from os.path import join, exists\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import random\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('results/ood_prob_dict.pickle', 'rb') as f:\n",
    "    ood_prob_dict = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('results/HR_prob_dict.pickle', 'rb') as f:\n",
    "    HR_prob_dict = pickle.load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate anchor image dict for deriving image pair (ref + target) dataset\n",
    "\n",
    "A \"anchor_images_dict\" maps a target image in a sequence to all its reference images in this sequence. This is needed to be run before deploying the LR model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "dir_list = ['demo_sequences']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "root_data_dir = 'data/sequences'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "HR_threshold = 0.85"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def anchor_model_1_get_all_images(prob_dict, fname_list, threshold=0.5):\n",
    "    \"\"\"\n",
    "    prob_dict: {fname: probability}, e.g. {'2006_0.png': 0.001, '2007_1.png': 0.132, ...}\n",
    "    fname_list: a candidate list of filenames\n",
    "    \"\"\"\n",
    "    anchor_images = []\n",
    "    for fname in fname_list:\n",
    "        prob = prob_dict[fname]\n",
    "        if prob >= threshold:\n",
    "            anchor_images.append(fname)\n",
    "    if len(anchor_images) == 0:\n",
    "        anchor_images.append(fname_list[-1])\n",
    "    return anchor_images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "anchor_images_dict = dict()\n",
    "for subdir in dir_list:\n",
    "    data_dir = join(root_data_dir, subdir)\n",
    "    for folder in os.listdir(data_dir):\n",
    "        idx = folder.split('_')[0]\n",
    "        folder_dir = join(data_dir, folder)\n",
    "        fname_list = []\n",
    "        for f in os.listdir(folder_dir):\n",
    "            if f[-4:] == '.png':\n",
    "                fname_list.append(f)\n",
    "        HR_prob_dict_sub = HR_prob_dict[idx]\n",
    "        anchor_images = anchor_model_1_get_all_images(HR_prob_dict_sub, fname_list, threshold=HR_threshold)\n",
    "        anchor_images_dict[idx] = anchor_images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('results/anchor_images_dict.pickle', 'wb') as f:\n",
    "    pickle.dump(anchor_images_dict, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Environment (conda_tensorflow_p36)",
   "language": "python",
   "name": "conda_tensorflow_p36"
  },
  "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}