<jupyter_start><jupyter_text>Training and evaluating the RhythmNet modelThis is where everything regarding the RhythmNet model happens. All the spatial temporal maps of the videos should be generated with `generate_st_maps.py`Pipeline0. Get HR as correct signal from pyVHR1. Extract spatial temporal map2. Train3. Eval4. Predict<jupyter_code>import os
from utils.dataset import DataLoaderRhythmNet
from utils.model_utils import plot_loss, load_model_if_checkpointed, save_model_checkpoint
from models.rhythmNet import RhythmNet
from loss_func.rhythmnet_loss import RhythmNetLoss
import torch
import config
from torch.utils.tensorboard import SummaryWriter

import utils.train_val as train_val
from utils.plot_scripts import plot_train_test_curves, bland_altman_plot, gt_vs_est, create_plot_for_tensorboard
import numpy as np
import pickle
import matplotlib.pyplot as plt
# check path to checkpoint directory
if config.CHECKPOINT_PATH:
    if not os.path.exists(config.CHECKPOINT_PATH):
        os.makedirs(config.CHECKPOINT_PATH)
        print("Output directory is created")

model = RhythmNet()

if torch.cuda.is_available():
    print('GPU available... using GPU')
    torch.cuda.manual_seed_all(42)
else:
    print("GPU not available, using CPU")

if config.CHECKPOINT_PATH:
    checkpoint_path = os.path.join(os.getcwd(), config.CHECKPOINT_PATH)
    if not os.path.exists(checkpoint_path):
        os.makedirs(checkpoint_path)
        print("Output directory is created")

device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')

model.to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, factor=0.8, patience=5, verbose=True
)
# loss_fn = nn.L1Loss()
loss_fn = RhythmNetLoss()

testset = trainset = None

# Initialize SummaryWriter object
writer = SummaryWriter()
# video files are st_maps
test_video_dir = '../data/st_maps/test'
train_video_dir = '../data/st_maps/train'
video_files_train = os.listdir("../data/st_maps/train/")
video_files_test = os.listdir("../data/st_maps/test/")
test_video_items = [test_video_dir+'/'+filename for filename in os.listdir(test_video_dir)] 
train_video_items = [train_video_dir+'/'+filename for filename in os.listdir(train_video_dir)] 
TARGET_SIGNAL_DIR = '../data/hr_data'
# --------------------------------------
# Build Dataloaders
# --------------------------------------
def collate_fn(batch):
    '''the collate function ensures that the maps of different sizes are stackable'''
    batched_st_map, batched_targets = [], []
    # for data in batch:
    #     batched_st_map.append(data["st_maps"])
    #     batched_targets.append(data["target"])
    # # torch.stack(batched_output_per_clip, dim=0).transpose_(0, 1)
    return batch

train_set = DataLoaderRhythmNet(st_maps_path=train_video_items, target_signal_path=TARGET_SIGNAL_DIR)
batch_size = 1
train_loader = torch.utils.data.DataLoader(
    dataset=train_set,
    batch_size=batch_size,
    num_workers=0,
    shuffle=False,
    collate_fn=collate_fn
)
print('\nTrain DataLoader constructed successfully!')

# Code to use multiple GPUs (if available)
if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    model = torch.nn.DataParallel(model)
# --------------------------------------
# Load checkpointed model (if  present)
# --------------------------------------
if device == "cpu":
    load_on_cpu = True
else:
    load_on_cpu = False
model, optimizer, checkpointed_loss, checkpoint_flag = load_model_if_checkpointed(model, optimizer, checkpoint_path, load_on_cpu=load_on_cpu)
if checkpoint_flag:
    print(f"Checkpoint Found! Loading from checkpoint :: LOSS={checkpointed_loss}")
else:
    print("Checkpoint Not Found! Training from beginning")
def rmse(l1, l2):
    return np.sqrt(np.mean((l1-l2)**2))

def mae(l1, l2):
    return np.mean([abs(item1-item2)for item1, item2 in zip(l1, l2)])

