python_implementation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import sys "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this cell is specifically for PoolNet (ground truth images)\n",
    "\n",
    "gt_list = []\n",
    "gt_paths = []\n",
    "path = './data/SOD/Imgs'  # ground image path directory\n",
    "test = os.listdir(path)\n",
    "for i in test: \n",
    "    if(i[-3:] == 'png'):\n",
    "        gt_list.append(i)\n",
    "gt_list = sorted(gt_list)\n",
    "for i in range(len(gt_list)):\n",
    "    gt_paths.append(os.path.join(path, gt_list[i]))\n",
    "\n",
    "gt_paths = sorted(gt_paths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this cell is specifically for PoolNet (predicted images)\n",
    "\n",
    "pred_paths = []\n",
    "path = './results/run-1-sal-s' # predicted image directory\n",
    "test_pred = sorted(os.listdir(path))\n",
    "for i in range(len(test_pred)):\n",
    "    pred_paths.append(os.path.join(path, test_pred[i]))\n",
    "\n",
    "pred_paths = sorted(pred_paths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prec_rec(y_true, y_pred, beta2):\n",
    "    \n",
    "    eps = sys.float_info.epsilon\n",
    "    tp = torch.sum(y_true * y_pred)\n",
    "    all_p_pred = torch.sum(y_pred)\n",
    "    all_p_true = torch.sum(y_true)\n",
    "    \n",
    "    prec = (tp + eps) / (all_p_pred + eps)\n",
    "    rec = (tp + eps) / (all_p_true + eps)\n",
    "    # print(prec)\n",
    "    # print(rec)\n",
    "    \n",
    "    return prec, rec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "overall_mae = 0\n",
    "for j in range(len(gt_paths)):\n",
    "\n",
    "        gt = np.array(Image.open(gt_paths[j]).convert('LA')) / 255\n",
    "        pred = np.array(Image.open(pred_paths[j]).convert('LA')) / 255 \n",
    "        mae = np.sum(np.abs(pred - gt)) / (pred.shape[:2][0] * pred.shape[:2][1])\n",
    "        \n",
    "        gt_arr = torch.from_numpy(np.array(gt)).float()\n",
    "        pred_arr = torch.from_numpy(np.array(pred)).float()\n",
    "        threshold = 216\n",
    "        y_pred = torch.ge(pred_arr, threshold).float()\n",
    "        y_true = torch.ge(gt_arr, 128).float()\n",
    "        y_true1 = torch.reshape(y_true, (1,-1))\n",
    "        y_pred1 = torch.reshape(y_pred, (1,-1))\n",
    "        \n",
    "        prec, rec = prec_rec(y_true1, y_pred1,0.3)\n",
    "\n",
    "        total_prec = total_prec + prec\n",
    "        total_rec = total_rec + rec\n",
    "\n",
    "        overall_mae = overall_mae + mae\n",
    "\n",
    "        print(j)\n",
    "\n",
    "beta2 = 0.3 \n",
    "overall_fb = (1+beta2) * (total_prec * total_rec) / ((beta2 * total_prec + total_rec) * len(gt_paths))\n",
    "print('overall_mae', overall_mae )\n",
    "print('overall_fb', overall_fb )"
   ]
  }
 ],
 "metadata": {
  "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}