utils.py

import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import Dataset, DataLoader
import pickle
from PIL import Image
from tqdm import tqdm
# from model import *
import h5py
import numpy as np
import imageio.v2 as imageio
import os
import matplotlib.pyplot as plt
import torch.optim.lr_scheduler as lr_scheduler
import sys
from torch.nn import DataParallel
import pandas as pd
from torchvision.utils import make_grid
import sklearn.metrics as metrics
import pdb
import torch.nn.init as init
import random
import warnings
warnings.filterwarnings("ignore")




def weight_init(m):
    if isinstance(m, (nn.Conv2d, nn.Linear)):
        init.kaiming_uniform_(m.weight, nonlinearity='relu')
        if m.bias is not None:
            init.constant_(m.bias, 0)

            
class PartNet2(nn.Module):
    def __init__(self):
        super(PartNet2, self).__init__()
        self.conv1_img = nn.Conv2d(3, 16, kernel_size=5, stride=1, padding=2)
        self.conv1_colorcheck = nn.Conv2d(72, 16, kernel_size=1, stride=1)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.conv2 = nn.Conv2d(32, 64, kernel_size=5, stride=1, padding=2)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.adaptive_pool = nn.AdaptiveAvgPool2d((10, 10))

        self.fc1 = nn.Linear(64 * 10 * 10, 64)
        self.fc2 = nn.Linear(64, 1)



    def forward(self, img, colorcheck):
        img = img.to(torch.float32)
        colorcheck = colorcheck.to(torch.float32)

        
        # colorcheck = colorcheck.view(colorcheck.size(0), -1, 10, 10)
        colorcheck = colorcheck.reshape(colorcheck.size(0), -1, 10, 10)


        img_features = self.pool1(self.relu1(self.conv1_img(img)))
        colorcheck_features = self.pool1(self.relu1(self.conv1_colorcheck(colorcheck)))

        
        colorcheck_features = F.adaptive_avg_pool2d(colorcheck_features, (img_features.size(2), img_features.size(3)))

        merged_features = torch.cat((img_features, colorcheck_features), dim=1)

        x = self.pool2(self.relu2(self.conv2(merged_features)))
        x = self.adaptive_pool(x)
        x = x.view(x.size(0), -1)

        x = self.fc2(torch.relu(self.fc1(x)))
        # x = F.relu(x) + 70# Initialize output to 95


        return x

class NPYDataset(Dataset):
    def __init__(self, h5_file_path, data_keys):
        self.all_frames = []
        self.all_labels = []
        self.all_colorchecks = []
        with h5py.File(h5_file_path, 'r') as f:
            for key in data_keys:
                if f'{key}/frames' in f and f'{key}/spo2_value' in f and f'{key}/colorcheck' in f:
                    self.all_frames.append(f[f'{key}/frames'][:])  # shape: (540, height, width, 3)
                    self.all_labels.append(f[f'{key}/spo2_value'][:])  # shape: (540,)
                    self.all_colorchecks.append(f[f'{key}/colorcheck'][:])  # shape: (24, 540, height, width, 3)
                else:
                    print(f"Key {key} not found in the file.")

    def __len__(self):
        return sum([len(frames) for frames in self.all_frames])

    def __getitem__(self, idx):
        frame_number = 0
        for video_frames, video_labels, video_colorchecks in zip(self.all_frames, self.all_labels, self.all_colorchecks):
            if idx < frame_number + len(video_frames):
                local_idx = idx - frame_number
                frame = video_frames[local_idx].astype(np.float32).transpose([2, 0, 1])
                label = video_labels[local_idx]
                # Correct handling for colorcheck indexing
                colorcheck_frames = video_colorchecks[:, local_idx, :, :, :]  # Extract all colorcheck data for this frame
                colorcheck = colorcheck_frames.transpose((0, 3, 1, 2)).astype(np.float32)  # Reformat to (24, 3, height, width)
                return torch.from_numpy(frame), torch.tensor(label, dtype=torch.float), torch.from_numpy(colorcheck)
            frame_number += len(video_frames)
        raise IndexError(f"Index {idx} is out of range for dataset with length {self.__len__()}")

    def get_subset(self, indices):
        frames = [self.all_frames[i-1] for i in indices]
        labels = [self.all_labels[i-1] for i in indices]
        colorchecks = [self.all_colorchecks[i-1] for i in indices]

        frames = np.vstack(frames)
        labels = np.concatenate(labels)
        # print('1 colorchecks.shape:',colorchecks.shape)
        colorchecks = np.concatenate(colorchecks, axis=1)
        # change from (24, 58860, 10, 10, 3) to (58860,24 , 3, 10, 10)
        colorchecks = colorchecks.transpose((1, 0, 4, 2, 3))
        # print('2 colorchecks.shape:',colorchecks.shape)

        return NPYDatasetFromArrays(frames, labels, colorchecks)