def compute_criteria(target_hr_list, predicted_hr_list):
    pearson_per_signal = []
    HR_MAE = mae(np.array(predicted_hr_list), np.array(target_hr_list))
    HR_RMSE = rmse(np.array(predicted_hr_list), np.array(target_hr_list))

    # for (gt_signal, predicted_signal) in zip(target_hr_list, predicted_hr_list):
    #     r, p_value = pearsonr(predicted_signal, gt_signal)
    #     pearson_per_signal.append(r)

    # return {"MAE": np.mean(HR_MAE), "RMSE": HR_RMSE, "Pearson": np.mean(pearson_per_signal)}
    return {"MAE": np.mean(HR_MAE), "RMSE": HR_RMSE}
# --------------------------------------
# Load checkpointed model (if  present)
# --------------------------------------
if config.DEVICE == "cpu":
    load_on_cpu = True
else:
    load_on_cpu = False
model, optimizer, checkpointed_loss, checkpoint_flag = load_model_if_checkpointed(model, optimizer,
                                                                                    checkpoint_path,
                                                                                    load_on_cpu=load_on_cpu)
if checkpoint_flag:
    print(f"Checkpoint Found! Loading from checkpoint :: LOSS={checkpointed_loss}")
    with open('models/train_loss.pkl', "rb") as f:
        pickle.load(f)
else:
    print("Checkpoint Not Found! Training from beginning")

test_signal_path = ['../data/st_maps/test/']
test_set = DataLoaderRhythmNet(st_maps_path=test_video_items, target_signal_path=test_signal_path)
test_loader = torch.utils.data.DataLoader(
    dataset=test_set,
    batch_size=batch_size,
    num_workers=0,
    shuffle=False,
    collate_fn=collate_fn
)
print('\nEvaluation DataLoader constructed successfully!')<jupyter_output>Evaluation DataLoader constructed successfully!<jupyter_text>Train<jupyter_code># Training
EPOCHS = 5
if not checkpoint_flag:
    train_losses = []
    train_loss_per_epoch = []
    eval_loss_per_epoch = []
    test_losses = []
    for epoch in range(EPOCHS):
        # short-circuit for evaluation
        
        target_hr_list, predicted_hr_list, train_loss = train_val.train_fn(model, train_loader, optimizer, loss_fn, batch_size=batch_size)
        train_losses.append(train_loss)
        # Save model with final train loss (script to save the best weights?)
        if checkpointed_loss != 0.0:
            if train_loss < checkpointed_loss:
                save_model_checkpoint(model, optimizer, train_loss, checkpoint_path)
                checkpointed_loss = train_loss
            else:
                pass
        else:
            if len(train_loss_per_epoch) > 0:
                if train_loss < min(train_loss_per_epoch):
                    save_model_checkpoint(model, optimizer, train_loss, checkpoint_path)
            else:
                save_model_checkpoint(model, optimizer, train_loss, checkpoint_path)

        metrics = compute_criteria(target_hr_list, predicted_hr_list)

        for metric in metrics.keys():
            writer.add_scalar(f"Train/{metric}", metrics[metric], epoch)

        print(f"\nFinished Train [Epoch: {epoch + 1}/{EPOCHS}]",
                "\nTraining Loss: {:.3f} |".format(train_loss),
                "HR_MAE : {:.3f} |".format(metrics["MAE"]),
                "HR_RMSE : {:.3f} |".format(metrics["RMSE"]),)
                # "Pearsonr : {:.3f} |".format(metrics["Pearson"]), )

        train_loss_per_epoch.append(train_loss)
        writer.add_scalar("Loss/train", train_loss, epoch+1)

        
        
        # validation
        target_hr_list, predicted_hr_list, test_loss = train_val.eval_fn(model, test_loader, loss_fn, batch_size=batch_size)
        test_losses.append(test_loss)
        # truth_hr_list.append(target)
        # estimated_hr_list.append(predicted)
        metrics = compute_criteria(target_hr_list, predicted_hr_list)
        for metric in metrics.keys():
            writer.add_scalar(f"Test/{metric}", metrics[metric], epoch)

        print(f"\nFinished Test [Epoch: {epoch + 1}/{EPOCHS}]",
                "\nTest Loss: {:.3f} |".format(test_loss),
                "HR_MAE : {:.3f} |".format(metrics["MAE"]),
                "HR_RMSE : {:.3f} |".format(metrics["RMSE"]),)

        writer.add_scalar("Loss/test", test_loss, epoch)

        # Plots on tensorboard
        # ba_plot_image = create_plot_for_tensorboard('bland_altman', target_hr_list, predicted_hr_list)
        # gtvsest_plot_image = create_plot_for_tensorboard('gt_vs_est', target_hr_list, predicted_hr_list)
        # writer.add_image('BA_plot', ba_plot_image, epoch)
        # writer.add_image('gtvsest_plot', gtvsest_plot_image, epoch)

    mean_loss = np.mean(train_loss_per_epoch)
    # Save the mean_loss value for each video instance to the writer
    print(f"Avg Training Loss: {np.mean(mean_loss)} for {EPOCHS} epochs")
    writer.flush()
    
