{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "892a0987",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys \n",
    "sys.path.insert(0, \"../../src\")\n",
    "from normalizevideos import CURRENT_DIR, data_dir\n",
    "\n",
    "from CorInfoMaxBSS import OnlineCorInfoMax, CorInfoMaxVideoSeparation\n",
    "from general_utils import display_matrix\n",
    "from bss_utils import ZeroOneNormalizeColumns\n",
    "from visualization_utils import Subplot_RGB_images\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.image as mpimg\n",
    "import numpy as np\n",
    "import os\n",
    "import cv2\n",
    "import pylab as pl\n",
    "from tqdm import tqdm\n",
    "from IPython.display import display, Latex, Math, clear_output\n",
    "from IPython import display as display1\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "from PIL import Image\n",
    "import zipfile\n",
    "\n",
    "FRAMES = 300"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "764610a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "ship_img = []\n",
    "bird_img = []\n",
    "guitar_img = []\n",
    "\n",
    "for i in range(FRAMES):\n",
    "    ship_img.append(mpimg.imread(os.path.join(CURRENT_DIR, \"Frames\", \"SHIP\", f\"{i+1}.jpg\")))\n",
    "    bird_img.append(mpimg.imread(os.path.join(CURRENT_DIR, \"Frames\", \"BIRD\", f\"{i+1}.jpg\")))\n",
    "    guitar_img.append(mpimg.imread(os.path.join(CURRENT_DIR, \"Frames\", \"GUITAR\", f\"{i+1}.jpg\")))\n",
    "\n",
    "ship_img = np.array(ship_img)\n",
    "bird_img = np.array(bird_img)\n",
    "guitar_img = np.array(guitar_img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f7f6ca1b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 300/300 [00:03<00:00, 83.13it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(300, 3, 691200)\n"
     ]
    }
   ],
   "source": [
    "frame_sizes = [360, 640]\n",
    "ship_image_shape = ship_img.shape\n",
    "bird_image_shape = bird_img.shape\n",
    "guitar_image_shape = guitar_img.shape\n",
    "# An array to hold the corresponding frames of the videos\n",
    "# for ex: [bird_img[0], drift_img[0], guitar_img[0]]\n",
    "same_frame = []\n",
    "small_images = []\n",
    "for i in tqdm(range(FRAMES)):\n",
    "    ship_img_ = cv2.resize(ship_img[i], # original image\n",
    "                           (0,0), # set fx and fy, not the final size\n",
    "                           fx=frame_sizes[0]/ship_image_shape[1],\n",
    "                           fy=frame_sizes[1]/ship_image_shape[2], \n",
    "                           interpolation=cv2.INTER_NEAREST)\n",
    "    \n",
    "    bird_img_ = cv2.resize(bird_img[i], # original image\n",
    "                           (0,0), # set fx and fy, not the final size\n",
    "                           fx=frame_sizes[0]/bird_image_shape[1],\n",
    "                           fy=frame_sizes[1]/bird_image_shape[2], \n",
    "                           interpolation=cv2.INTER_NEAREST)\n",
    "    \n",
    "    guitar_img_ = cv2.resize(guitar_img[i], # original image\n",
    "                           (0,0), # set fx and fy, not the final size\n",
    "                           fx=frame_sizes[0]/guitar_image_shape[1],\n",
    "                           fy=frame_sizes[1]/guitar_image_shape[2], \n",
    "                           interpolation=cv2.INTER_NEAREST)\n",
    "    \n",
    "    same_frame.append([ship_img_, bird_img_, guitar_img_])\n",
    "    small_images.append([same_frame[i][0].reshape(-1,)/255, same_frame[i][1].reshape(-1,)/255, same_frame[i][2].reshape(-1,)/255])\n",
    "    \n",
    "small_images = np.array(small_images)\n",
    "print(small_images.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "61547e36",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Coefficient Correlation Matrix of the first frames:\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{bmatrix} \n",
       " 1.000 & -0.024 & -0.047 & \\\\\n",
       " -0.024 & 1.000 & 0.060 & \\\\\n",
       " -0.047 & 0.060 & 1.000 & \\\\\n",
       "\\end{bmatrix}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "test_s = small_images[0][[0, 1, 2]]\n",
    "print(\"Coefficient Correlation Matrix of the first frames:\")\n",
    "display_matrix(np.corrcoef(test_s))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d0a7ca5f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mixture Matrix:\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{bmatrix} \n",
       " 0.100 & 0.343 & 1.153 & \\\\\n",
       " 1.597 & 0.981 & 0.514 & \\\\\n",
       " 0.221 & 0.780 & 1.660 & \\\\\n",
       " 0.255 & 1.392 & 0.435 & \\\\\n",
       " 1.266 & 0.817 & 0.673 & \\\\\n",
       " 1.745 & 1.318 & 1.030 & \\\\\n",
       "\\end{bmatrix}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Take three sources and create six mixtures\n",
    "n_mixtures = 6\n",
    "n_sources = 3\n",
    "\n",
    "np.random.seed(100)\n",
    "\n",
    "A = np.random.standard_normal((n_mixtures, n_sources))\n",
    "minA = np.min(A)\n",
    "min_abs = np.abs(minA)\n",
    "\n",
    "for i in range(n_mixtures):\n",
    "    for j in range(n_sources):\n",
    "        if A[i, j] <= 0:\n",
    "            A[i, j] += min_abs + 0.1\n",
    "        elif A[i, j] + minA >= 0:\n",
    "            A[i, j] += minA\n",
    "        \n",
    "# A = ortho_group.rvs(3) # Make it an orthogonal matrix instead\n",
    "# A = np.random.uniform(0.1, 2, (n_mixtures, n_sources))\n",
    "\n",
    "print(f\"Mixture Matrix:\")\n",
    "display_matrix(A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "43d4523e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 300/300 [00:21<00:00, 14.00it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(300, 6, 691200)\n",
      "(300, 3, 691200)\n",
      "(300, 6, 691200)\n",
      "(300, 3, 691200)\n"
     ]
    }
   ],
   "source": [
    "X_arr = [] # Mixtures for each frame\n",
    "X_plot_arr = []\n",
    "S_arr = [] # Keeps the three S matrices\n",
    "\n",
    "for i in tqdm(range(FRAMES)):\n",
    "    temp = []\n",
    "\n",
    "    s = small_images[i][[0, 1, 2]]\n",
    "    S = s\n",
    "    S_arr.append(S)\n",
    "\n",
    "    X = A @ S\n",
    "\n",
    "    X_plot_arr.append(ZeroOneNormalizeColumns(X.T).T)\n",
    "    temp.append(X)\n",
    "    \n",
    "    X_arr.append(np.array(temp))\n",
    "\n",
    "X_arr = np.array(X_arr)\n",
    "X_arr = X_arr.reshape(FRAMES, n_mixtures, 691200)\n",
    "S_arr = np.array(S_arr)\n",
    "X_plot_arr = np.array(X_plot_arr)\n",
    "S_plot_arr = ZeroOneNormalizeColumns(S_arr)\n",
    "print(X_arr.shape)\n",
    "print(S_arr.shape)\n",
    "print(X_plot_arr.shape)\n",
    "print(S_plot_arr.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b6134340",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 251964 characters]\n",
      "text/plain": [
       "<Figure size 1080x216 with 6 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 255456 characters]\n",
      "text/plain": [
       "<Figure size 1080x216 with 6 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 253032 characters]\n",
      "text/plain": [
       "<Figure size 1080x216 with 6 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "Subplot_RGB_images(X_plot_arr[60], imsize = [360, 640], height=3, width=15) # Second 2\n",
    "Subplot_RGB_images(X_plot_arr[120], imsize = [360, 640], height=3, width=15) # Second 4\n",
    "Subplot_RGB_images(X_plot_arr[180], imsize = [360, 640], height=3, width=15) # Second 6"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05d10350",
   "metadata": {},
   "source": [
    "# Saving 3 Mixture Videos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7aac4894",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 300/300 [00:04<00:00, 65.17it/s]\n"
     ]
    }
   ],
   "source": [
    "mix_dir = os.path.join(CURRENT_DIR, \"Mix-Videos\")\n",
    "if not os.path.exists(mix_dir):\n",
    "    os.mkdir(mix_dir)\n",
    "mix1 = cv2.VideoWriter(os.path.join(mix_dir, \"mix1.avi\"), cv2.VideoWriter_fourcc(*\"MJPG\"), 30, (640, 360))\n",
    "mix2 = cv2.VideoWriter(os.path.join(mix_dir, \"mix2.avi\"), cv2.VideoWriter_fourcc(*\"MJPG\"), 30, (640, 360))\n",
    "mix3 = cv2.VideoWriter(os.path.join(mix_dir, \"mix3.avi\"), cv2.VideoWriter_fourcc(*\"MJPG\"), 30, (640, 360))\n",
    "\n",
    "for i in tqdm(range(FRAMES)):\n",
    "    mix1.write(cv2.normalize(cv2.cvtColor(X_plot_arr[i][0].reshape(360, 640, 3).astype(\"float32\"), cv2.COLOR_RGB2BGR), None, 255, 0, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U))\n",
    "    mix2.write(cv2.normalize(cv2.cvtColor(X_plot_arr[i][1].reshape(360, 640, 3).astype(\"float32\"), cv2.COLOR_RGB2BGR), None, 255, 0, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U))\n",
    "    mix3.write(cv2.normalize(cv2.cvtColor(X_plot_arr[i][2].reshape(360, 640, 3).astype(\"float32\"), cv2.COLOR_RGB2BGR), None, 255, 0, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U))\n",
    "\n",
    "mix1.release()\n",
    "mix2.release()\n",
    "mix3.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "dc3870b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# class LDMIVid(OnlineCorInfoMax):\n",
    "    \n",
    "#     def seperate_videos(self, Wf_list, X, imsize = [360, 640], n_pixel_per_frame=5, \n",
    "#                         n_iter = 1000000, neural_dynamic_iterations = 250, \n",
    "#                         neural_lr_start = 0.9, neural_lr_stop = 1e-3, \n",
    "#                         neural_loop_lr_rule=\"divide_by_loop_index\",\n",
    "#                         neural_lr_decay_multiplier=0.005,\n",
    "#                         use_error_corr_structured_connectivity=False,\n",
    "#                         shuffle=False,\n",
    "#                         debug_iteration_point=1000,\n",
    "#                         plot_in_jupyter=False,):\n",
    "        \n",
    "#         lambday, lambdae, muW, gamy, game, W, By, Be = (\n",
    "#             self.lambday,\n",
    "#             self.lambdae,\n",
    "#             self.muW,\n",
    "#             self.gamy,\n",
    "#             self.game,\n",
    "#             self.W,\n",
    "#             self.By,\n",
    "#             self.Be,\n",
    "#         )        \n",
    "#         neural_dynamic_tol = self.neural_OUTPUT_COMP_TOL\n",
    "#         debugging = self.set_ground_truth\n",
    "#         number_of_mixtures = X.shape[1]\n",
    "#         number_of_frames = X.shape[0]\n",
    "\n",
    "        \n",
    "#         samples = X.shape[2]\n",
    "                \n",
    "#         if debugging:\n",
    "#             SIRlist = []\n",
    "#             A = self.A\n",
    "            \n",
    "#         sample_counter = 0\n",
    "#         for i_sample in tqdm(range(n_iter)):\n",
    "#             x_current = X_arr[np.mod(sample_counter // n_pixel_per_frame, 10),:, np.random.randint(samples)]\n",
    "#             y = np.zeros(self.s_dim)\n",
    "\n",
    "#             y = self.run_neural_dynamics_nnantisparse(\n",
    "#                                                     x_current,\n",
    "#                                                     y,\n",
    "#                                                     W,\n",
    "#                                                     By,\n",
    "#                                                     Be,\n",
    "#                                                     gamy,\n",
    "#                                                     game,\n",
    "#                                                     lr_start=neural_lr_start,\n",
    "#                                                     lr_stop=neural_lr_stop,\n",
    "#                                                     lr_rule=neural_loop_lr_rule,\n",
    "#                                                     lr_decay_multiplier=neural_lr_decay_multiplier,\n",
    "#                                                     neural_dynamic_iterations=neural_dynamic_iterations,\n",
    "#                                                     neural_OUTPUT_COMP_TOL=neural_dynamic_tol,\n",
    "#                                                 )\n",
    "\n",
    "#             e = y - W @ x_current\n",
    "\n",
    "#             W = W + muW * np.outer(e, x_current)\n",
    "\n",
    "#             z = By @ y\n",
    "#             By = (1/lambday) * (By - gamy * np.outer(z, z))\n",
    "            \n",
    "#             if debugging:\n",
    "#                 if (i_sample % debug_iteration_point) == 0:\n",
    "#                     self.W = W\n",
    "#                     self.By = By\n",
    "#                     self.Be = Be\n",
    "#                     Wf = self.compute_overall_mapping(return_mapping = True)\n",
    "#                     Wf_list.append(Wf)\n",
    "#                     SIR = self.CalculateSIR(A, Wf)[0]\n",
    "#                     SIRlist.append(SIR)\n",
    "#                     self.SIR_list = SIRlist\n",
    "#                     if plot_in_jupyter:\n",
    "#                         pl.clf()\n",
    "#                         pl.plot(np.array(SIRlist), linewidth = 3)\n",
    "#                         pl.xlabel(\"Number of Iterations / {}\".format(debug_iteration_point), fontsize = 15)\n",
    "#                         pl.ylabel(\"SIR (dB)\", fontsize = 15)\n",
    "#                         pl.title(\"SIR Behaviour\", fontsize = 15)\n",
    "#                         pl.grid()\n",
    "#                         clear_output(wait=True)\n",
    "#                         display(pl.gcf())   \n",
    "#             sample_counter += 1\n",
    "#         return X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0ad8fb97",
   "metadata": {},
   "outputs": [],
   "source": [
    "lambday = 1 - 1e-1/5\n",
    "lambdae = 0.3\n",
    "\n",
    "s_dim = S.shape[0]\n",
    "x_dim = X.shape[0]\n",
    "\n",
    "n_iter = 352500\n",
    "ppf = 20\n",
    "Wf_arr = []\n",
    "\n",
    "# Inverse output covariance\n",
    "By = 1 * np.eye(s_dim)\n",
    "# Inverse error covariance\n",
    "Be = 65 * np.eye(s_dim)\n",
    "\n",
    "debug_iteration_point = n_iter // (FRAMES - 1)\n",
    "model = CorInfoMaxVideoSeparation(  s_dim = s_dim, x_dim = x_dim, muW = 60*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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "b6666fca",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 22908 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 352500/352500 [03:32<00:00, 1654.96it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 22908 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "aa = model.seperate_videos(Wf_list = Wf_arr, n_pixel_per_frame=ppf, X = X_arr, \n",
    "                           n_iter = n_iter, neural_dynamic_iterations = 500,\n",
    "                           plot_in_jupyter = True, neural_lr_start = 0.25,\n",
    "                           neural_lr_stop = 0.001, debug_iteration_point = debug_iteration_point, \n",
    "                           shuffle = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "66db4c31",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "30.3179485488743"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.SIR_list[299]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d48d714a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((3, 691200), (6, 691200), (3, 691200), (641, 3, 6))"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X = X_arr[0]\n",
    "S = S_arr[0]\n",
    "# Wf = model.compute_overall_mapping(return_mapping = True)\n",
    "Wf_arr = np.array(Wf_arr)\n",
    "Y = Wf_arr[-1] @ X\n",
    "Y.shape, X.shape, S.shape, Wf_arr.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "39841932",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 300/300 [00:06<00:00, 45.24it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(300, 3, 691200)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Y_reconstructed_frames = [ZeroOneNormalizeColumns(-(Wf_arr[i] @ X_arr[i,:,:]).T).T for i in tqdm(range(FRAMES))]\n",
    "Y_reconstructed_frames = np.array(Y_reconstructed_frames)\n",
    "Y_reconstructed_frames.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "244cceca",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 454644 characters]\n",
      "text/plain": [
       "<Figure size 1080x1080 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 452588 characters]\n",
      "text/plain": [
       "<Figure size 1080x1080 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 453480 characters]\n",
      "text/plain": [
       "<Figure size 1080x1080 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 460880 characters]\n",
      "text/plain": [
       "<Figure size 1080x1080 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 461240 characters]\n",
      "text/plain": [
       "<Figure size 1080x1080 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 467792 characters]\n",
      "text/plain": [
       "<Figure size 1080x1080 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "Subplot_RGB_images(Y_reconstructed_frames[1], imsize = [360, 640])\n",
    "Subplot_RGB_images(Y_reconstructed_frames[25], imsize = [360, 640])\n",
    "Subplot_RGB_images(Y_reconstructed_frames[50], imsize = [360, 640])\n",
    "Subplot_RGB_images(Y_reconstructed_frames[80], imsize = [360, 640])\n",
    "Subplot_RGB_images(Y_reconstructed_frames[100], imsize = [360, 640])\n",
    "Subplot_RGB_images(Y_reconstructed_frames[-1], imsize = [360, 640])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba868671",
   "metadata": {},
   "source": [
    "# Save Seperated Videos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "60aa28c8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 300/300 [00:06<00:00, 45.37it/s]\n"
     ]
    }
   ],
   "source": [
    "sep_dir = os.path.join(CURRENT_DIR, \"Separated-Videos\")\n",
    "if not os.path.exists(sep_dir):\n",
    "    os.mkdir(sep_dir)\n",
    "sep1 = cv2.VideoWriter(os.path.join(sep_dir, \"sep1.avi\"), cv2.VideoWriter_fourcc(*\"MJPG\"), 30, (640, 360))\n",
    "sep2 = cv2.VideoWriter(os.path.join(sep_dir, \"sep2.avi\"), cv2.VideoWriter_fourcc(*\"MJPG\"), 30, (640, 360))\n",
    "sep3 = cv2.VideoWriter(os.path.join(sep_dir, \"sep3.avi\"), cv2.VideoWriter_fourcc(*\"MJPG\"), 30, (640, 360))\n",
    "\n",
    "for i in tqdm(range(FRAMES)):\n",
    "    sep1.write(cv2.normalize(cv2.cvtColor(Y_reconstructed_frames[i][0].reshape(360, 640, 3).astype(\"float32\"), cv2.COLOR_RGB2BGR), None, 255, 0, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U))\n",
    "    sep2.write(cv2.normalize(cv2.cvtColor(Y_reconstructed_frames[i][1].reshape(360, 640, 3).astype(\"float32\"), cv2.COLOR_RGB2BGR), None, 255, 0, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U))\n",
    "    sep3.write(cv2.normalize(cv2.cvtColor(Y_reconstructed_frames[i][2].reshape(360, 640, 3).astype(\"float32\"), cv2.COLOR_RGB2BGR), None, 255, 0, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U))\n",
    "    \n",
    "sep1.release()\n",
    "sep2.release()\n",
    "sep3.release()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8aaf49f1",
   "metadata": {},
   "source": [
    "# Creating the Final Video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "0a2e0857",
   "metadata": {},
   "outputs": [],
   "source": [
    "DPI = 162\n",
    "\n",
    "# Modified Subplot_RGB_Images()\n",
    "def final_vid_subplot(I1, I2, I3, imsize = [360, 640], height = 6, width = 9, title = ''):\n",
    "    n_images = 3\n",
    "    images = (I1, I2, I3)\n",
    "    titles = (\"Source Video\", \"Mixture Video\", \"Recovered Video\")\n",
    "    \n",
    "    fig, ax = plt.subplots(3, n_images, figsize = (9, 16)) \n",
    "    fig.patch.set_facecolor(\"black\")\n",
    "    fig.suptitle(title)\n",
    "    fig.set_figheight(height)\n",
    "    fig.set_figwidth(width) \n",
    "    for j in range(n_images):\n",
    "        for i in range(n_images):\n",
    "            Im = [images[j][i,:].reshape(imsize[0],imsize[1], 3) for i in range(n_images)]\n",
    "            ax[j][i].set_title(f\"{titles[j]} {i+1}\", color=\"white\", fontsize=15)\n",
    "            ax[j][i].imshow(Im[i])\n",
    "            ax[j][i].axes.xaxis.set_visible(False)\n",
    "            ax[j][i].axes.yaxis.set_visible(False)\n",
    "    plt.subplots_adjust(right=0.97,\\\n",
    "                    left=0.03,\\\n",
    "                    bottom=0.03,\\\n",
    "                    top=0.9,\\\n",
    "                    wspace=0.1,\\\n",
    "                    hspace=0.4)\n",
    "    fig.dpi = DPI \n",
    "    return fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "fd43df30",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 300/300 [05:15<00:00,  1.05s/it]\n"
     ]
    }
   ],
   "source": [
    "final_vid = cv2.VideoWriter(\n",
    "    os.path.join(CURRENT_DIR, \"FINAL.avi\"), \n",
    "    cv2.VideoWriter_fourcc(*\"MJPG\"), 30, (DPI * 9, DPI * 6)) # dpi * 9, dpi * 6\n",
    "\n",
    "for i in tqdm(range(FRAMES)):\n",
    "    fig = final_vid_subplot(S_plot_arr[i], X_plot_arr[i], Y_reconstructed_frames[i])\n",
    "    plt.savefig(os.path.join(data_dir, f\"FINAL-{i+1}.png\"))\n",
    "\n",
    "    img = Image.open(os.path.join(data_dir, f\"FINAL-{i+1}.png\"))\n",
    "    rgb_img = img.convert('RGB')\n",
    "    rgb_img.save(os.path.join(data_dir, f\"FINAL-{i+1}.jpg\"))\n",
    "    os.remove(os.path.join(data_dir, f\"FINAL-{i+1}.png\"))\n",
    "    \n",
    "    final_vid.write(cv2.imread(os.path.join(data_dir, f\"FINAL-{i+1}.jpg\")))\n",
    "    os.remove(os.path.join(data_dir, f\"FINAL-{i+1}.jpg\"))\n",
    "    plt.close(fig)\n",
    "final_vid.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34878962",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}