class NPYDatasetFromArrays(Dataset):
    def __init__(self, frames, labels, colorchecks):
        self.frames = frames
        self.labels = labels
        self.colorchecks = colorchecks

    def __len__(self):
        return len(self.frames)

    def __getitem__(self, idx):
        frame = self.frames[idx].astype(np.float32).transpose([2, 0, 1])
        label = self.labels[idx]
        colorcheck = self.colorchecks[idx].astype(np.float32)
        return torch.from_numpy(frame), torch.tensor(label, dtype=torch.float), torch.from_numpy(colorcheck)


class Combined_Loss(nn.Module):
    def __init__(self):
        super(Combined_Loss, self).__init__()

    def forward(self, output, target):
        loss1 = torch.sum((output - target) ** 2 * torch.abs(1 - target / 100))
        loss = loss1
        return loss

def train(args, model, device, train_loader, criterion, optimizer, epoch, train_losses):
    model.train()
    running_loss = 0.0
    progress_bar = tqdm(enumerate(train_loader), total=len(train_loader), desc=f'Epoch {epoch} Training')
    for batch_idx, (img, target, colorcheck) in progress_bar:
        # 75,400,400,3
        img = img.to(torch.float32)
        colorcheck = colorcheck.to(torch.float32)
        img, target, colorcheck = img.to(device), target.to(device), colorcheck.to(device)
        optimizer.zero_grad()
        output = model(img,colorcheck).squeeze()
        loss = criterion(output.float(), target.float())
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        progress_bar.set_postfix({'average_loss': running_loss / (batch_idx + 1)})
    average_loss = running_loss / len(train_loader)
    train_losses.append(average_loss)
    print(f'\033[92mTrain Epoch: {epoch} \tAverage Loss: {average_loss:.6f}\033[0m')
    return train_losses

def lowpass_filter(data, cutoff, fs):
    from scipy.signal import butter, filtfilt
    nyquist = 0.5 * fs
    normal_cutoff = cutoff / nyquist
    b, a = butter(1, normal_cutoff, btype='low', analog=False)
    y = filtfilt(b, a, data)
    return y



from scipy.signal import butter, filtfilt
from sklearn import metrics
from scipy.optimize import curve_fit

def fit_alpha_beta(first_10_targets, first_10_predictions, fixed_alpha):
    first_10_targets = np.array(first_10_targets)
    first_10_predictions = np.array(first_10_predictions)
    
    def fit_func(pred, alpha, beta):
        return alpha * pred + beta

    
    popt, _ = curve_fit(fit_func, first_10_predictions, first_10_targets)
    alpha, beta = popt

    
    if alpha < 0.1:
        
        beta = np.mean(first_10_targets - fixed_alpha * first_10_predictions)
        alpha = fixed_alpha

    return alpha, beta

def from_filtered_to_fit(all_predictions_flitered,all_targets,alpha,beta):
    predictions_fit = alpha * all_predictions_flitered + beta
    predictions_fit = np.clip(predictions_fit, 80, 100)
    predictions_fit = predictions_fit.flatten()
    
    predictions_fit_filtered = lowpass_filter(predictions_fit, 0.025 , 1)
    all_predictions_fit=[]
    all_predictions_fit.extend(predictions_fit_filtered)
    mean_average_error = np.mean(np.abs(np.array(all_predictions_fit) - np.array(all_targets)))
    
    correlation = np.corrcoef(all_predictions_fit, all_targets)[0, 1]
    return all_predictions_fit,mean_average_error,correlation

# search for the first 5 different targets from the start, if there is less than 5 different targets, use the different targets;re turn the index and the targets
def find_first_5_diff_targets(targets):
    diff_targets = []
    diff_targets_index = []
    for i, target in enumerate(targets):
        if target not in diff_targets:
            diff_targets.append(target)
            diff_targets_index.append(i)
        if len(diff_targets) == 5:
            break

    while len(diff_targets) < 5:
        diff_targets.append(diff_targets[-1])
        diff_targets_index.append(diff_targets_index[-1])
    return diff_targets_index, diff_targets

