Merge pull request #88 from Keasys/transformer

Transformer

pyVHR/BPM/BPM.py

@@ -263,11 +263,11 @@ class BPM:
 
     BVP signal must be a float32 numpy.ndarray with shape [num_estimators, num_frames].
     """
-    def __init__(self, data, fps, startTime=0, minHz=0.65, maxHz=4., verb=False):
+    def __init__(self, data, fps, startTime=0, minHz=0.65, maxHz=4.0, verb=False):
         """
         Input 'data' is a BVP signal defined as a float32 Numpy.ndarray with shape [num_estimators, num_frames]
         """
-        self.nFFT = 2048//1  # freq. resolution for STFTs
+        self.nFFT = 2048  # freq. resolution for STFTs
         if len(data.shape) == 1:
             self.data = data.reshape(1, -1)  # 2D array raw-wise
         else:
@@ -289,7 +289,7 @@ def BVP_to_BPM(self):
         """
         if self.data.shape[0] == 0:
             return np.float32(0.0)
-        Pfreqs, Power = Welch(self.data, self.fps, self.minHz, self.maxHz, self.nFFT)
+        Pfreqs, Power = Welch(self.data, self.fps)
         # -- BPM estimate
         Pmax = np.argmax(Power, axis=1)  # power max
         return Pfreqs[Pmax.squeeze()]

pyVHR/deepRPPG/TRANSFORMER/__init__.py

@@ -0,0 +1 @@
+from . import *

pyVHR/deepRPPG/TRANSFORMER/model.py

@@ -0,0 +1,15 @@
+import torch
+import torch.nn as nn
+
+class rPPGTransformer(nn.Transformer):
+    """Documentation for rPPGTransformer
+
+    """
+    def __init__(self, emb_dim, nhead, num_encoder_layers=12, num_decoder_layers=12):
+        super(rPPGTransformer, self).__init__(emb_dim, nhead, num_decoder_layers, num_encoder_layers)
+
+
+    def forward(self, X):
+        Y = super(rPPGTransformer,self).forward(X,X)
+        return torch.mean(Y, axis=1)
+

pyVHR/deepRPPG/__init__.py

@@ -1,2 +1,3 @@
 from .mtts_can import *
-from .hr_cnn import *
+from .hr_cnn import *
+from .transformer import *

pyVHR/deepRPPG/transformer.py

@@ -0,0 +1,65 @@
+import pyVHR
+import numpy as np
+import os
+import torch
+import time
+import torchvision.transforms as transforms
+
+from torch.utils.data import DataLoader
+
+from .TRANSFORMER.model import rPPGTransformer
+from .HR_CNN.utils import butter_bandpass_filter
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+def RPPG_TRANSFORMER_bvp_pred(frames):
+    model_path = pyVHR.__path__[0] + "/deepRPPG/TRANSFORMER/model_vipl_next_50_1.dct"
+    model_name = model_path.split('/')[-1].split('.')[0]
+    if not os.path.isfile(model_path):
+        url = ""
+        print('Downloading rPPG Transformer model...')
+        r = requests.get(url, allow_redirects=True)
+        open(model_path, 'wb').write(r.content)
+    model = rPPGTransformer(250, nhead=1, num_encoder_layers=12, num_decoder_layers=12).to(device) 
+    model.load_state_dict(torch.load(model_path))
+
+
+    model.eval()
+
+    frames = torch.as_tensor(frames)
+    frames = torch.utils.data.TensorDataset(frames.to(device, dtype=torch.float))
+    frames_loader = torch.utils.data.DataLoader(frames,
+                                                batch_size=16,
+                                                shuffle=True,
+                                                drop_last=True)
+
+    start = time.time()
+
+    outputs = []
+
+    for i, chrom_patches in enumerate(frames_loader):
+        output = model(chrom_patches[0]).detach().squeeze().cpu().numpy()
+        outputs.append(output)
+    end = time.time()
+    print("processing time: ", end - start)
+
+    outputs=np.array(outputs)
+
+    mean = np.mean(outputs)
+    outputs = (outputs - np.mean(outputs)) / np.std(outputs)
+
+    fs = 30
+
+    lowcut = 0.8
+    highcut = 6
+
+    filtered_outputs = butter_bandpass_filter(outputs, lowcut, highcut, fs, order=4)
+    filtered_outputs = (filtered_outputs - np.mean(filtered_outputs)) / np.std(filtered_outputs)
+
+    # rearrange output unbatched
+    filtered_outputs = np.array(filtered_outputs)
+    filtered_outputs = np.concatenate(np.concatenate(filtered_outputs, axis=0), axis=0)
+
+    return filtered_outputs
+
+

pyVHR/utils/errors.py

@@ -200,7 +200,7 @@ def get_SNR(bvps, fps, reference_hrs, timesES):
     nfft = np.ceil((60*2*NyquistF)/FResBPM)
     SNRs = []
     for idx, bvp in enumerate(bvps):
-        curr_ref = reference_hrs[int(timesES[idx])]
+        curr_ref = reference_hrs[idx]
         pfreqs, power = Welch(bvp, fps, nfft=nfft)
         GTMask1 = np.logical_and(pfreqs>=curr_ref-interv1, pfreqs<=curr_ref+interv1)
         GTMask2 = np.logical_and(pfreqs>=(curr_ref*2)-interv2, pfreqs<=(curr_ref*2)+interv2)