inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zjnYK-jckO5J"
   },
   "source": [
    "# Start from here if git repo exists (or if you restart runtime)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 313,
     "status": "ok",
     "timestamp": 1672889892044,
     "user": {
      "displayName": "Pradyumna Chari",
      "userId": "06555723454452599039"
     },
     "user_tz": 480
    },
    "id": "3fDtcMQZb23t",
    "outputId": "d619e7a4-17c3-420b-ad0f-661799cab14a"
   },
   "outputs": [],
   "source": [
    "# Checkout to the branch containing the correct code\n",
    "# (Currently old_method_all_perm). See the output to be sure\n",
    "%cd mmhealth_postprocess\n",
    "!git checkout old_method_all_perm\n",
    "!git branch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pJE91R9OkYnv"
   },
   "outputs": [],
   "source": [
    "# !pip install imageio-ffmpeg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-D9s_s2wj-eI"
   },
   "source": [
    "# Python code starts here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 2321,
     "status": "ok",
     "timestamp": 1672970717147,
     "user": {
      "displayName": "Pradyumna Chari",
      "userId": "06555723454452599039"
     },
     "user_tz": 480
    },
    "id": "IlPV8Ajddobz",
    "outputId": "9785a7e6-d9a2-40ed-8b9e-b0e91cab74a6"
   },
   "outputs": [],
   "source": [
    "import numpy as np \n",
    "import matplotlib.pyplot as plt \n",
    "import os\n",
    "import torch\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "import torchvision.transforms as T\n",
    "import h5py\n",
    "from tqdm import tqdm\n",
    "import imageio\n",
    "import scipy\n",
    "import random\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"2\"\n",
    "\n",
    "from uni_modal.cameras.PhysNet_NewLoader.heart_rate_extract import prpsd, getErrors\n",
    "\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "# put a test tensor on the device\n",
    "x = torch.tensor([1., 2.], device=device)\n",
    "print('Running on device: {}'.format(device))\n",
    "print(torch.cuda.current_device())\n",
    "\n",
    "#seeds\n",
    "sd = 42\n",
    "# torch.manual_seed(seed)\n",
    "# np.random.seed(seed)\n",
    "\n",
    "np.random.seed(sd)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "torch.manual_seed(sd)\n",
    "random.seed(sd)\n",
    "if torch.cuda.is_available():\n",
    "  torch.cuda.manual_seed_all(sd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BVH-WL5kxA4f"
   },
   "outputs": [],
   "source": [
    "import math\n",
    "import torch.nn as nn\n",
    "from torch.nn.modules.utils import _triple\n",
    "import torch\n",
    "\n",
    "from __future__ import print_function, division\n",
    "from torchvision import transforms\n",
    "\n",
    "class SNRLossOnPreComputedAndWindowedFFT_base(nn.Module):\n",
    "    def __init__(self, start_idx, window_fraction=0.02, device=torch.device('cpu')):\n",
    "        super(SNRLossOnPreComputedAndWindowedFFT_base, self).__init__()\n",
    "        self.start_idx = start_idx\n",
    "        self.window_fraction = window_fraction\n",
    "        self.device = device\n",
    "        \n",
    "    def forward(self, outputs: torch.Tensor, targets: torch.Tensor):\n",
    "        assert outputs.shape == targets.shape, \"The predicted output and the target labels have different shapes\"\n",
    "        if not outputs.is_cuda:\n",
    "            torch.backends.mkl.is_available()\n",
    "        # The window length for the calculating the singal to noise ratio of the output's PSD\n",
    "        # Here we take 5% in tota, i.e. 2.5% on both sides\n",
    "        window_size_base_harmonic = int(self.window_fraction * targets.shape[1]) + 1\n",
    "\n",
    "        # Get the strongest peaks from the output, i.e. the heart beats from the PSD of the targets\n",
    "        Y2 = torch.abs(targets) ** 2\n",
    "        HRixs = torch.argmax(Y2,axis=1)\n",
    "\n",
    "        # get the PSD of the outputs of the neural network\n",
    "        X2 = torch.abs(outputs) ** 2\n",
    "\n",
    "        # calc SNR for each batch\n",
    "        losses = torch.empty((X2.shape[0],), dtype=torch.float32)\n",
    "        for count, ref_idx in enumerate(HRixs):\n",
    "            # Compute the power around the heart beat idx and its fist harmonic => Signal\n",
    "            base_start = max([0, ref_idx - window_size_base_harmonic])\n",
    "            base_end   = ref_idx + window_size_base_harmonic + 1\n",
    "            harmonic_start = max([0, (2 * ref_idx) + self.start_idx - window_size_base_harmonic])\n",
    "            harmonic_end   = (2 * ref_idx) + self.start_idx + window_size_base_harmonic + 1\n",
    "            pulse_freq_amp = torch.sum(X2[count, base_start : base_end])#+torch.sum(X2[count, harmonic_start : harmonic_end])\n",
    "            # Compute the power outisde the above windows => Noise\n",
    "            other_avrg = torch.sum(X2[count, :base_start]) + torch.sum(X2[count, base_end:])#+torch.sum(X2[count, base_end : harmonic_start]) + torch.sum(X2[count, harmonic_end:]))\n",
    "            # Take the SNR loss in decibels\n",
    "            losses[count] = -10*torch.log10(pulse_freq_amp/(other_avrg+1e-7))\n",
    "        losses.to(self.device)\n",
    "        return torch.mean(losses)\n",
    "\n",
    "\n",
    "\n",
    "'''\n",
    "Code of 'Remote Photoplethysmograph Signal Measurement from Facial Videos Using Spatio-Temporal Networks' \n",
    "By Zitong Yu, 2019/05/05\n",
    "If you use the code, please cite:\n",
    "@inproceedings{yu2019remote,\n",
    "    title={Remote Photoplethysmograph Signal Measurement from Facial Videos Using Spatio-Temporal Networks},\n",
    "    author={Yu, Zitong and Li, Xiaobai and Zhao, Guoying},\n",
    "    booktitle= {British Machine Vision Conference (BMVC)},\n",
    "    year = {2019}\n",
    "}\n",
    "Only for research purpose, and commercial use is not allowed.\n",
    "MIT License\n",
    "Copyright (c) 2019 \n",
    "      How to use it\n",
    "    #1. Inference the model\n",
    "    rPPG, x_visual, x_visual3232, x_visual1616 = model(inputs)\n",
    "    \n",
    "    #2. Normalized the Predicted rPPG signal and GroundTruth BVP signal\n",
    "    rPPG = (rPPG-torch.mean(rPPG)) /torch.std(rPPG)\t \t# normalize\n",
    "    BVP_label = (BVP_label-torch.mean(BVP_label)) /torch.std(BVP_label)\t \t# normalize\n",
    "    \n",
    "    #3. Calculate the loss\n",
    "    loss_ecg = Neg_Pearson(rPPG, BVP_label)\n",
    "'''\n",
    "########################################\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "class Neg_Pearson2(torch.nn.Module):    # Pearson range [-1, 1] so if < 0, abs|loss| ; if >0, 1- loss\n",
    "    def __init__(self):\n",
    "        super(Neg_Pearson2,self).__init__()\n",
    "\n",
    "        self.epsilon = 1e-2\n",
    "\n",
    "        return\n",
    "    def forward(self, preds, labels):       # tensor [Batch, Temporal]\n",
    "        loss = 0\n",
    "        for i in range(preds.shape[0]):\n",
    "            # print(labels[i])\n",
    "            # print(preds[i])\n",
    "            x = normalize_signal2(preds[i])\n",
    "            x = torch.nan_to_num(x,nan=0,posinf=0,neginf=0)\n",
    "            y = normalize_signal2(labels[i])\n",
    "            y = torch.nan_to_num(y,nan=0,posinf=0,neginf=0)\n",
    "\n",
    "            sum_x = torch.nansum(x)                # x\n",
    "            sum_y = torch.nansum(y)               # y\n",
    "            sum_xy = torch.nansum(x*y)         # xy\n",
    "            sum_x2 = torch.nansum(torch.pow(x,2))  # x^2\n",
    "            sum_y2 = torch.nansum(torch.pow(y,2)) # y^2\n",
    "            N = preds.shape[1]\n",
    "            pearson = (N*sum_xy - sum_x*sum_y)/(torch.sqrt((N*sum_x2 - torch.pow(sum_x,2)+self.epsilon)*(N*sum_y2 - torch.pow(sum_y,2)+self.epsilon)))\n",
    "\n",
    "            # print(torch.sqrt((N*sum_x2 - torch.pow(sum_x,2))), (N*sum_y2 - torch.pow(sum_y,2)))\n",
    "            # print(pearson, end=\" \")\n",
    "            #if (pearson>=0).data.cpu().numpy():    # torch.cuda.ByteTensor -->  numpy\n",
    "            #    loss += 1 - pearson\n",
    "            #else:\n",
    "            #    loss += 1 - torch.abs(pearson)\n",
    "            \n",
    "            loss += (1 - pearson)**2\n",
    "            \n",
    "        # print(loss) \n",
    "        loss = loss/preds.shape[0]\n",
    "        return loss\n",
    "\n",
    "\n",
    "def normalize_signal2(sig):\n",
    "    return (sig-torch.nanmean(sig)) / (torch.std(sig)+1.00e-6)\n",
    "\n",
    "\n",
    "def str2array(str_arr):\n",
    "    temp = [float(i[0:6])*100 for i in str_arr]\n",
    "    return np.array(temp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wpboqDmlSGHS"
   },
   "outputs": [],
   "source": [
    "##UNIMODAL DATALOADER\n",
    "class RppgData2(Dataset):\n",
    "    def __init__(self, datapath, datapath2, datapaths, recording_str, video_length = 900, num_segments = 3, frame_length = 300, fs=30, l_freq_bpm=45, u_freq_bpm=180, fft_resolution = 48) -> None:\n",
    "        #REMOVE RECORDING STRING\n",
    "        #TODO TODO this line is temporary \n",
    "        self.ppg_offset = 0 #25\n",
    "        self.num_samps = 30\n",
    "\n",
    "        self.no_att_frms = 64\n",
    "        #Data structure for videos\n",
    "        self.video_length = video_length\n",
    "        self.num_segments = num_segments\n",
    "        self.part_length = 300#int(video_length/num_segments)\n",
    "        self.datapath = datapath\n",
    "        self.datapath2 = datapath2\n",
    "\n",
    "        self.fs = fs\n",
    "        self.l_freq_bpm = l_freq_bpm\n",
    "        self.u_freq_bpm = u_freq_bpm\n",
    "        self.fft_resolution = fft_resolution\n",
    "        # TODO TODO TODO : To be corrected\n",
    "        # The PPG files for the RGB are stored as rgbd_ppg and not rgbd_rgb_ppg\n",
    "\n",
    "        #load videos and signals\n",
    "        self.video_list = datapaths\n",
    "        self.signal_list = []\n",
    "        self.item_list = []\n",
    "        #load signals\n",
    "        remove_folders = []\n",
    "        for folder in self.video_list:\n",
    "            file_path1 = os.path.join(self.datapath, folder)\n",
    "            file_path1_1 = os.path.join(self.datapath2, 'residual_0_'+folder+'.npy')\n",
    "            file_path2_1 = os.path.join(self.datapath2, 'residual_300_'+folder+'.npy')\n",
    "            file_path3_1 = os.path.join(self.datapath2, 'residual_600_'+folder+'.npy')\n",
    "            # print(file_path)\n",
    "            if (os.path.exists(file_path1) and os.path.exists(file_path1_1) and os.path.exists(file_path2_1) and os.path.exists(file_path3_1)):\n",
    "                #Load the time series signal\n",
    "                # avg_temp = np.zeros((3,self.video_length))\n",
    "                # for frame_start in range(self.num_segments):\n",
    "                #     vid = np.load(os.path.join(self.datapath, 'residual_'+str(int(frame_start*self.part_length))+'_'+folder+'.npy'))\n",
    "                #     temp = np.mean(np.mean(vid,axis=1),axis=1)\n",
    "                #     temp = np.transpose(temp,(1,0))\n",
    "\n",
    "                #     avg_temp[:,frame_start*self.part_length:(frame_start+1)*self.part_length] = temp\n",
    "\n",
    "                #     # for img_idx in range(self.part_length):\n",
    "                #     #     image = vid.get_data(img_idx) #frame_start*self.part_length+\n",
    "\n",
    "                #     #     avg_temp[:,int(frame_start*self.part_length+img_idx)] = np.mean(np.mean(image,axis=0),axis=0)\n",
    "\n",
    "                # # avg_temp = (avg_temp-np.mean(avg_temp,axis=1,keepdims=True))/np.std(avg_temp,axis=1,keepdims=True)\n",
    "                # self.item_list.append(avg_temp)\n",
    "\n",
    "                if(os.path.exists(os.path.join(self.datapath,folder, f\"rgbd_ppg.npy\"))):\n",
    "                    signal = np.load(os.path.join(self.datapath,folder, f\"rgbd_ppg.npy\"))\n",
    "                    signal = str2array(signal)\n",
    "                    mean_temp = np.mean(signal)\n",
    "                    std_temp = np.std(signal)\n",
    "                    signal = (signal - mean_temp)/std_temp\n",
    "                    self.signal_list.append(signal[:self.video_length])\n",
    "                else:\n",
    "                    print(folder, \"ppg doesn't exist.\")\n",
    "                    remove_folders.append(folder)\n",
    "            else:\n",
    "                print(folder, \" doesn't exist.\")\n",
    "                remove_folders.append(folder)\n",
    "\n",
    "        for i in remove_folders:\n",
    "            self.video_list.remove(i)    \n",
    "            print(\"removed\", i)\n",
    "        #process_vital_signs\n",
    "        ### TODO TODO TODO detrend signals\n",
    "        self.signal_list = np.array(self.signal_list)\n",
    "\n",
    "        # self.item_list = np.array(self.item_list)\n",
    "\n",
    "        # Create a list of video number and valid frame nuber to extract the datad from.\n",
    "        self.frame_length = frame_length\n",
    "        self.video_nums = np.arange(0, len(self.video_list))\n",
    "        # self.frame_nums = np.arange(0, self.part_length - frame_length - self.ppg_offset)\n",
    "\n",
    "        #create all possible sampling combinations and put in self.all_idxs\n",
    "        self.all_idxs = []\n",
    "        for num in self.video_nums:\n",
    "            cur_frame_nums = np.random.choice(np.arange(self.video_length - self.frame_length - self.ppg_offset), size=self.num_samps, replace=False)\n",
    "            # cur_frame_nums = np.random.randint(\n",
    "            #   low=0, high = self.video_length - self.frame_length - self.ppg_offset, size = self.num_samps)\n",
    "            \n",
    "            for cur_frame_num in cur_frame_nums:\n",
    "                self.all_idxs.append((num,cur_frame_num))\n",
    "\n",
    "\n",
    "        #FFT parts\n",
    "        # seq_len = self.frame_length*self.fft_resolution\n",
    "        \n",
    "        # freqs_bpm = np.fft.fftfreq(int(seq_len), d=1/self.fs) * 60\n",
    "        # self.l_freq_idx = np.argmin(np.abs(freqs_bpm - self.l_freq_bpm))\n",
    "        # self.u_freq_idx = np.argmin(np.abs(freqs_bpm - self.u_freq_bpm))\n",
    "        # print(self.l_freq_idx, self.u_freq_idx)\n",
    "        # print(freqs_bpm[self.l_freq_idx], freqs_bpm[self.u_freq_idx])\n",
    "        # assert self.l_freq_idx < self.u_freq_idx\n",
    "\n",
    "            \n",
    "            \n",
    "    def __len__(self):\n",
    "        return int(len(self.all_idxs))\n",
    "    def __getitem__(self, idx):\n",
    "        video_number, frame_start = self.all_idxs[idx]\n",
    "        # video_number = self.video_nums[idx]\n",
    "        # frame_start = np.random.choice(self.frame_nums)\n",
    "        #Get video_frames\n",
    "\n",
    "        # item = self.item_list[int(video_number)][:,int(frame_start):int(frame_start+self.frame_length)]\n",
    "        # item = (item-np.mean(item,axis=1,keepdims=True))/np.std(item,axis=1,keepdims=True)\n",
    "        #Get signal\n",
    "        item_sig = self.signal_list[int(video_number)][int(frame_start):int(frame_start+self.frame_length)]\n",
    "\n",
    "        ##Load the 64 necessary frames for attention\n",
    "        folder = self.video_list[video_number]\n",
    "        temp_stor = np.zeros((self.frame_length,128,128,3))\n",
    "        \n",
    "        temp_stor2 = np.zeros((self.video_length,128,128,3))\n",
    "        \n",
    "        ixx = 0\n",
    "        for frm in range(frame_start, frame_start+self.frame_length):\n",
    "            im = imageio.imread(os.path.join(self.datapath, folder,'rgbd_rgb_'+str(frm)+'.png'))\n",
    "            # im = imageio.imread(os.path.join(self.datapath, folder,'rgb_'+str(frm)+'.png'))\n",
    "\n",
    "            temp_stor[ixx] = im\n",
    "            \n",
    "            ixx+=1\n",
    "            \n",
    "        frmOuts = temp_stor\n",
    "            \n",
    "        for fst in range(self.num_segments):\n",
    "            vid = np.load(os.path.join(self.datapath2, 'residual_'+str(int(fst*self.part_length))+'_'+folder+'.npy'))\n",
    "            \n",
    "            if fst<self.num_segments-1:\n",
    "                temp_stor2[fst*self.part_length:(fst+1)*self.part_length] = vid\n",
    "            else:\n",
    "#                 print(fst)\n",
    "#                 print(vid.shape)\n",
    "                temp_stor2[fst*self.part_length:(fst+1)*self.part_length-20] = vid[20:]\n",
    "\n",
    "        frmOuts2 = temp_stor2[frame_start:frame_start+self.frame_length]\n",
    "        \n",
    "        frmOuts = np.concatenate((frmOuts, frmOuts2), axis=-1)\n",
    "\n",
    "\n",
    "        \n",
    "        \n",
    "#         ##Do the random scramble\n",
    "#         w_box = 30\n",
    "#         margin = 0#20\n",
    "        \n",
    "#         ix1 = np.random.randint(margin, high=128-margin-w_box)\n",
    "#         ix2 = np.random.randint(margin, high=128-margin-w_box)\n",
    "        \n",
    "#         #These are box indices\n",
    "        \n",
    "#         temp_shuf = frmOuts[:,ix1:ix1+w_box,ix2:ix2+w_box]\n",
    "        \n",
    "#         np.random.shuffle(temp_shuf)\n",
    "        \n",
    "#         frmOuts[:,ix1:ix1+w_box,ix2:ix2+w_box] = temp_shuf\n",
    "        \n",
    "        \n",
    "\n",
    "        return frmOuts, item_sig\n",
    "        \n",
    "        \n",
    "        # #Patch for torch constructor not accepting uint16 datatypes\n",
    "        # if(item.dtype == np.uint16):\n",
    "        #     item = item.astype(np.int32)\n",
    "        # return np.array(item), np.array(item_sig)\n",
    "\n",
    "\n",
    "    # def lowPassFilter(self, BVP, butter_order=4):\n",
    "    #     [b, a] = sig.butter(butter_order, [self.l_freq_bpm/60, self.u_freq_bpm/60], btype='bandpass', fs = self.fs)\n",
    "    #     filtered_BVP = sig.filtfilt(b, a, np.double(BVP))\n",
    "    #     return filtered_BVP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7z0WpPowS0wg"
   },
   "outputs": [],
   "source": [
    "def extract_video(path, path2, cur_session):\n",
    "    length_seg = 300\n",
    "    vid = np.zeros((900,128,128,3))\n",
    "    \n",
    "    for j in range(900): #number of segments\n",
    "        im = imageio.imread(os.path.join(path, cur_session,'rgbd_rgb_'+str(j)+'.png'))\n",
    "        vid[j] = im\n",
    "    \n",
    "    vid2 = np.zeros((900,128,128,3))\n",
    "    for j in range(3): #number of segments\n",
    "        video_rd = np.load(os.path.join(path2, 'residual_'+str(j*length_seg)+'_'+cur_session+'.npy'))\n",
    "        \n",
    "        if j<5:\n",
    "            vid2[j*length_seg:(j+1)*length_seg] = video_rd\n",
    "        else:\n",
    "            vid2[j*length_seg:(j+1)*length_seg-20] = video_rd[20:]\n",
    "        \n",
    "    #concatenate\n",
    "    vid = np.concatenate((vid, vid2), axis=-1)\n",
    "#             vid.append(np.mean(np.mean(image,axis=0),axis=0))\n",
    "\n",
    "    return vid #should be of shape 3x900\n",
    "\n",
    "##Eval model\n",
    "def eval_model(root_dir, root_dir2, session_names, model, model2, in_frames=64, \n",
    "               hr_window_size = 300, stride = 128, video_fps = 30, ppg_offset = 0, fft_resolution=48):\n",
    "    model.eval()#!!!\n",
    "    model2.eval()\n",
    "    video_samples = []\n",
    "    for cur_session in session_names:\n",
    "        video_sample = {\"video_path\" : root_dir, \"video_path2\" : root_dir2, \"cur_session\" : cur_session}\n",
    "        video_samples.append(video_sample)\n",
    "\n",
    "\n",
    "    #Get indices for FFT\n",
    "    #TODO: band limit dependencies. clean this up\n",
    "    l_freq_bpm = 45\n",
    "    u_freq_bpm = 180\n",
    "\n",
    "    for cur_video_sample in tqdm(video_samples):\n",
    "        cur_video_path = cur_video_sample[\"video_path\"]\n",
    "        cur_video_path2 = cur_video_sample[\"video_path2\"]\n",
    "        cur_session = cur_video_sample[\"cur_session\"]\n",
    "\n",
    "        frames = extract_video(path=cur_video_path, path2=cur_video_path2, cur_session=cur_session) # (900, 128, 128, 3)\n",
    "        target = np.load(os.path.join(cur_video_path, cur_session,'rgbd_ppg.npy'))\n",
    "        target = str2array(target)\n",
    "        target = target[:900]\n",
    "        ##Apply offset to target\n",
    "#         target = target[ppg_offset:]\n",
    "\n",
    "\n",
    "\n",
    "        #Normalize target\n",
    "        target = (target-np.mean(target,axis=0,keepdims=True))/np.std(target,axis=0,keepdims=True)\n",
    "\n",
    "        #get the start indices\n",
    "        start_indices = np.arange(0,frames.shape[0]+1-ppg_offset-hr_window_size,stride)\n",
    "#         print(start_indices)\n",
    "        \n",
    "#         print(frames[890:1200,0,0,0])\n",
    "\n",
    "        batched_ip = []\n",
    "        batched_tgt = []\n",
    "\n",
    "        for ix in start_indices:\n",
    "            temp_ip = frames[ix:ix+hr_window_size]\n",
    "#             temp_ip = (temp_ip-np.mean(temp_ip,axis=1,keepdims=True))/np.std(temp_ip,axis=1,keepdims=True)\n",
    "            #Take FFT\n",
    "            \n",
    "\n",
    "            batched_ip.append(temp_ip)\n",
    "\n",
    "            temp_tgt = target[ix:ix+hr_window_size]\n",
    "            \n",
    "            batched_tgt.append(temp_tgt)\n",
    "\n",
    "        batched_ip = torch.Tensor(np.array(batched_ip)).to(device)\n",
    "        # batched_tgt = torch.Tensor(np.array(batched_tgt)).to(device)\n",
    "        batched_tgt = np.array(batched_tgt) #B,fft_size\n",
    "        # plt.plot(batched_tgt[:,0])\n",
    "        # plt.show()\n",
    "        \n",
    "        \n",
    "\n",
    "\n",
    "        # Potential High GPU usage\n",
    "        # A batch_size of approx 14\n",
    "        \n",
    "\n",
    "        with torch.no_grad():\n",
    "            msk = model2(batched_ip[:,0:64])\n",
    "            est_ppg_ffts = model(batched_ip,msk)\n",
    "            \n",
    "            # (14, 64)\n",
    "            est_ppg_ffts = est_ppg_ffts.squeeze().cpu().numpy() #Size: B,fft_size\n",
    "\n",
    "       # Save\n",
    "        cur_video_sample['est_ppg_ffts'] = est_ppg_ffts\n",
    "        cur_video_sample['gt_ppg_ffts'] = batched_tgt\n",
    "        \n",
    "#         print(est_ppg_ffts[:,0])\n",
    "    print('All finished!')\n",
    "\n",
    "    #Estimate using waveforms\n",
    "    mae_list = []\n",
    "    all_hr_est = []\n",
    "    all_hr_gt = []\n",
    "    for index, cur_video_sample in enumerate(video_samples):\n",
    "        cur_video_path = cur_video_sample['video_path']\n",
    "        # print('Video:', cur_video_path)\n",
    "        est_ppg_ffts = cur_video_sample['est_ppg_ffts']\n",
    "        # Load GT\n",
    "        gt_ppg_ffts = cur_video_sample['gt_ppg_ffts']\n",
    "\n",
    "        #Just need to iterate over batch dimension and\n",
    "\n",
    "        # Get est HR for each window\n",
    "        hr_est_temp = []\n",
    "        hr_gt_temp = []\n",
    "        for ixx in range(est_ppg_ffts.shape[0]):\n",
    "            # plt.plot(est_ppg_ffts[ixx], label='Est')\n",
    "            # plt.plot(gt_ppg_ffts[ixx], label='GT')\n",
    "            # plt.legend()\n",
    "            # plt.show()\n",
    "            est_hr = prpsd2(est_ppg_ffts[ixx], 30, 45, 150, BUTTER_ORDER=6, DETREND=False)\n",
    "            gt_hr = prpsd2(gt_ppg_ffts[ixx], 30, 45, 150, BUTTER_ORDER=6, DETREND=False)\n",
    "            hr_est_temp.append(est_hr)\n",
    "            hr_gt_temp.append(gt_hr)\n",
    "\n",
    "        hr_est_windowed = np.array([hr_est_temp])\n",
    "        hr_gt_windowed = np.array(hr_gt_temp)\n",
    "        all_hr_est.append(hr_est_temp)\n",
    "        all_hr_gt.append(hr_gt_temp)\n",
    "\n",
    "        # Errors\n",
    "        RMSE, MAE, MAX, PCC = getErrors(hr_est_windowed, hr_gt_windowed)\n",
    "\n",
    "        mae_list.append(MAE)\n",
    "    print('Mean MAE:', np.mean(np.array(mae_list)))\n",
    "    return np.array(mae_list), (all_hr_est, all_hr_gt)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # mae_loss_list, hrs = eval_model(root_dir=destination_folder, root_dir2=destination_folder2, session_names=val, model=model, model2=model2)\n",
    "        \n",
    "# # print(hrs[0][50],hrs[1][50])\n",
    "\n",
    "# ## Load numpys'\n",
    "# aa = np.load('/home/pradyumnachari/Documents/rPPG_Implicit_Expt/vitals/dataset_res_npys/residual_900_subject_1_1.npy')\n",
    "\n",
    "# print(aa[:,0,0,0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "U1H2MspBw_Ck"
   },
   "outputs": [],
   "source": [
    "def prpsdFFT(BVP, FS, LL_PR, UL_PR, WIN_SZ, FFT_RES):\n",
    "    #BVP has shape (k,)\n",
    "    \n",
    "    l_freq_bpm = LL_PR\n",
    "    u_freq_bpm = UL_PR\n",
    "    seq_len = WIN_SZ*FFT_RES\n",
    "    freqs_bpm = np.fft.fftfreq(int(seq_len), d=1/FS) * 60\n",
    "\n",
    "    base_ix = np.argmin(np.abs(freqs_bpm - l_freq_bpm))\n",
    "    #We can now index the above to get the final index\n",
    "    #Input is a PSD. We just find the argmax\n",
    "    max_ix = np.argmax(BVP)\n",
    "\n",
    "    hr = freqs_bpm[base_ix+max_ix]\n",
    "            \n",
    "    return hr\n",
    "\n",
    "def prpsd2(BVP, FS, LL_PR, UL_PR, BUTTER_ORDER=6, DETREND=False, PlotTF=False, FResBPM = 0.1):\n",
    "    '''\n",
    "    Estimates pulse rate from the power spectral density a BVP signal\n",
    "    \n",
    "    Inputs\n",
    "        BVP              : A BVP timeseries. (1d numpy array)\n",
    "        fs               : The sample rate of the BVP time series (Hz/fps). (int)\n",
    "        lower_cutoff_bpm : The lower limit for pulse rate (bpm). (int)\n",
    "        upper_cutoff_bpm : The upper limit for pulse rate (bpm). (int)\n",
    "        butter_order     : Order of the Butterworth Filter. (int)\n",
    "        detrend          : Detrend the input signal. (bool)\n",
    "        FResBPM          : Resolution (bpm) of bins in power spectrum used to determine pulse rate and SNR. (float)\n",
    "    \n",
    "    Outputs\n",
    "        pulse_rate       : The estimated pulse rate in BPM. (float)\n",
    "    \n",
    "    Daniel McDuff, Ethan Blackford, January 2019\n",
    "    Copyright (c)\n",
    "    Licensed under the MIT License and the RAIL AI License.\n",
    "    '''\n",
    "    from scipy.signal import butter\n",
    "    from scipy import signal\n",
    "    import numpy as np\n",
    "\n",
    "    N = (60*FS)/FResBPM\n",
    "\n",
    "    # Detrending + nth order butterworth + periodogram\n",
    "#     if DETREND:\n",
    "#         BVP = detrend(np.cumsum(BVP), 100)\n",
    "#     if BUTTER_ORDER:\n",
    "    [b, a] = signal.butter(BUTTER_ORDER, [LL_PR/60, UL_PR/60], btype='bandpass', fs = FS)\n",
    "    \n",
    "    BVP = signal.filtfilt(b, a, np.double(BVP))\n",
    "    \n",
    "    # Calculate the PSD and the mask for the desired range\n",
    "#     if detrend:\n",
    "    F, Pxx = signal.periodogram(x=BVP,  nfft=N, fs=FS, detrend=False);  \n",
    "#     else:\n",
    "#         F, Pxx = signal.periodogram(x=BVP, window=np.hanning(len(BVP)), nfft=N, fs=FS)\n",
    "    FMask = (F >= (LL_PR/60)) & (F <= (UL_PR/60))\n",
    "    \n",
    "    # Calculate predicted pulse rate:\n",
    "    FRange = F * FMask\n",
    "    PRange = Pxx * FMask\n",
    "    MaxInd = np.argmax(PRange)\n",
    "    pulse_rate_freq = FRange[MaxInd]\n",
    "    pulse_rate = pulse_rate_freq*60\n",
    "\n",
    "    # Optionally Plot the PSD and peak frequency\n",
    "    if PlotTF:\n",
    "        # Plot PSD (in dB) and peak frequency\n",
    "        plt.figure()\n",
    "        plt.plot(F, 10 * np.log10(Pxx))\n",
    "        plt.plot(pulse_rate_freq, 10 * np.log10(PRange[MaxInd]),'ro')\n",
    "        plt.xlabel('Frequency (Hz)')\n",
    "        plt.ylabel('Power (dB)')\n",
    "        plt.xlim([0, 4.5])\n",
    "        plt.title('Power Spectrum and Peak Frequency')\n",
    "            \n",
    "    return pulse_rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import pickle\n",
    "\n",
    "# with open(\"/home/pradyumnachari/Documents/ImplicitPPG/SIGGRAPH_Data/fitzpatrick_labels.pkl\", \"rb\") as fpf:\n",
    "#         out = pickle.load(fpf)\n",
    "        \n",
    "# # print(out[0]['v_1'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# print(out[2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# aa = np.load('/home/pradyumnachari/Documents/rPPG_Implicit_Expt/vitals/dataset/subject_1_1/ppg/ppg.npy')\n",
    "\n",
    "\n",
    "\n",
    "# print(str2array(aa)[15])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 213856,
     "status": "ok",
     "timestamp": 1672890126548,
     "user": {
      "displayName": "Pradyumna Chari",
      "userId": "06555723454452599039"
     },
     "user_tz": 480
    },
    "id": "lsfi3Q5OgPgp",
    "outputId": "45760260-c7db-4676-9f60-ef6fe8da2314",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "destination_folder = \"/media/pradyumnachari/OSA_HDD_3/FastImplicitPPG/data\"\n",
    "destination_folder = \"/home/pradyumnachari/Documents/FastImplicitPleth/data\"\n",
    "\n",
    "destination_folder2 = \"/home/pradyumnachari/Documents/FastImplicitPleth/residual_npy/\"\n",
    "\n",
    "# fitz_labels_path = \"/home/pradyumnachari/Documents/ImplicitPPG/SIGGRAPH_Data/fitzpatrick_labels.pkl\"\n",
    "\n",
    "with open(\"/media/pradyumnachari/OSA_HDD_3/FastImplicitPPG/folds/split_0.pkl\", \"rb\") as fpf:\n",
    "        out = pickle.load(fpf)\n",
    "\n",
    "\n",
    "train = out[\"train\"]\n",
    "val = out[\"val\"]\n",
    "test = out[\"test\"]\n",
    "\n",
    "print(train)\n",
    "print(test)\n",
    "print(len(train), len(test), len(val))\n",
    "#Dataset\n",
    "dataset = RppgData2(datapath=destination_folder, datapath2=destination_folder2, datapaths=train, recording_str=\"rgb\")\n",
    "dataset_val = RppgData2(datapath=destination_folder, datapath2=destination_folder2, datapaths=val, recording_str=\"rgb\")\n",
    "dataset_test = RppgData2(datapath=destination_folder, datapath2=destination_folder2, datapaths=test, recording_str=\"rgb\")\n",
    "\n",
    "print(len(train), len(test), len(val))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read the npy file: /home/pradyumnachari/Documents/rPPG_Implicit_Expt/vitals/dataset_res_npys/residual_0_subject_1_1.npy\n",
    "aa = np.load('/home/pradyumnachari/Documents/rPPG_Implicit_Expt/vitals/dataset_res_npys/residual_0_subject_1_1.npy')\n",
    "print(aa.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "45IFN6rggeWx"
   },
   "outputs": [],
   "source": [
    "def collate_fn(batch):\n",
    "    try:\n",
    "        print(batch.shape)\n",
    "    except:\n",
    "        print(len(batch))\n",
    "    # Get the number of workers\n",
    "    num_workers = len(batch[0])\n",
    "    # Initialize the batch\n",
    "    batch_size = len(batch)\n",
    "    batched_data = []\n",
    "    \n",
    "    # Iterate over each worker\n",
    "    for worker_idx in range(num_workers):\n",
    "        # Initialize the worker's data\n",
    "        worker_data = []\n",
    "        \n",
    "        # Iterate over each sample in the batch\n",
    "        for sample_idx in range(batch_size):\n",
    "            # Get the sample\n",
    "            sample = batch[sample_idx][worker_idx]\n",
    "            \n",
    "            # Add the sample to the worker's data\n",
    "            worker_data.append(sample)\n",
    "        \n",
    "        # Convert the worker's data to a numpy array and add it to the batch\n",
    "        batched_data.append(np.array(worker_data))\n",
    "    \n",
    "    # Return the batched data\n",
    "    return tuple(batched_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train_dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=8, shuffle=True, num_workers = 2)\n",
    "# test_dataloader  = DataLoader(dataset_test, collate_fn=collate_fn, batch_size=8, shuffle=True, num_workers = 2)\n",
    "\n",
    "train_dataloader = DataLoader(dataset, batch_size=4, shuffle=True)\n",
    "test_dataloader  = DataLoader(dataset_test, batch_size=4, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OL8TYkWjg3DB"
   },
   "outputs": [],
   "source": [
    "# print(dataset.vital_mean, dataset.vital_std)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "executionInfo": {
     "elapsed": 652,
     "status": "ok",
     "timestamp": 1672890127356,
     "user": {
      "displayName": "Pradyumna Chari",
      "userId": "06555723454452599039"
     },
     "user_tz": 480
    },
    "id": "bxHL640Jg6U1",
    "outputId": "9ac59d15-d521-433c-e713-d571a0369987"
   },
   "outputs": [],
   "source": [
    "print(len(train_dataloader), len(test_dataloader))\n",
    "\n",
    "# Visualize some examples\n",
    "for i in range(len(dataset)):\n",
    "    train_batch, train_batch_sig = dataset[i]\n",
    "    break\n",
    "\n",
    "print(train_batch.shape, train_batch.dtype, train_batch_sig.shape)\n",
    "\n",
    "# plt.figure()\n",
    "# plt.plot(np.transpose(train_batch, (1,0)))\n",
    "# plt.figure()\n",
    "# plt.plot(train_batch_sig)\n",
    "\n",
    "for i in range(len(dataset_test)):\n",
    "    test_batch, test_batch_sig = dataset_test[i]\n",
    "    break\n",
    "\n",
    "print(test_batch.shape, test_batch_sig.shape)\n",
    "\n",
    "# plt.figure()\n",
    "# plt.plot(np.transpose(test_batch, (1,0)))\n",
    "# plt.figure()\n",
    "# plt.plot(test_batch_sig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 9,
     "status": "ok",
     "timestamp": 1672890127476,
     "user": {
      "displayName": "Pradyumna Chari",
      "userId": "06555723454452599039"
     },
     "user_tz": 480
    },
    "id": "UnUIKN4BlN60",
    "outputId": "be258801-3b6a-4d00-b761-e64b7687eee5"
   },
   "outputs": [],
   "source": [
    "#Set Checkpoint Directory\n",
    "colab_filename = 'VITAL_best_Fold1_May20_FinalSplit'\n",
    "# ckpt_parent_path = '/media/pradyumnachari/FastImplicitP/ImplicitPPG/checkpoints_01-11-expts'\n",
    "ckpt_path = '/media/pradyumnachari/OSA_HDD_3/FastImplicitPPG/checkpoints-02-16'\n",
    "# assert os.path.exists(ckpt_parent_path), \"Check folder to save checkpoint\"\n",
    "assert os.path.exists(ckpt_path), \"Check folder to save checkpoint\"\n",
    "\n",
    "# ckpt_path = os.path.join(ckpt_parent_path, colab_filename)\n",
    "os.makedirs(ckpt_path, exist_ok=True)\n",
    "print(f\"Checkpoints will be saved in {ckpt_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 8,
     "status": "ok",
     "timestamp": 1672890127477,
     "user": {
      "displayName": "Pradyumna Chari",
      "userId": "06555723454452599039"
     },
     "user_tz": 480
    },
    "id": "9l0ED1COWv63",
    "outputId": "d66cb8ac-994b-4429-bdca-1fd0f8021fd6"
   },
   "outputs": [],
   "source": [
    "#Check if Checkpoints exist\n",
    "paths = os.listdir(ckpt_path)\n",
    "if(len(paths) > 0):\n",
    "    paths.sort() \n",
    "    last_checkpoint = os.path.join(ckpt_path, paths[-1])\n",
    "\n",
    "    print(\"Checkpoints already exist-\")\n",
    "    print(paths)\n",
    "    print(\"Last checkpoint set to -\")\n",
    "    print(last_checkpoint)\n",
    "else:\n",
    "  print(\"No checkpoints found, starting from scratch!\")\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_BxlO_PRUP23"
   },
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "class PlethRegressor(nn.Module):\n",
    "    def __init__(self, inp_len, out_len, latent=512):\n",
    "        super(PlethRegressor, self).__init__()\n",
    "        self.inp_len = inp_len\n",
    "        self.latent = latent\n",
    "        self.out_len = out_len\n",
    "        self.Enc1 = nn.Sequential(\n",
    "            nn.Conv1d(6, 64, kernel_size=9, stride=1,padding=4),\n",
    "            nn.BatchNorm1d(64),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        self.Enc2 = nn.Sequential(\n",
    "            nn.Conv1d(64, 128, kernel_size=9, stride=1,padding=4),\n",
    "            nn.BatchNorm1d(128),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        self.Enc3 = nn.Sequential(\n",
    "            nn.Conv1d(128, 128, kernel_size=9, stride=1,padding=4),\n",
    "            nn.BatchNorm1d(128),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        self.Enc4 = nn.Sequential(\n",
    "            nn.Conv1d(128, 128, kernel_size=9, stride=1,padding=4),\n",
    "            nn.BatchNorm1d(128),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        \n",
    "        self.Dec1 = nn.Sequential(\n",
    "            nn.Conv1d(128, 128, kernel_size=9, stride=1,padding=4),\n",
    "            nn.BatchNorm1d(128),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        self.Dec2 = nn.Sequential(\n",
    "            nn.Conv1d(128, 128, kernel_size=9, stride=1,padding=4),\n",
    "            nn.BatchNorm1d(128),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        self.Dec3 = nn.Sequential(\n",
    "            nn.Conv1d(128, 64, kernel_size=9, stride=1,padding=4),\n",
    "            nn.BatchNorm1d(64),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        self.Dec4 = nn.Sequential(\n",
    "            nn.Conv1d(64, 16, kernel_size=9, stride=1,padding=4),\n",
    "            nn.BatchNorm1d(16),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        self.FinalLayer   = nn.Conv1d(16, 1, kernel_size=1, stride=1)\n",
    "    def forward(self, inp, wts):\n",
    "        #wts has shape B,128,128\n",
    "        wts = wts.reshape(wts.shape[0],1,wts.shape[1],wts.shape[2],1)\n",
    "        #inp has shape [B,300,128,128,3]\n",
    "        inp = torch.sum(inp*wts,(2,3))/torch.sum(wts,(2,3))\n",
    "        inp = inp.permute(0,2,1)\n",
    "        \n",
    "        ##Mean and std normalization\n",
    "        inp = (inp - torch.mean(inp, 2, True))/torch.std(inp, dim=2, keepdim=True)\n",
    "        # inp = (inp - torch.mean(inp, (1,2), True))/torch.std(inp, dim=(1,2), keepdim=True)\n",
    "\n",
    "        inp = self.Enc1(inp)\n",
    "        inp = self.Enc2(inp)\n",
    "        inp = self.Enc3(inp)\n",
    "        inp = self.Enc4(inp)\n",
    "        inp = self.Dec1(inp)\n",
    "        inp = self.Dec2(inp)\n",
    "        inp = self.Dec3(inp)\n",
    "        inp = self.Dec4(inp)\n",
    "        output_signal = self.FinalLayer(inp)\n",
    "        return torch.squeeze(output_signal, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 1354,
     "status": "ok",
     "timestamp": 1672968088467,
     "user": {
      "displayName": "Pradyumna Chari",
      "userId": "06555723454452599039"
     },
     "user_tz": 480
    },
    "id": "YN85G93rqO4M",
    "outputId": "82622334-27cc-47e5-bf8a-d6794dca6cfe"
   },
   "outputs": [],
   "source": [
    "# model1 = PlethRegressor(inp_len=300, out_len=300)\n",
    "\n",
    "# out = model1(torch.randn(5,300,128,128,3),torch.randn(5,128,128))\n",
    "# print(out.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5GchXcSsU8q2"
   },
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "class CNN3D(nn.Module):\n",
    "    def __init__(self, frames=64, sidelen = 128, channels=3):  \n",
    "        super(CNN3D, self).__init__()\n",
    "        \n",
    "        self.ConvBlock1 = nn.Sequential(\n",
    "            nn.Conv3d(channels, 16, [1,5,5],stride=1, padding=[0,2,2]),\n",
    "            nn.BatchNorm3d(16),\n",
    "            nn.ReLU(inplace=True),\n",
    "        )\n",
    "\n",
    "        self.ConvBlock2 = nn.Sequential(\n",
    "            nn.Conv3d(16, 32, [3, 3, 3], stride=1, padding=1),\n",
    "            nn.BatchNorm3d(32),\n",
    "            nn.ReLU(inplace=True),\n",
    "        )\n",
    "        self.ConvBlock3 = nn.Sequential(\n",
    "            nn.Conv3d(32, 64, [3, 3, 3], stride=1, padding=1),\n",
    "            nn.BatchNorm3d(64),\n",
    "            nn.ReLU(inplace=True),\n",
    "        )\n",
    "        \n",
    "        self.ConvBlock4 = nn.Sequential(\n",
    "            nn.Conv3d(64, 64, [3, 3, 3], stride=1, padding=1),\n",
    "            nn.BatchNorm3d(64),\n",
    "            nn.ReLU(inplace=True),\n",
    "        )\n",
    "        # self.ConvBlock5 = nn.Sequential(\n",
    "        #     nn.Conv3d(64, 64, [3, 3, 3], stride=1, padding=1),\n",
    "        #     nn.BatchNorm3d(64),\n",
    "        #     nn.ReLU(inplace=True),\n",
    "        # )\n",
    "        # self.ConvBlock6 = nn.Sequential(\n",
    "        #     nn.Conv3d(64, 64, [3, 3, 3], stride=1, padding=1),\n",
    "        #     nn.BatchNorm3d(64),\n",
    "        #     nn.ReLU(inplace=True),\n",
    "        # )\n",
    "        # self.ConvBlock7 = nn.Sequential(\n",
    "        #     nn.Conv3d(64, 64, [3, 3, 3], stride=1, padding=1),\n",
    "        #     nn.BatchNorm3d(64),\n",
    "        #     nn.ReLU(inplace=True),\n",
    "        # )\n",
    "        # self.ConvBlock8 = nn.Sequential(\n",
    "        #     nn.Conv3d(64, 64, [3, 3, 3], stride=1, padding=1),\n",
    "        #     nn.BatchNorm3d(64),\n",
    "        #     nn.ReLU(inplace=True),\n",
    "        # )\n",
    "        # self.ConvBlock9 = nn.Sequential(\n",
    "        #     nn.Conv3d(64, 64, [3, 3, 3], stride=1, padding=1),\n",
    "        #     nn.BatchNorm3d(64),\n",
    "        #     nn.ReLU(inplace=True),\n",
    "        # )\n",
    "        \n",
    "        # self.upsample = nn.Sequential(\n",
    "        #     nn.ConvTranspose3d(in_channels=64, out_channels=64, kernel_size=[4,1,1], stride=[2,1,1], padding=[1,0,0]),   #[1, 128, 32]\n",
    "        #     nn.BatchNorm3d(64),\n",
    "        #     nn.ELU(),\n",
    "        # )\n",
    "        # self.upsample2 = nn.Sequential(\n",
    "        #     nn.ConvTranspose3d(in_channels=64, out_channels=64, kernel_size=[4,1,1], stride=[2,1,1], padding=[1,0,0]),   #[1, 128, 32]\n",
    "        #     nn.BatchNorm3d(64),\n",
    "        #     nn.ELU(),\n",
    "        # )\n",
    " \n",
    "        self.ConvBlock10 = nn.Conv3d(64, 1, [1,1,1],stride=1, padding=0)\n",
    "        \n",
    "        # self.MaxpoolSpa = nn.MaxPool3d((1, 2, 2), stride=(1, 2, 2))\n",
    "        # self.MaxpoolSpaTem = nn.MaxPool3d((2, 2, 2), stride=2)\n",
    "        self.MaxpoolTem = nn.MaxPool3d((2, 1, 1), stride=(2, 1, 1))\n",
    "        \n",
    "        \n",
    "        #self.poolspa = nn.AdaptiveMaxPool3d((frames,1,1))    # pool only spatial space \n",
    "        self.poolspa = nn.AdaptiveAvgPool3d((1,sidelen,sidelen))\n",
    "\n",
    "        \n",
    "    def forward(self, x):\t    \t# x [B, T, 128,128, 3]\n",
    "        # x_visual = x\n",
    "\n",
    "        x = x.permute(0,4,1,2,3)\n",
    "#         print(x.shape)\n",
    "        [batch,channel,length,width,height] = x.shape\n",
    "        # print(length)\n",
    "        # print(torch.mean(x, dim=(1,2,3,4)))\n",
    "          \n",
    "        x = self.ConvBlock1(x)\t\t     # x [3, T, 128,128]\n",
    "        # x = self.MaxpoolSpa(x)       # x [16, T, 64,64]\n",
    "\n",
    "        x = self.ConvBlock2(x)\t\t    # x [32, T, 64,64]\n",
    "        x_visual6464 = self.ConvBlock3(x)\t    \t# x [32, T, 64,64]\n",
    "        x = self.MaxpoolTem(x_visual6464)      # x [32, T/2, 128,128]    Temporal halve\n",
    "\n",
    "        x = self.ConvBlock4(x)\t\t    # x [64, T/2, 128,128]\n",
    "        ###MIGHT NEED THE BELOW\n",
    "        # x_visual3232 = self.ConvBlock5(x)\t    \t# x [64, T/2, 128,128]\n",
    "        # x = self.MaxpoolTem(x_visual3232)      # x [64, T/4, 128,128]\n",
    "\n",
    "\n",
    "        # x = self.ConvBlock6(x)\t\t    # x [64, T/4, 128,128]\n",
    "        # x_visual1616 = self.ConvBlock7(x)\t    \t# x [64, T/4, 128,128]\n",
    "        ###END NEED HERE\n",
    "\n",
    "\n",
    "\n",
    "        # x = self.MaxpoolSpa(x_visual1616)      # x [64, T/4, 8,8]\n",
    "\n",
    "\n",
    "        # x = self.ConvBlock8(x)\t\t    # x [64, T/4, 8, 8]\n",
    "        # x = self.ConvBlock9(x)\t\t    # x [64, T/4, 8, 8]\n",
    "        # x = self.upsample(x)\t\t    # x [64, T/2, 8, 8]\n",
    "        # x = self.upsample2(x)\t\t    # x [64, T, 8, 8]\n",
    "\n",
    "        # print(x.shape)\n",
    "        \n",
    "        x = self.poolspa(x)     # x [64, T, 1,1]    -->  groundtruth left and right - 7 \n",
    "        x = torch.sigmoid(self.ConvBlock10(x))    # x [1, T, 1,1]\n",
    "\n",
    "        # print(x.shape)\n",
    "        # print(torch.mean(x, dim=(1,2,3,4)))\n",
    "         \n",
    "\n",
    "        return x.reshape(x.shape[0],x.shape[3],x.shape[4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 27660,
     "status": "ok",
     "timestamp": 1672970756948,
     "user": {
      "displayName": "Pradyumna Chari",
      "userId": "06555723454452599039"
     },
     "user_tz": 480
    },
    "id": "r5TtmD00q4a0",
    "outputId": "92316de5-1f6f-4f4f-9cb8-19b1c6d228a9"
   },
   "outputs": [],
   "source": [
    "# model2 = CNN3D(frames=64, sidelen = 128, channels=3)\n",
    "\n",
    "# out = model2(torch.randn(5,64,128,128,3))\n",
    "# print(out.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FDyYvb2LMzhA"
   },
   "outputs": [],
   "source": [
    "import math\n",
    "import torch.nn as nn\n",
    "from torch.nn.modules.utils import _triple\n",
    "import torch as tr\n",
    "import pdb\n",
    "\n",
    "class SNRLoss_dB_Signals(nn.Module):\n",
    "  def __init__(self):\n",
    "    super(SNRLoss_dB_Signals, self).__init__()\n",
    "  def forward(self, outputs: tr.Tensor, targets: tr.Tensor, Fs=30):\n",
    "    device = outputs.device\n",
    "    if not outputs.is_cuda:\n",
    "      torch.backends.mkl.is_available()\n",
    "    N = 600#3*outputs.shape[-1] #1000#\n",
    "    # N_samp = outputs.shape[-1]\n",
    "    N_samp = 300\n",
    "    #print(N)\n",
    "    pulse_band = tr.tensor([45/60., 180/60.], dtype=tr.float32).to(device)\n",
    "    # wind_sz = int(1*N/64)#20\n",
    "    wind_sz = 3\n",
    "\n",
    "    f = tr.linspace(0, Fs/2, int(N/2)+1, dtype=tr.float32).to(device)\n",
    "    # print(60*(f[1]-f[0]))\n",
    "    min_idx = tr.argmin(tr.abs(f - pulse_band[0]))\n",
    "    max_idx = tr.argmin(tr.abs(f - pulse_band[1]))\n",
    "    # print(min_idx,max_idx)\n",
    "\n",
    "    outputs = outputs.view(-1, N_samp)\n",
    "    targets = targets.view(-1, N_samp)\n",
    "\n",
    "    #Generate GT heart indices from GT signals\n",
    "    Y = torch.fft.rfft(targets, n=N, dim=1, norm='forward')\n",
    "    Y2 = tr.abs(Y) ** 2\n",
    "    HRixs = tr.argmax(Y2[:,min_idx:max_idx],axis=1)+min_idx\n",
    "\n",
    "    #print(outputs.shape)\n",
    "    X = torch.fft.rfft(outputs, n=N, dim=1, norm='forward')\n",
    "\n",
    "    P1 = tr.abs(X) ** 2\n",
    "    # print(HRixs)\n",
    "    # plt.figure()\n",
    "    # plt.plot(f,P1[0])\n",
    "    # plt.plot(f[HRixs[0]],P1[0,HRixs[0]],'rx')\n",
    "    # plt.plot(f[HRixs[0]-wind_sz],P1[0,HRixs[0]-wind_sz],'gx')\n",
    "    # plt.plot(f[HRixs[0]+wind_sz],P1[0,HRixs[0]+wind_sz],'gx')\n",
    "    # plt.grid('on')\n",
    "    # plt.show()\n",
    "\n",
    "    # calc SNR for each batch\n",
    "    losses = tr.empty((X.shape[0],), dtype=tr.float32)#.to(device)\n",
    "    for count, ref_idx in enumerate(HRixs):\n",
    "      pulse_freq_amp = tr.sum(P1[count, ref_idx-wind_sz:ref_idx+wind_sz])+tr.sum(P1[count, 2*ref_idx-wind_sz:2*ref_idx+wind_sz])\n",
    "      # other_avrg = (tr.sum(P1[count, min_idx:ref_idx-wind_sz]) + tr.sum(P1[count, ref_idx+wind_sz:max_idx]))\n",
    "      other_avrg = (tr.sum(P1[count, min_idx:ref_idx-wind_sz])+tr.sum(P1[count, ref_idx+wind_sz:2*ref_idx-wind_sz]) + tr.sum(P1[count, 2*ref_idx+wind_sz:max_idx]))\n",
    "      # print(pulse_freq_amp,other_avrg)\n",
    "      losses[count] = -10*tr.log10(pulse_freq_amp/(other_avrg+1e-7))\n",
    "    losses.to(device)\n",
    "    return tr.mean(losses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def total_variation_loss(img):\n",
    "     bs_img, h_img, w_img = img.size()\n",
    "     tv_h = torch.pow(img[:,1:,:]-img[:,:-1,:], 2).sum()\n",
    "     tv_w = torch.pow(img[:,:,1:]-img[:,:,:-1], 2).sum()\n",
    "     return (tv_h+tv_w)/(bs_img*h_img*w_img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(os.path.join(os.getcwd(), f\"{ckpt_path}/latest_context.pth\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "executionInfo": {
     "elapsed": 795162,
     "status": "error",
     "timestamp": 1672890922635,
     "user": {
      "displayName": "Pradyumna Chari",
      "userId": "06555723454452599039"
     },
     "user_tz": 480
    },
    "id": "8_ri0LHHhdW0",
    "outputId": "cd82f79d-1f57-4232-c61a-1c46e21ff560",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "import imageio as iio\n",
    "\n",
    "#CONTEXT PATH\n",
    "PATH = os.path.join(os.getcwd(), f\"{ckpt_path}/latest_context.pth\")\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = PlethRegressor(inp_len=300, out_len=300).to(device)\n",
    "model2 = CNN3D(frames=64, sidelen = 128, channels=6).to(device)\n",
    "\n",
    "# model.load_state_dict(torch.load(last_checkpoint))\n",
    "loss_fn  = Neg_Pearson2()\n",
    "loss_fn2  = SNRLoss_dB_Signals()#Neg_Pearson()\n",
    "\n",
    "# lam = 3 # Neg_Pearson2\n",
    "# lam2 = 4\n",
    "\n",
    "learning_rate = 2e-4\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-6)\n",
    "optimizer2 = torch.optim.Adam(model2.parameters(), lr=learning_rate, weight_decay=1e-6)\n",
    "\n",
    "# Train configurations\n",
    "epochs = 3\n",
    "checkpoint_period = 1\n",
    "epoch_start = 1\n",
    "mae_best_loss = 1000\n",
    "# if os.path.exists(PATH):\n",
    "#     print('Context checkpoint exists. Loading state dictionaries.')\n",
    "#     checkpoint = torch.load(PATH)\n",
    "#     model.load_state_dict(checkpoint['model_state_dict'])\n",
    "#     optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "#     epoch_start = checkpoint['epoch']\n",
    "#     epoch_start+=1\n",
    "\n",
    "loss_list = []\n",
    "\n",
    "# train_dataloader = train_dataloader[:100]\n",
    "\n",
    "for epoch in range(epoch_start, epochs+1):\n",
    "    # Training Phase\n",
    "    loss_train = 0\n",
    "    no_batches = 0\n",
    "    print(\"Starting Epoch: {}\".format(epoch))\n",
    "    for batch, (imgs, signal) in tqdm(enumerate(train_dataloader), total=len(train_dataloader)):\n",
    "        model.train()\n",
    "        model2.train()\n",
    "        # Convert to the appropriate format and mount on the specified device\n",
    "        imgs = imgs.float().to(device)\n",
    "        signal = signal.float().to(device)\n",
    "        # convert to float and then put it on device\n",
    "\n",
    "        #Get attention mask\n",
    "        msk = model2(imgs[:,0:64])\n",
    "\n",
    "        # Predict the PPG signal and find ther loss\n",
    "        pred_signal = model(imgs,msk)\n",
    "        \n",
    "        # print(pred_signal)\n",
    "        # print(signal)\n",
    "        loss1 = loss_fn(pred_signal, signal)\n",
    "        loss2 = loss_fn2(pred_signal, signal)\n",
    "        loss3 = total_variation_loss(msk)\n",
    "\n",
    "        # loss = 3*loss1 + 5*loss3\n",
    "        # if epoch == 0 and batch > len(train_dataloader)//2:\n",
    "        #     loss = loss2 + 3*loss1 + 5*loss3\n",
    "\n",
    "\n",
    "\n",
    "        # if its the first epoch and \n",
    "        # if epoch == 0 and batch < len(train_dataloader)//2:\n",
    "        #     loss = 3*loss1 + 5*loss3\n",
    "        # else:\n",
    "        #     loss = loss2 + 3*loss1 + 5*loss3\n",
    "        \n",
    "#         loss = loss2+lam2*loss3\n",
    "        \n",
    "#         print(loss,loss1,loss2,loss3)\n",
    "        loss = loss2+3*loss1+5*loss3\n",
    "        \n",
    "\n",
    "        # Backprop\n",
    "        optimizer.zero_grad()\n",
    "        optimizer2.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer2.step()\n",
    "        \n",
    "        # Accumulate the total loss\n",
    "        loss_train += loss.item()\n",
    "        no_batches+=1\n",
    "\n",
    "        loss_list.append(loss.item())\n",
    "\n",
    "        # show mask\n",
    "        if batch%10 == 0:\n",
    "            ixx = 0\n",
    "            og_msk = msk[ixx].cpu().detach().numpy()\n",
    "            # plt.imshow(og_msk)\n",
    "            # save to file: \"Mask_samples/epoch_{}_batch_{}.png\".format(epoch, batch)\n",
    "            # plt.imsave(\"./Mask_samples/epoch_{}_batch_{}_loss{}.png\".format(epoch, batch, round(loss.item(), 4)), og_msk)\n",
    "            iio.imwrite(\"./Mask_samples/epoch_{}_batch_{}_loss{}.png\".format(epoch, batch, round(loss.item(), 4)), og_msk)\n",
    "            # plt.show()\n",
    "            # plot the targets\n",
    "            # plt.figure()\n",
    "            # # evaluate the model\n",
    "            # plt.plot(signal[ixx].cpu().detach().numpy(), label='GT')\n",
    "            # plt.plot(pred_signal[ixx].cpu().detach().numpy(), label='Pred')\n",
    "            # plt.legend()\n",
    "            # plt.show()\n",
    "\n",
    "            # dont plot the signal, but save it to file\n",
    "            np.save(\"./Mask_samples/pred_epoch_{}_batch_{}_loss{}.npy\".format(epoch, batch, round(loss.item(), 4)), pred_signal[ixx].cpu().detach().numpy())\n",
    "            np.save(\"./Mask_samples/gt_epoch_{}_batch_{}_loss{}.npy\".format(epoch, batch, round(loss.item(), 4)), signal[ixx].cpu().detach().numpy())\n",
    "\n",
    "    # Save the model every few epochs\n",
    "    if(epoch % checkpoint_period == 0):\n",
    "        torch.save(model.state_dict(), os.path.join(os.getcwd(), f\"{ckpt_path}/PhysNet_state_dict_{epoch}_epochs.pth\"))\n",
    "        torch.save(model2.state_dict(), os.path.join(os.getcwd(), f\"{ckpt_path}/PhysNetAtt_state_dict_{epoch}_epochs.pth\"))\n",
    "        #See if best checkpoint\n",
    "        mae_loss_list, hrs = eval_model(root_dir=destination_folder, root_dir2=destination_folder2, session_names=val, model=model, model2=model2)\n",
    "        current_loss = np.mean(mae_loss_list) \n",
    "        if(current_loss < mae_best_loss):\n",
    "            mae_best_loss = current_loss\n",
    "            torch.save(model.state_dict(), os.path.join(os.getcwd(), f\"{ckpt_path}/PhysNet_state_dict_best.pth\"))\n",
    "            torch.save(model2.state_dict(), os.path.join(os.getcwd(), f\"{ckpt_path}/PhysNetAtt_state_dict_best.pth\"))\n",
    "            print(\"Best checkpoint saved!\")\n",
    "        print(\"Saved Checkpoint!\")\n",
    "\n",
    "    print(f\"Epoch: {epoch} ; Loss: {loss_train/no_batches:>7f}\")\n",
    "    # plot loss list so far\n",
    "    plt.figure()\n",
    "    plt.plot(loss_list)\n",
    "    plt.show()\n",
    "    \n",
    "    #plot an example frame\n",
    "    ixx = 0\n",
    "    og_msk = msk[ixx].cpu().detach().numpy()\n",
    "    \n",
    "    plt.figure()\n",
    "    plt.imshow(og_msk)\n",
    "    plt.show()\n",
    "    \n",
    "    #SAVE CONTEXT AFTER EPOCH\n",
    "    torch.save({\n",
    "            'epoch': epoch,\n",
    "            'model_state_dict': model.state_dict(),\n",
    "            'model2_state_dict': model2.state_dict(),\n",
    "            'optimizer_state_dict': optimizer.state_dict(),\n",
    "            }, PATH)\n",
    "\n",
    "        \n",
    "# plot the loss\n",
    "plt.figure()\n",
    "plt.plot(loss_list)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.plot(loss_list)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check what the devide ordinal is\n",
    "print(torch.cuda.current_device())\n",
    "torch.cuda.get_device_name(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check how much memory is available on the cuda gpu\n",
    "print(torch.cuda.memory_allocated(0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# mae_loss_list, hrs = eval_model(root_dir=destination_folder, root_dir2=destination_folder2, session_names=val, model=model, model2=model2)\n",
    "        \n",
    "# # print(hrs[0][50],hrs[1][50])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RkfVgB1puryP"
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from collections import Counter\n",
    "import copy\n",
    "import sklearn.metrics\n",
    "\n",
    "def eval_clinical_performance(hr_est, hr_gt, fitz_labels_path, session_names):\n",
    "    l_m_d_arr = distribute_l_m_d(fitz_labels_path , session_names)\n",
    "    l_m_d_arr = np.array(l_m_d_arr)\n",
    "    #absolute percentage error\n",
    "    # print(hr_gt.shape, hr_est.shape)\n",
    "    apes = np.abs(hr_gt - hr_est)/hr_gt*100\n",
    "    # print(apes)\n",
    "    l_apes = np.reshape(apes[np.where(l_m_d_arr==1)], (-1))\n",
    "    d_apes = np.reshape(apes[np.where(l_m_d_arr==2)], (-1))\n",
    "\n",
    "    l_5 = len(l_apes[l_apes <= 5])/len(l_apes)*100 \n",
    "    d_5 = len(d_apes[d_apes <= 5])/len(d_apes)*100\n",
    "    \n",
    "    l_10 = len(l_apes[l_apes <= 10])/len(l_apes)*100\n",
    "    d_10 = len(d_apes[d_apes <= 10])/len(d_apes)*100\n",
    "\n",
    "    print(\"AAMI Standard - L,D\")\n",
    "    print(l_10, d_10)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "########################################### Performance ##################################################\n",
    "def eval_performance(hr_est, hr_gt):\n",
    "    hr_est = np.reshape(hr_est, (-1))\n",
    "    hr_gt  = np.reshape(hr_gt, (-1))\n",
    "    r = scipy.stats.pearsonr(hr_est, hr_gt)\n",
    "    mae = np.sum(np.abs(hr_est - hr_gt))/len(hr_est)\n",
    "    hr_std = np.std(hr_est - hr_gt)\n",
    "    hr_rmse = np.sqrt(np.sum(np.square(hr_est-hr_gt))/len(hr_est))\n",
    "    hr_mape = sklearn.metrics.mean_absolute_percentage_error(hr_est, hr_gt)\n",
    "\n",
    "    return mae, hr_mape, hr_rmse, hr_std, r[0]\n",
    "\n",
    "def eval_performance_bias(hr_est, hr_gt, fitz_labels_path, session_names):\n",
    "    l_m_d_arr = distribute_l_m_d(fitz_labels_path , session_names)\n",
    "    l_m_d_arr = np.array(l_m_d_arr)\n",
    "\n",
    "    general_performance = eval_performance(hr_est, hr_gt)\n",
    "    l_p = np.array(eval_performance(hr_est[np.where(l_m_d_arr == 1)], hr_gt[np.where(l_m_d_arr == 1)]))\n",
    "    d_p = np.array(eval_performance(hr_est[np.where(l_m_d_arr == 2)], hr_gt[np.where(l_m_d_arr == 2)]))\n",
    "\n",
    "    performance_diffs = np.array([l_p-d_p])\n",
    "    performance_diffs = np.abs(performance_diffs)\n",
    "    performance_max_diffs = performance_diffs.max(axis=0)\n",
    "\n",
    "    print(\"General Performance\")\n",
    "    print(general_performance)\n",
    "    print(\"Performance Max Differences\")\n",
    "    print(performance_max_diffs)\n",
    "    print(\"Performance By Skin Tone\")\n",
    "    print(\"Light - \", l_p)\n",
    "    print(\"Dark - \", d_p)\n",
    "\n",
    "    return general_performance, performance_max_diffs\n",
    "\n",
    "\n",
    "\n",
    "def eval_models_performance(list_of_errors, list_of_labels):\n",
    "    err_arr = np.array(list_of_errors)\n",
    "    err_str_arr = [['%0.3f' %x for x in y] for y in err_arr]\n",
    "\n",
    "    fig, ax = plt.subplots()\n",
    "    fig.patch.set_visible(False)\n",
    "    fig.set_size_inches(fig_size[0], fig_size[1])\n",
    "    ax.axis('off')\n",
    "    ax.axis('tight')\n",
    "    table = ax.table(cellText=err_str_arr, \n",
    "                    colLabels=[\"MAE\", \"RMSE\", \"STD\", \"MAPE\", \"Pearson Correlation\"], \n",
    "                    rowLabels=list_of_labels)\n",
    "    fig.tight_layout()\n",
    "    ax.set_title(f'Performance Metrics', weight='bold', size=14, color='k')\n",
    "    table.scale(1, 6)\n",
    "    table.set_fontsize(20)\n",
    "    plt.show()\n",
    "\n",
    "from scipy.stats import wasserstein_distance\n",
    "def eval_wasserstein(errors, fitz_labels_path, session_names):\n",
    "    l_m_d_arr = distribute_l_m_d(fitz_labels_path , session_names)\n",
    "    l_m_d_arr = np.array(l_m_d_arr)\n",
    "\n",
    "    l_p = np.array(errors[np.where(l_m_d_arr == 1)[0]])\n",
    "    d_p = np.array(errors[np.where(l_m_d_arr == 2)[0]])\n",
    "\n",
    "    wass_distances = wasserstein_distance(l_p, d_p)\n",
    "    print(\"Wasserstein L/D\")\n",
    "    print(wass_distances)\n",
    "    print(\"Num Light:\", len(l_p), \"Num Dark\", len(d_p))\n",
    "\n",
    "    return wass_distances\n",
    "\n",
    "\n",
    "\n",
    "############################################# Fairness ###################################################\n",
    "\n",
    "def gen_cdf(data, cutoff_mae = 2.5):\n",
    "    # sort the data:\n",
    "    data_sorted = np.sort(data)\n",
    "    # calculate the proportional values of samples\n",
    "    p = (np.arange(len(data))+0.5)/len(data)\n",
    "\n",
    "    return data_sorted, p\n",
    "\n",
    "def gen_cdf_sp(data, freq_res=0.1, max_error = 100):\n",
    "    #round to nearest freq_resolution\n",
    "    data = np.round(data/freq_res)*freq_res\n",
    "    data = np.sort(data)\n",
    "    #count num in each freq_resolution\n",
    "    data_counts = list(Counter(data).items())\n",
    "    freqs = [i[0] for i in data_counts]\n",
    "    counts = [i[1] for i in data_counts]\n",
    "    counts_cum = np.cumsum(counts)\n",
    "    sum_counts = np.sum(counts)\n",
    "    cdf = [i / sum_counts for i in counts_cum]\n",
    "    #add missing freqs\n",
    "    i = 0\n",
    "    #account for not having zero bin\n",
    "    if(freqs[0] != 0):\n",
    "        freqs.insert(0,0.00)\n",
    "        cdf.insert(0,0.00)\n",
    "    while len(freqs) != int(100/freq_res):\n",
    "        if(len(freqs) > i):\n",
    "            if(freqs[i] == freq_res*i):\n",
    "                i += 1\n",
    "            else:\n",
    "                freqs.insert(i, freq_res*i)\n",
    "                cdf.insert(i, cdf[i-1])\n",
    "                # print(freqs[i], freq_res*i,cdf[i])\n",
    "                i +=1\n",
    "        else:\n",
    "            # print(freqs[i], freq_res*i,cdf[i])\n",
    "            freqs.insert(i, freq_res*i)\n",
    "            cdf.insert(i, cdf[i-1])\n",
    "            # print(freqs[i], freq_res*i,cdf[i])\n",
    "            i +=1\n",
    "\n",
    "\n",
    "    return freqs, cdf\n",
    "\n",
    "def calc_sp_ks_test(cdf1, cdf2):\n",
    "    cdf1 = np.array(cdf1)\n",
    "    cdf2 = np.array(cdf2)\n",
    "    max = np.max(np.abs(cdf1-cdf2))\n",
    "    return max\n",
    "\n",
    "def calc_sp_diff_expected(cdf1, cdf2, freqs):\n",
    "    cdf1 = copy.deepcopy(cdf1)\n",
    "    cdf2 = copy.deepcopy(cdf2)\n",
    "    cdf1.insert(0,0)\n",
    "    cdf2.insert(0,0)\n",
    "    cdf1 = np.array(cdf1)\n",
    "    cdf2 = np.array(cdf2)\n",
    "    pdf1 = np.diff(cdf1)\n",
    "    pdf2 = np.diff(cdf2)\n",
    "\n",
    "    diff_pdf = pdf1-pdf2\n",
    "    group_fairness_expectation = np.sum(diff_pdf * freqs)\n",
    "\n",
    "    return group_fairness_expectation\n",
    "\n",
    "def distribute_l_m_d(fitz_labels_path, session_names):\n",
    "    with open(fitz_labels_path, \"rb\") as fpf:\n",
    "        out = pickle.load(fpf)\n",
    "\n",
    "    #mae_list\n",
    "    #session_names\n",
    "    sess_w_fitz = []\n",
    "    fitz_dict = dict(out)\n",
    "    l_m_d_arr = []\n",
    "    for i, sess in enumerate(session_names):\n",
    "        pid = sess.split(\"_\")\n",
    "        pid = pid[0] + \"_\" + pid[1]\n",
    "        fitz_id = fitz_dict[pid]\n",
    "        if(fitz_id < 3):\n",
    "            l_m_d_arr.append(1)\n",
    "        elif(fitz_id < 5): \n",
    "            l_m_d_arr.append(-1)\n",
    "        else:\n",
    "            l_m_d_arr.append(2)\n",
    "    return l_m_d_arr\n",
    "\n",
    "def calc_anderson_k_sample_test(k_samples : list, ):\n",
    "    statistic, crit_vals, p_value = scipy.stats.anderson_ksamp(k_samples, midrank=False)\n",
    "\n",
    "    return statistic, crit_vals, p_value\n",
    "\n",
    "def calc_ks_test(samples : list, samples_labels : list):\n",
    "    #KS Pairwise Sample Test \n",
    "    #return (dark/light),(dark/medium),(medium/light)\n",
    "    ks_full = scipy.stats.ks_2samp(samples[np.where(samples_labels == 2)[0]], samples[np.where(samples_labels == 1)[0]])\n",
    "\n",
    "    return ks_full\n",
    "\n",
    "def dTest(l1,l2):\n",
    "    n1 = len(l1)\n",
    "    n2 = len(l2)\n",
    "    l1 = np.array(l1)\n",
    "    l2 = np.array(l2)\n",
    "    s = np.sqrt(((n1-1)*np.var(l1,ddof=1)+(n2-1)*np.var(l2,ddof=1))/(n1+n2-2))\n",
    "    stat = np.abs(np.mean(l1)-np.mean(l2))/s\n",
    "    print('DTest output is ', stat)\n",
    "    return\n",
    "\n",
    "def calc_d_test(samples : list, samples_labels : list):\n",
    "    #D Pairwise Sample Test \n",
    "    #return (dark/light),(dark/medium),(medium/light)\n",
    "    d_full = dTest(samples[np.where(samples_labels == 2)[0]], samples[np.where(samples_labels == 1)[0]])\n",
    "\n",
    "    return d_full\n",
    "\n",
    "def eval_model_fairness(maes, fitz_labels_path, session_names, name):\n",
    "    l_m_d_arr = distribute_l_m_d(fitz_labels_path , session_names)\n",
    "    l_m_d_arr = np.array(l_m_d_arr)\n",
    "    # maes = maes[:,0]\n",
    "\n",
    "    #Fused\n",
    "    data, p = gen_cdf_sp(maes)\n",
    "    data1, p1 = gen_cdf_sp(maes[np.where(l_m_d_arr == 1)[0]])\n",
    "    data2, p2 = gen_cdf_sp(maes[np.where(l_m_d_arr == 2)[0]])\n",
    "\n",
    "    print(\"L/D, M/D, L/M\")\n",
    "    ####\n",
    "    sp_diff = calc_sp_diff_expected(p2,p1, data)\n",
    "    print(\"Expected Diff\")\n",
    "    print(sp_diff)\n",
    "    #KS Pairwise Sample Test\n",
    "    print(\"KS Test\")\n",
    "    ks_full = calc_ks_test(samples=maes, samples_labels=l_m_d_arr)\n",
    "    print(ks_full)\n",
    "    #D_test\n",
    "    print(\"D_test\")\n",
    "    d_test_results = calc_d_test(maes, l_m_d_arr)\n",
    "    print(d_test_results)\n",
    "    \n",
    "\n",
    "\n",
    "import pickle\n",
    "\n",
    "def eval_model_fitz(fitz_labels_path, mae_list, session_names, name=\"RF\"):\n",
    "    with open(fitz_labels_path, \"rb\") as fpf:\n",
    "        out = pickle.load(fpf)\n",
    "    #mae_list\n",
    "    #session_names\n",
    "    sess_w_fitz = []\n",
    "    fitz_dict = dict(out)\n",
    "    for sess in session_names:\n",
    "        pid = sess.split(\"_\")\n",
    "        pid = pid[0] + \"_\" + pid[1]\n",
    "        sess_w_fitz.append((sess, fitz_dict[pid]))\n",
    "    plt.figure()\n",
    "    plt.hist([i[1] for i in sess_w_fitz], bins=np.arange(1,8))\n",
    "    plt.title(f\"Histogram: {name} Video Count with Fitzpatrick Scale\")\n",
    "    plt.ylabel(\"Num Videos\")\n",
    "    plt.xlabel(\"Fitzpatrick Scale\")\n",
    "\n",
    "    fitz_errors = [[], [], [], [], [], []]\n",
    "    for pid, error in zip(sess_w_fitz, mae_list):\n",
    "        fitz_errors[pid[1]-1].append(error[0])\n",
    "    plt.figure()\n",
    "    plt.plot(np.arange(1,7), [np.mean(a) for a in fitz_errors])\n",
    "    plt.title(f\"{name} Mean Absolute Error as a function of Fitzpatrick Scale\")\n",
    "    plt.ylabel(\"Mean Absolute Error\")\n",
    "    plt.xlabel(\"Fitzpatrick Scale\")\n",
    "    plt.show()\n",
    "\n",
    "    l_m_d_counts = []\n",
    "    for i in sess_w_fitz:\n",
    "        if(i[1]<3):\n",
    "            l_m_d_counts.append(1)\n",
    "        elif(i[1] < 5 and i[1] > 2):\n",
    "            l_m_d_counts.append(2)\n",
    "        else:\n",
    "            l_m_d_counts.append(3)\n",
    "\n",
    "    plt.figure()\n",
    "    plt.hist(l_m_d_counts, bins=np.arange(1,5))\n",
    "    plt.title(f\"Histogram: {name} Video Count with Fitzpatrick Scale\")\n",
    "    plt.ylabel(\"Num Videos\")\n",
    "    plt.xlabel(\"Fitzpatrick Scale\")\n",
    "\n",
    "    cat_errors = [np.mean(a) for a in fitz_errors]\n",
    "    l_m_d_errors = [(cat_errors[0]+cat_errors[1])/2, (cat_errors[2]+cat_errors[3])/2, (cat_errors[4]+cat_errors[5])/2]\n",
    "    plt.figure()\n",
    "    plt.plot(np.arange(1,4), l_m_d_errors)\n",
    "    plt.title(f\"{name} Mean Absolute Error as a function of Light/Medium/Dark Skin Tone Averaged Fitz\")\n",
    "    plt.ylabel(\"Mean Absolute Error\")\n",
    "    plt.xlabel(\"Fitzpatrick Scale\")\n",
    "    plt.show()\n",
    "\n",
    "    #3 histogram plots\n",
    "    plt.figure()\n",
    "    plt.subplot(1, 3, 1)\n",
    "    plt.hist(np.concatenate((fitz_errors[0], fitz_errors[1])), 20)\n",
    "    plt.title(\"Light\")\n",
    "    plt.subplot(1, 3, 2)\n",
    "    plt.xlabel(\"MAE\")\n",
    "    plt.ylabel(\"Counts\")\n",
    "    plt.hist(np.concatenate((fitz_errors[2], fitz_errors[3])), 20)\n",
    "    plt.title(\"Medium\")\n",
    "    plt.subplot(1, 3, 3)\n",
    "    plt.hist(np.concatenate((fitz_errors[4], fitz_errors[5])), 20)\n",
    "    plt.title(\"Dark\")\n",
    "    \n",
    "    l_m_d_errors = [(np.sum(fitz_errors[0]) +np.sum(fitz_errors[1]))/(len(fitz_errors[0])+len(fitz_errors[1])), \n",
    "                    (np.sum(fitz_errors[2]) +np.sum(fitz_errors[3]))/(len(fitz_errors[2])+len(fitz_errors[3])), \n",
    "                    (np.sum(fitz_errors[4]) +np.sum(fitz_errors[5]))/(len(fitz_errors[4])+len(fitz_errors[5]))]\n",
    "    plt.figure()\n",
    "    plt.plot(np.arange(1,4), l_m_d_errors)\n",
    "    plt.title(f\"{name} Mean Absolute Error as a function of Light/Medium/Dark Skin Tone Absolute\")\n",
    "    plt.ylabel(\"Mean Absolute Error\")\n",
    "    plt.xlabel(\"Fitzpatrick Scale\")\n",
    "    plt.show()\n",
    "\n",
    "    print(len(fitz_errors[2])+len(fitz_errors[3]))\n",
    "\n",
    "    print('Mean MAE:', np.mean(np.array(mae_list)))\n",
    "    return [np.mean(a) for a in fitz_errors]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 6,
     "status": "ok",
     "timestamp": 1672890927953,
     "user": {
      "displayName": "Pradyumna Chari",
      "userId": "06555723454452599039"
     },
     "user_tz": 480
    },
    "id": "pDuFzSqVZrGa",
    "outputId": "1f6008f0-9df2-444e-f6e6-00b825c8d47b"
   },
   "outputs": [],
   "source": [
    "model = PlethRegressor(inp_len=300, out_len=300).to(device)\n",
    "model.load_state_dict(torch.load(f\"{ckpt_path}/PhysNet_state_dict_best.pth\"))\n",
    "\n",
    "model2 = CNN3D(frames=64, sidelen = 128, channels=6).to(device)\n",
    "model2.load_state_dict(torch.load(f\"{ckpt_path}/PhysNetAtt_state_dict_best.pth\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 141299,
     "status": "ok",
     "timestamp": 1672891069628,
     "user": {
      "displayName": "Pradyumna Chari",
      "userId": "06555723454452599039"
     },
     "user_tz": 480
    },
    "id": "0UGwflRoZtN4",
    "outputId": "49bbcca8-1d82-4db6-b7f1-1feb33310d04"
   },
   "outputs": [],
   "source": [
    "maes_test, hr_test = eval_model(root_dir=destination_folder, root_dir2=destination_folder2, session_names=train[0:3], model=model, model2=model2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 321,
     "status": "ok",
     "timestamp": 1672891069937,
     "user": {
      "displayName": "Pradyumna Chari",
      "userId": "06555723454452599039"
     },
     "user_tz": 480
    },
    "id": "Zrp5HohBemBk",
    "outputId": "437e7f6a-8753-4943-c6d6-8dd57421b21d"
   },
   "outputs": [],
   "source": [
    "fitz_labels_path = \"/media/pradyumnachari/OSA_HDD_3/FastImplicitPPG/FpST labels.pkl\"\n",
    "eval_clinical_performance(hr_est=np.array(hr_test[0]), hr_gt=np.array(hr_test[1]), fitz_labels_path=fitz_labels_path, session_names=test)\n",
    "print(100*\"*\")\n",
    "eval_performance_bias(hr_est=np.array(hr_test[0]), hr_gt=np.array(hr_test[1]), fitz_labels_path=fitz_labels_path, session_names=test)\n",
    "print(100*\"*\")\n",
    "eval_model_fairness(maes= maes_test[:,0], fitz_labels_path=fitz_labels_path, session_names=test, name=\"rand\")\n",
    "print(100*\"*\")\n",
    "eval_wasserstein(maes_test[:,0], fitz_labels_path, test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "executionInfo": {
     "elapsed": 1037,
     "status": "ok",
     "timestamp": 1672891070971,
     "user": {
      "displayName": "Pradyumna Chari",
      "userId": "06555723454452599039"
     },
     "user_tz": 480
    },
    "id": "lAW7sr__GNjV",
    "outputId": "f3f1f778-444a-4c18-abff-6653b890d84a"
   },
   "outputs": [],
   "source": [
    "eval_model_fitz(fitz_labels_path, maes_test, test, \" \")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": [
    {
     "file_id": "1C9C7CGXx-5-Y-GLQaHor0AyTR5ebB9T_",
     "timestamp": 1639989285816
    }
   ]
  },
  "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}