def test(args, model, device, test_loader, criterion, test_losses, epoch, test_file, flag):
    model.eval()
    test_loss = 0.0
    all_targets = []
    all_predictions = []
    all_predictions_fit = []
    all_predictions_flitered = []
    with torch.no_grad():
        progress_bar = tqdm(enumerate(test_loader), total=len(test_loader), desc=f'Epoch {epoch} Testing')
        for batch_idx, (img, target, colorcheck) in progress_bar:
            img, target, colorcheck = img.to(device), target.to(device), colorcheck.to(device)
            output = model(img, colorcheck)
            test_loss += criterion(output, target).item()
            progress_bar.set_postfix({'average_loss': test_loss / (batch_idx + 1)})
            
            
            all_targets.extend(target.cpu().numpy())
            all_predictions.extend(output.cpu().numpy())
    all_targets = np.array(all_targets)
    all_predictions = np.array(all_predictions)
    
    predictions_prefliter = 1 * all_predictions
    # predictions_prefliter = np.clip(predictions_prefliter, 50, 150)
    
    fs = 1  
    cutoff = 0.025  

    predictions_prefliter = predictions_prefliter.flatten()
    predictions_flitered = lowpass_filter(predictions_prefliter, cutoff, fs)
    all_predictions_flitered.extend(predictions_flitered)

    all_predictions_flitered = np.array(all_predictions_flitered)

    first_5_targets = all_targets[:5]
    first_5_predictions = all_predictions_flitered[:5]

    first_5_diff_index, first_5_diff_targets = find_first_5_diff_targets(first_5_targets)
    first_5_diff_predictions = [all_predictions_flitered[i] for i in first_5_diff_index]

    first_270_targets = all_targets[:270]
    first_270_predictions = all_predictions_flitered[:270]
    fixed_alpha = 1
    alpha, beta = fit_alpha_beta(first_270_targets, first_270_predictions, fixed_alpha)
    alpha_5, beta_5 = fit_alpha_beta(first_5_targets, first_5_predictions, fixed_alpha)
    alpha_5_diff, beta_5_diff = fit_alpha_beta(first_5_diff_targets, first_5_diff_predictions, fixed_alpha)

    average_loss = test_loss / len(test_loader)
    test_losses.append(average_loss)
    
    all_predictions_fit,mean_average_error,correlation = from_filtered_to_fit(all_predictions_flitered,all_targets,alpha,beta)
    
    all_predictions_fit_5,mean_average_error_5,correlation_5 = from_filtered_to_fit(all_predictions_flitered,all_targets,alpha_5,beta_5)

    save_test_file = str(test_file).replace('/', '_').replace('[', '').replace(']', '')

    
    print(f'\033[91m\nTest set: Average loss: {average_loss:.4f}\033[0m')
    print(f'\033[91m\nTest set: MAE: {mean_average_error:.4f}\033[0m')

    
    plt.figure(figsize=(15, 10))
    plt.plot(all_targets, label='spo2', color='red', marker='o', linestyle='-', markersize=2)
    plt.plot(all_predictions_flitered, label='Predicted SpO2', color='blue', marker='o', linestyle='-', markersize=2)
    plt.plot(all_predictions_fit, label=f'Fitted SpO2 from Predicted SpO2 (α={alpha}, β={beta:.2f})', color='green', alpha=0.7, marker='x', linestyle='-', markersize=2)
    plt.title(f'{str(test_file)}  Prediction vs spo2 MAE : {mean_average_error:.4f}   correlation : {correlation:.4f}')
    plt.xlabel('Time')
    plt.ylabel('Value')
    plt.legend(loc='upper left', bbox_to_anchor=(1, 1))
    plt.tight_layout()
    plt.ylim(75, 105)  
    plt.savefig(os.path.join(args.save_fig_dir, f'{test_file[0].replace("/", "_")}_Prediction_vs_Target.png'))
    plt.close()

    return all_predictions_fit, all_targets, mean_average_error, all_predictions_flitered, all_predictions_fit_5 

def save_metrics(all_metrics_spo2_pred_fit,args,description,calibration_num):
    
    avg_metrics_spo2_pred_fit = average_metrics(all_metrics_spo2_pred_fit)

    metrics_names = ['MAE', 'MAPE', 'RMSE', 'Pearson', 'Pearson_fix', 'SNR', 'MACC']
    stats_df = pd.DataFrame({
        'Metric': [f'{metric} Fitted SpO2' for metric in metrics_names],
        'Average': [
            f"{avg_metrics_spo2_pred_fit[metric][0]:.2f} +/- {avg_metrics_spo2_pred_fit[metric][1]:.2f}" for metric in metrics_names
        ]
    })
    
    fig, ax = plt.subplots(figsize=(12, 6))
    ax.axis('tight')
    ax.axis('off')

    
    plt.title(str(calibration_num)+description, fontsize=16, pad=20)

    
    table = ax.table(cellText=stats_df.values, colLabels=stats_df.columns, cellLoc='center', loc='center')
    table.auto_set_font_size(False)
    table.set_fontsize(12)
    table.scale(1.2, 1.8)  
    plt.savefig(os.path.join(args.save_fig_dir, str(calibration_num)+description+'.png'))
    plt.close()
    
    csv_file_path = os.path.join(args.save_fig_dir, str(calibration_num)+description+'.csv')
    stats_df.to_csv(csv_file_path, index=False)



