{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b870371e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "sys.path.append(\"../../src\")\n",
    "sys.path.append(\"../../data\")\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.image as mpimg\n",
    "import cv2\n",
    "from IPython import display\n",
    "import pylab as pl\n",
    "\n",
    "from CorInfoMaxBSS import *\n",
    "from general_utils import *\n",
    "from visualization_utils import * \n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "notebook_name = 'Image Separation'\n",
    "\n",
    "# np.random.seed(250)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3f896d7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "results_path = \"Image_Separation_Results\"\n",
    "if not os.path.exists(results_path):\n",
    "    os.mkdir(results_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a8a7d2e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_paths = '../../data/TestImages'\n",
    "images = []\n",
    "for im_dir in os.listdir(image_paths):\n",
    "    try:\n",
    "        images.append(mpimg.imread(os.path.join(image_paths,im_dir)))\n",
    "    except:\n",
    "        pass\n",
    "images = np.array(images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9c238871",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(454, 605, 3)\n",
      "(486, 648, 3)\n",
      "(486, 648, 3)\n",
      "(486, 648, 3)\n",
      "(454, 605, 3)\n",
      "(454, 605, 3)\n",
      "(486, 648, 3)\n",
      "(454, 605, 3)\n",
      "(454, 605, 3)\n",
      "(486, 648, 3)\n",
      "(486, 648, 3)\n",
      "(454, 605, 3)\n",
      "(454, 605, 3)\n",
      "(486, 648, 3)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(14, 824010)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "image_height_and_width = [454, 605]\n",
    "small_images = np.zeros((images.shape[0],image_height_and_width[0]* image_height_and_width[1]* 3))\n",
    "small_to_large_image_size_ratio = 0.15\n",
    "for i in range(images.shape[0]):\n",
    "    small_img = cv2.resize(images[i], # original image\n",
    "                           (0,0), # set fx and fy, not the final size\n",
    "                           fx=small_to_large_image_size_ratio, \n",
    "                           fy=small_to_large_image_size_ratio, \n",
    "                           interpolation=cv2.INTER_NEAREST)\n",
    "    print(small_img.shape)\n",
    "    if small_img.shape[0] == 486:\n",
    "        small_img = cv2.resize(small_img, (0,0), fx = 454/486, fy = 605/648)\n",
    "\n",
    "    small_images[i] = small_img.reshape(-1,)/255\n",
    "\n",
    "small_images.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6d4f2140",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{bmatrix} \n",
       " 1.000 & 0.076 & 0.262 & \\\\\n",
       " 0.076 & 1.000 & 0.240 & \\\\\n",
       " 0.262 & 0.240 & 1.000 & \\\\\n",
       "\\end{bmatrix}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 786936 characters]\n",
      "text/plain": [
       "<Figure size 1296x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "S = small_images[[12,6,3]]\n",
    "Subplot_RGB_images(S, imsize = image_height_and_width, height = 4, width = 18)\n",
    "display_matrix(np.corrcoef(S))\n",
    "# plt.savefig('Original_Images.pdf', format = 'pdf', dpi = 1500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "53b8cc6c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The following is the mixture matrix A\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{bmatrix} \n",
       " -0.363 & 0.650 & 1.757 & \\\\\n",
       " 1.100 & 1.568 & 1.487 & \\\\\n",
       " -1.266 & 0.032 & -0.417 & \\\\\n",
       " -0.822 & 0.643 & 1.260 & \\\\\n",
       " -0.023 & -0.752 & 0.661 & \\\\\n",
       "\\end{bmatrix}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input SNR is : 39.998179949657654\n",
      "Row standard deviation of mixtures : [0.49470436 0.68985465 0.35937612 0.38933594 0.20646696]\n"
     ]
    }
   ],
   "source": [
    "np.random.seed(650)\n",
    "NumberofMixtures = 5\n",
    "NumberofSources = 3\n",
    "\n",
    "A = np.random.randn(NumberofMixtures,NumberofSources)\n",
    "X = np.dot(A,S)\n",
    "SNR = 40\n",
    "\n",
    "X, NoisePart = addWGN(X, SNR, return_noise=True)\n",
    "\n",
    "SNRinp = 10 * np.log10(\n",
    "    np.sum(np.mean((X - NoisePart) ** 2, axis=1))\n",
    "    / np.sum(np.mean(NoisePart**2, axis=1))\n",
    ")\n",
    "\n",
    "X_ = ZeroOneNormalizeColumns(X.T).T\n",
    "print(\"The following is the mixture matrix A\")\n",
    "display_matrix(A)\n",
    "print(\"Input SNR is : {}\".format(SNRinp))\n",
    "\n",
    "print(\"Row standard deviation of mixtures : {}\".format(X.std(1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a3517060",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 527452 characters]\n",
      "text/plain": [
       "<Figure size 1296x216 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "Subplot_RGB_images(X_, imsize = image_height_and_width, height = 3, width = 18)\n",
    "# plt.savefig('Mixture_Images.pdf', format = 'pdf', dpi = 1500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "31e09efc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 23728 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 824010/824010 [07:18<00:00, 1881.00it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 23728 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "lambday = 1 - 1e-1/15\n",
    "lambdae = 0.5\n",
    "\n",
    "s_dim = S.shape[0]\n",
    "x_dim = X.shape[0]\n",
    "\n",
    "n_samples = X.shape[1]\n",
    "n_iter = n_samples\n",
    "ppf = X.shape[1]\n",
    "Wf_arr = []\n",
    "\n",
    "# Inverse output covariance\n",
    "By = 1 * np.eye(s_dim)\n",
    "# Inverse error covariance\n",
    "Be = 100 * np.eye(s_dim)\n",
    "\n",
    "debug_iteration_point = 2500\n",
    "model = CorInfoMaxVideoSeparation(  s_dim = s_dim, x_dim = x_dim, muW = 50*1e-3, lambday = lambday,\n",
    "                                    lambdae = lambdae, By = By, Be = Be, neural_OUTPUT_COMP_TOL = 1e-6,\n",
    "                                    set_ground_truth = True, S = S, A = A)\n",
    "\n",
    "_ = model.seperate_videos( Wf_list = Wf_arr, n_pixel_per_frame = ppf, X = X.reshape(1,NumberofMixtures, -1), \n",
    "                       n_iter = n_iter, neural_dynamic_iterations = 500,\n",
    "                       plot_in_jupyter = True, neural_lr_start = 0.5,\n",
    "                       neural_lr_stop = 0.001, debug_iteration_point = debug_iteration_point, \n",
    "                       shuffle = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "37a93bd4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Component PSNR Values : [32.45042381987587, 29.724934762179863, 32.377582098129764]\n",
      "\n",
      "Overall SINR : 30.363729263125514\n"
     ]
    }
   ],
   "source": [
    "Wf = -1*model.compute_overall_mapping(return_mapping = True)\n",
    "OverallMatrix = Wf @ A\n",
    "perm = np.argmax(OverallMatrix, axis = 0)\n",
    "Y = (Wf @ X)[perm]\n",
    "Y = ZeroOneNormalizeColumns(Y.T).T\n",
    "PSNR_levels = []\n",
    "for kk in range(S.shape[0]):\n",
    "    PSNR_levels.append(psnr(S[kk], Y[kk]))\n",
    "\n",
    "SINR = 10*np.log10(CalculateSINRjit(Y, S, False)[0])\n",
    "print(\"Component PSNR Values : {}\\n\".format(PSNR_levels))\n",
    "print(\"Overall SINR : {}\".format(SINR))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8db70e51",
   "metadata": {},
   "source": [
    "# Mixtures vs Reconstructed vs Ground Truth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "72194f25",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 527452 characters]\n",
      "text/plain": [
       "<Figure size 1296x216 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "Subplot_RGB_images(X_, imsize = image_height_and_width, height = 3, width = 18)\n",
    "plt.savefig(os.path.join(results_path,'Image_Mixtures.pdf'), format = 'pdf', dpi = 1500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e2a17b9b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 816808 characters]\n",
      "text/plain": [
       "<Figure size 1296x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "Subplot_RGB_images(ZeroOneNormalizeColumns(Y.T).T, imsize = image_height_and_width, height = 4, width = 18)\n",
    "plt.savefig(os.path.join(results_path,'CorInfoMax_Output.pdf'), format = 'pdf', dpi = 1500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "15dd75ef",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 816808 characters]\n",
      "text/plain": [
       "<Figure size 1296x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "Subplot_RGB_images(np.clip(Y, 0,1), imsize = image_height_and_width, height = 4, width = 18)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "bc4f1d09",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 816808 characters]\n",
      "text/plain": [
       "<Figure size 1296x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "Subplot_RGB_images(Y, imsize = image_height_and_width, height = 4, width = 18)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f0e27759",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 786936 characters]\n",
      "text/plain": [
       "<Figure size 1296x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "Subplot_RGB_images(S, imsize = image_height_and_width, height = 4, width = 18)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