if not checkpoint_flag:
    with open('models/train_loss.pkl', "wb") as f:
        pickle.dump(train_losses, f)
    with open('models/test_loss.pkl', "wb") as f:
        pickle.dump(test_losses, f)
else:
    with open('models/train_loss.pkl', "rb") as f:
        train_losses = pickle.load(f)
    with open('models/test_loss.pkl', "rb") as f:
        test_losses = pickle.load(f)
fig, ax = plt.subplots()
ax.plot(range(1,6), train_losses, label='train loss')
ax.plot(range(1,6), test_losses, label='test loss')
ax.set_xlabel('Epoch')
ax.set_ylabel('Loss')
ax.set_xticks(range(1,6))
ax.legend();
plt.savefig('plots/loss.pdf')<jupyter_output><empty_output><jupyter_text>Test<jupyter_code>scenarios = ['talk', 'gym', 'rotation', 'resting'] # just hardcode the scenarios according to the folder, this only works since shuffling is off
def gt_vs_est(data1, data2, plot_path=None):
    data1 = np.asarray(data1)
    data2 = np.asarray(data2)
    fig, ax = plt.subplots()
    scenario_arr = np.array(scenarios)
    for scenario in set(scenarios):
        mask = scenario_arr == scenario
        ax.scatter(data1[mask],data2[mask], label=scenario)
    ax.plot([np.min(data1),np.max(data1)],[np.min(data1),np.max(data1)],'r',alpha=0.2)
    ax.set_title('true labels vs estimated')
    ax.set_ylabel('estimated HR')
    ax.set_xlabel('true HR')
    ax.legend()

    fig.savefig(plot_path + f'/true_vs_est.pdf')

def bland_altman_plot(data1, data2, plot_path=None):
    data1 = np.array(data1)
    data2 = np.array(data2)
    mean = np.mean([data1, data2], axis=0)
    diff = data1 - data2  # Difference between data1 and data2
    md = np.mean(diff)  # Mean of the difference
    sd = np.std(diff, axis=0)  # Standard deviation of the difference

    fig, ax = plt.subplots()
    scenario_arr = np.array(scenarios)
    for scenario in set(scenarios):
        mask = scenario_arr == scenario
        ax.scatter(mean[mask], diff[mask], label=scenario)
    ax.axhline(md, color='gray', linestyle='--')
    ax.axhline(md + 1.96 * sd, color='gray', linestyle='--')
    ax.axhline(md - 1.96 * sd, color='gray', linestyle='--')
    ax.set_xlabel(r'$\frac{1}{2}(HR_{gt}-HR_{est})$')
    ax.set_ylabel(r'$HR_{gt}-HR_{est}$')
    ax.legend()

    fig.savefig(plot_path + f'/bland-altman_new.pdf')