def plot_losses(train_losses, test_losses, epoch, flag, args):
    plt.figure(figsize=(10, 5))
    plt.plot(train_losses, label='Training Loss')
    plt.plot(test_losses, label='Test Loss')
    plt.title(f'Losses up to Epoch {epoch}')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    # y set to [0-100]
    plt.ylim(0,100)
    plt.legend()
    plt.savefig(os.path.join(args.save_fig_dir,str(flag)+f'_loss_plot_epoch_{epoch}.png'))
    plt.close()

from scipy.optimize import minimize
from scipy.interpolate import interp1d
from scipy.ndimage import shift
from statsmodels.tsa.stattools import ccovf


def equalize_array_size(array1, array2):
    len1, len2 = len(array1), len(array2)
    dif_length = len1 - len2
    if dif_length < 0:
        array2 = array2[int(np.floor(-dif_length/2)):len2-int(np.ceil(-dif_length/2))]
    elif dif_length > 0:
        array1 = array1[int(np.floor(dif_length/2)):len1-int(np.ceil(dif_length/2))]
    return array1, array2, dif_length

def chisqr_align(reference, target, roi=None, order=1, init=0.1, bound=1):
    reference, target, dif_length = equalize_array_size(reference, target)
    if roi is None: roi = [0, len(reference) - 1]
    ROI = slice(int(roi[0]), int(roi[1]), 1)
    reference = reference / np.mean(reference[ROI])

    def fcn2min(x):
        shifted = shift(target, x, order=order)
        shifted = shifted / np.mean(shifted[ROI])
        return np.sum(((reference - shifted) ** 2)[ROI])

    minb = min([(init - bound), (init + bound)])
    maxb = max([(init - bound), (init + bound)])
    result = minimize(fcn2min, init, method='L-BFGS-B', bounds=[(minb, maxb)])
    return result.x[0] + int(np.floor(dif_length / 2))

def phase_align(reference, target, roi, res=100):
    ROI = slice(int(roi[0]), int(roi[1]), 1)
    x, r1 = highres(reference[ROI], kind='linear', res=res)
    x, r2 = highres(target[ROI], kind='linear', res=res)
    r1 -= r1.mean()
    r2 -= r2.mean()
    cc = ccovf(r1, r2, demean=False, unbiased=False)
    mod = 1 if np.argmax(cc) != 0 else -1
    return np.argmax(cc) * mod * (1. / res)

def highres(y, kind='cubic', res=100):
    y = np.array(y)
    x = np.arange(0, y.shape[0])
    f = interp1d(x, y, kind=kind)
    xnew = np.linspace(0, x.shape[0] - 1, x.shape[0] * res)
    ynew = f(xnew)
    return xnew, ynew

def calculate_metrics(predictions, true_values):
    metrics = {}
    errors = np.squeeze(predictions) - np.squeeze(true_values)
    metrics["MAE"] = np.mean(np.abs(errors))
    metrics["MAPE"] = np.mean(np.abs(errors / true_values)) * 100
    metrics["RMSE"] = np.sqrt(np.mean(np.square(errors)))
    metrics["Pearson"] = np.corrcoef(predictions, true_values)[0, 1]

    # Phase alignment
    s_phase = phase_align(true_values, predictions, [10, len(true_values)-10])
    aligned_predictions_phase = shift(predictions, s_phase, mode='nearest')
    correlation_phase = np.corrcoef(true_values, aligned_predictions_phase)[0, 1]

    # Chi-squared alignment
    s_chisqr = chisqr_align(true_values, predictions, [10, len(true_values)-10], init=-3.5, bound=2)
    aligned_predictions_chisqr = shift(predictions, s_chisqr, mode='nearest')
    correlation_chisqr = np.corrcoef(true_values, aligned_predictions_chisqr)[0, 1]

    # Choose the best alignment
    metrics["Pearson_fix"] = max(correlation_phase, correlation_chisqr)
    SNR_all = 10 * np.log10(np.mean(np.square(true_values)) / np.mean(np.square(errors)))
    metrics["SNR"] = np.mean(SNR_all)

    MACC_all = 1 - (np.mean(np.abs(errors)) / np.mean(true_values))
    metrics["MACC"] = np.mean(MACC_all)
    
    # change nan to np.nan
    for key in metrics.keys():
        if np.isnan(metrics[key]):
            metrics[key] = np.nan
    
    print("MAE=",metrics["MAE"])
    print("RMSE=",metrics["RMSE"])
    print("MAPE=",metrics["MAPE"])
    print("Pearson=",metrics["Pearson"])

    return metrics

def average_metrics(all_metrics):
    avg_metrics = {}
    for metric in all_metrics[0].keys():
        values = [m[metric] for m in all_metrics]
        avg_metrics[metric] = (np.nanmean(values), np.nanstd(values) / np.sqrt(len(values)))
    return avg_metrics

def set_random_seed(seed):
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False