{
 "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",
    "from dsp_utils import psnr\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": [
      " 36%|███▌      | 108/300 [00:01<00:02, 86.88it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The history saving thread hit an unexpected error (OperationalError('disk I/O error')).History will not be written to the database.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 300/300 [00:03<00:00, 86.99it/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",
       " 1.198 & 1.020 & 1.273 & \\\\\n",
       " 0.367 & 1.364 & 0.901 & \\\\\n",
       " 0.100 & 0.869 & 0.627 & \\\\\n",
       " 1.257 & 0.859 & 0.015 & \\\\\n",
       " 0.957 & 0.789 & 0.592 & \\\\\n",
       "\\end{bmatrix}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Take three sources and create five mixtures\n",
    "n_mixtures = 5\n",
    "n_sources = 3\n",
    "\n",
    "np.random.seed(250)\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:13<00:00, 22.67it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(300, 5, 691200)\n",
      "(300, 3, 691200)\n",
      "(300, 5, 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, 299536 characters]\n",
      "text/plain": [
       "<Figure size 1080x216 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 306532 characters]\n",
      "text/plain": [
       "<Figure size 1080x216 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 303876 characters]\n",
      "text/plain": [
       "<Figure size 1080x216 with 5 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": [
      "OpenCV: FFMPEG: tag 0x5634504d/'MP4V' is not supported with codec id 12 and format 'mp4 / MP4 (MPEG-4 Part 14)'\n",
      "OpenCV: FFMPEG: fallback to use tag 0x7634706d/'mp4v'\n",
      "OpenCV: FFMPEG: tag 0x5634504d/'MP4V' is not supported with codec id 12 and format 'mp4 / MP4 (MPEG-4 Part 14)'\n",
      "OpenCV: FFMPEG: fallback to use tag 0x7634706d/'mp4v'\n",
      "OpenCV: FFMPEG: tag 0x5634504d/'MP4V' is not supported with codec id 12 and format 'mp4 / MP4 (MPEG-4 Part 14)'\n",
      "OpenCV: FFMPEG: fallback to use tag 0x7634706d/'mp4v'\n",
      "100%|██████████| 300/300 [00:02<00:00, 100.16it/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",
    "mix1 = cv2.VideoWriter(os.path.join(mix_dir, \"mix1.mp4\"), cv2.VideoWriter_fourcc(*'MP4V'), 30, (640, 360))\n",
    "mix2 = cv2.VideoWriter(os.path.join(mix_dir, \"mix2.mp4\"), cv2.VideoWriter_fourcc(*'MP4V'), 30, (640, 360))\n",
    "mix3 = cv2.VideoWriter(os.path.join(mix_dir, \"mix3.mp4\"), cv2.VideoWriter_fourcc(*'MP4V'), 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": "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 = 650000\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": 10,
   "id": "b6666fca",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 21500 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 650000/650000 [03:11<00:00, 3396.79it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 21500 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.5,\n",
    "                           neural_lr_stop = 0.001, debug_iteration_point = debug_iteration_point, \n",
    "                           shuffle = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "66db4c31",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(32.77791435632643, 300)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.SIR_list[-1], len(model.SIR_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "d48d714a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((3, 691200), (5, 691200), (3, 691200), (300, 3, 5))"
      ]
     },
     "execution_count": 12,
     "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": 13,
   "id": "39841932",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 300/300 [00:05<00:00, 56.77it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(300, 3, 691200)"
      ]
     },
     "execution_count": 13,
     "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": 14,
   "id": "ac9525e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 300/300 [00:02<00:00, 136.12it/s]\n"
     ]
    }
   ],
   "source": [
    "psnr_list_intime = []\n",
    "for i in tqdm(range(FRAMES)):\n",
    "    psnr0 = psnr(Y_reconstructed_frames[i,0,:], S_arr[i,0,:])\n",
    "    psnr1 = psnr(Y_reconstructed_frames[i,1,:], S_arr[i,1,:])\n",
    "    psnr2 = psnr(Y_reconstructed_frames[i,2,:], S_arr[i,2,:])\n",
    "    psnr_list_intime.append([psnr0, psnr1, psnr2])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "92b53e93",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 35212 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(np.array(psnr_list_intime))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9b24a58e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 300/300 [00:04<00:00, 73.71it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(300, 3, 691200)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Y_reconstructed_frames_final = [ZeroOneNormalizeColumns(-(Wf_arr[-1] @ X_arr[i,:,:]).T).T for i in tqdm(range(FRAMES))]\n",
    "Y_reconstructed_frames_final = np.array(Y_reconstructed_frames_final)\n",
    "Y_reconstructed_frames_final.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "83a35962",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 300/300 [00:01<00:00, 186.80it/s]\n"
     ]
    }
   ],
   "source": [
    "psnr_list = []\n",
    "for i in tqdm(range(FRAMES)):\n",
    "    psnr0 = psnr(Y_reconstructed_frames_final[i,0,:], S_arr[i,0,:])\n",
    "    psnr1 = psnr(Y_reconstructed_frames_final[i,1,:], S_arr[i,1,:])\n",
    "    psnr2 = psnr(Y_reconstructed_frames_final[i,2,:], S_arr[i,2,:])\n",
    "    psnr_list.append([psnr0, psnr1, psnr2])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "cb6b1f46",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 39036 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(np.array(psnr_list))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "eea869dd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([35.60922568, 48.07245908, 44.58145597])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.mean(np.array(psnr_list), 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "244cceca",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 477904 characters]\n",
      "text/plain": [
       "<Figure size 1080x1080 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 479552 characters]\n",
      "text/plain": [
       "<Figure size 1080x1080 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 473652 characters]\n",
      "text/plain": [
       "<Figure size 1080x1080 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 479760 characters]\n",
      "text/plain": [
       "<Figure size 1080x1080 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 481020 characters]\n",
      "text/plain": [
       "<Figure size 1080x1080 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 468760 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": "code",
   "execution_count": 21,
   "id": "af4a9733",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(\"Figures\"):\n",
    "    os.mkdir(\"Figures\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "0f6cc8a4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 469928 characters]\n",
      "text/plain": [
       "<Figure size 1080x1080 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 304524 characters]\n",
      "text/plain": [
       "<Figure size 1080x1080 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 468760 characters]\n",
      "text/plain": [
       "<Figure size 1080x1080 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "Subplot_RGB_images(S_arr[-1], imsize = [360, 640])\n",
    "plt.savefig('Figures/Original_Videos_Last_Frames.pdf', format='pdf')\n",
    "Subplot_RGB_images(X_plot_arr[-1], imsize = [360, 640])\n",
    "plt.savefig('Figures/Mixed_Videos_Last_Frames.pdf', format='pdf')\n",
    "Subplot_RGB_images(Y_reconstructed_frames[-1], imsize = [360, 640])\n",
    "plt.savefig('Figures/Separated_Videos_Last_Frames.pdf', format='pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba868671",
   "metadata": {},
   "source": [
    "# Save Seperated Videos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "60aa28c8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "OpenCV: FFMPEG: tag 0x5634504d/'MP4V' is not supported with codec id 12 and format 'mp4 / MP4 (MPEG-4 Part 14)'\n",
      "OpenCV: FFMPEG: fallback to use tag 0x7634706d/'mp4v'\n",
      "OpenCV: FFMPEG: tag 0x5634504d/'MP4V' is not supported with codec id 12 and format 'mp4 / MP4 (MPEG-4 Part 14)'\n",
      "OpenCV: FFMPEG: fallback to use tag 0x7634706d/'mp4v'\n",
      "OpenCV: FFMPEG: tag 0x5634504d/'MP4V' is not supported with codec id 12 and format 'mp4 / MP4 (MPEG-4 Part 14)'\n",
      "OpenCV: FFMPEG: fallback to use tag 0x7634706d/'mp4v'\n",
      "100%|██████████| 300/300 [00:04<00:00, 69.62it/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",
    "sep1 = cv2.VideoWriter(os.path.join(sep_dir, \"sep1.mp4\"), cv2.VideoWriter_fourcc(*\"MP4V\"), 30, (640, 360))\n",
    "sep2 = cv2.VideoWriter(os.path.join(sep_dir, \"sep2.mp4\"), cv2.VideoWriter_fourcc(*\"MP4V\"), 30, (640, 360))\n",
    "sep3 = cv2.VideoWriter(os.path.join(sep_dir, \"sep3.mp4\"), cv2.VideoWriter_fourcc(*\"MP4V\"), 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": 24,
   "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": 25,
   "id": "fd43df30",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "OpenCV: FFMPEG: tag 0x5634504d/'MP4V' is not supported with codec id 12 and format 'mp4 / MP4 (MPEG-4 Part 14)'\n",
      "OpenCV: FFMPEG: fallback to use tag 0x7634706d/'mp4v'\n",
      "100%|██████████| 300/300 [04:03<00:00,  1.23it/s]\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",
    "final_vid = cv2.VideoWriter(\n",
    "    os.path.join(CURRENT_DIR, \"FINAL.mp4\"), \n",
    "    cv2.VideoWriter_fourcc(*\"MP4V\"), 30, (648, 432)) # 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}.jpg\"))\n",
    "\n",
    "    img = Image.open(os.path.join(data_dir, f\"FINAL-{i+1}.jpg\"))\n",
    "    rgb_img = img.convert('RGB')\n",
    "    rgb_img.save(os.path.join(data_dir, f\"FINAL-{i+1}.jpg\"))\n",
    "    \n",
    "#     final_vid.write(cv2.imread(os.path.join(data_dir, f\"FINAL-{i+1}.jpg\")))\n",
    "    final_vid.write(cv2.normalize(cv2.cvtColor(np.array(rgb_img).astype(\"float32\"), cv2.COLOR_RGB2BGR), None, 255, 0, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U))\n",
    "    os.remove(os.path.join(data_dir, f\"FINAL-{i+1}.jpg\"))\n",
    "    plt.close(fig)\n",
    "final_vid.release()"
   ]
  }
 ],
 "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
}