# this code is basically just to get the plots
EPOCHS_TEST = 1
print(f"Finished Training, Validating {len(video_files_test)} video files for {EPOCHS_TEST} Epochs")

eval_loss_per_epoch = []
test_losses = []
for epoch in range(EPOCHS_TEST):
    # validation
    target_hr_list, predicted_hr_list, test_loss = train_val.eval_fn(model, test_loader, loss_fn, batch_size=batch_size)
    test_losses.append(test_loss)
    metrics = compute_criteria(target_hr_list, predicted_hr_list)

    print(f"\nFinished Test [Epoch: {epoch + 1}/{EPOCHS_TEST}]",
            "\nTest Loss: {:.3f} |".format(test_loss),
            "HR_MAE : {:.3f} |".format(metrics["MAE"]),
            "HR_RMSE : {:.3f} |".format(metrics["RMSE"]),)

# Plots on the local storage.
gt_vs_est(target_hr_list, predicted_hr_list, plot_path='plots')
bland_altman_plot(target_hr_list, predicted_hr_list, plot_path='plots')
writer.close()
print("done")
preds = []
targets = []
model.eval()
for data in test_loader:
    st_maps, target = data[0]['st_maps'], data[0]['target']
    targets.append(target)
    with torch.no_grad():
        preds.append(model(st_maps, target)[0])
import matplotlib.pyplot as plt
fig, ax = plt.subplots()
i = 3
x = range(len(targets[i]))
ax.plot(preds[i].flatten(), label='predicted')
ax.plot(targets[i], label='gt')
ax.set_title(f"prediction for scenario '{scenarios[i]}'")
ax.set_xlabel('time (s)')
ax.set_ylabel('Heart rate (BPM)')
ax.legend();
plt.savefig('plots/predictedHR.pdf')
from pyVHR.utils.errors import getErrors, displayErrors
import pandas as pd
errors = []
for i, signal in enumerate(zip(preds,targets,scenarios)):
    RMSE, MAE, MAX, PCC, CCC = getErrors(np.array(signal[0]), np.array(signal[1]),None,None)
    errors.append((np.mean(RMSE), np.mean(MAE), np.mean(MAX), np.mean(PCC), np.mean(CCC), signal[2]))
df = pd.DataFrame(errors, columns=['RMSE', 'MAE', 'MAX', 'PCC', 'CCC', 'Scenario'])
df
dfg = df.groupby(['Scenario']).mean()
dfg.style.bar()
dfg_latex = pd.DataFrame(dfg,columns=['RMSE','MAE','PCC'])
tab = dfg_latex.style.format({'RMSE': '{:.2f}', 'MAE': '{:.2f}', 'PCC': '{:.4f}'}).to_latex(caption='Mean score for videos across scenarios',position_float="centering", label='tab:HR_est_scores',hrules=True)
print(tab)<jupyter_output>\begin{table}
\centering
\caption{Mean score for videos across scenarios}
\label{tab:HR_est_scores}
\begin{tabular}{lrrr}
\toprule
 & RMSE & MAE & PCC \\
Scenario &  &  &  \\
\midrule
gym & 46.26 & 44.29 & -0.0433 \\
resting & 3.44 & 2.83 & 0.2466 \\
rotation & 4.63 & 3.72 & 0.1231 \\
talk & 16.49 & 15.83 & 0.2083 \\
\bottomrule
\end{tabular}
\end{table}
{'fraction_non_alphanumeric': 0.14660414682394315, 'fraction_numerical': 0.01106308154443549, 'mean_word_length': 3.614604462474645, 'pattern_counts': {'":': 7, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '337090', 'license': 'MIT', 'n_tokens_mistral': 4808, 'n_tokens_neox': 4663, 'n_words': 935, 'path': 'RhythmNet/train.ipynb', 'repo_name': 'AndreasLH/Heart-Rate-Estimation', 'size': '13649'}