sarc-asm 0.2.2__py2.py3-none-any.whl

contraction_net/__init__.py

@@ -0,0 +1,4 @@
+from .contraction_net import ContractionNet
+from .training import Trainer
+from .data import DataProcess
+from .prediction import predict_contractions

contraction_net/annotation.py

@@ -0,0 +1,148 @@
+import random
+import os
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.widgets import Button
+
+
+class TimeSeriesAnnotator:
+    """
+    Class to annotate contraction intervals in multiple time-series files and save the results.
+
+    This class allows users to interactively annotate the start and end of contractions
+    in time-series data by clicking on the plot. The user can press the left mouse button
+    at the start of a contraction interval and release at the end of the contraction interval.
+    The results are saved to a specified output directory.
+
+    Parameters
+    ----------
+    file_list : list of str
+        List of file paths to the time-series data files to be annotated.
+    output_dir : str
+        Directory where the output files with annotated contraction intervals will be saved.
+    figsize : tuple, optional
+        Figure size in inches.
+
+    Methods
+    -------
+    load_next_file()
+        Load the next time-series file for annotation.
+    press_callback(event)
+        Record the x-coordinate (time) where the mouse button is pressed.
+    release_callback(event)
+        Record the x-coordinate (time) where the mouse button is released.
+    save_and_load_next(event)
+        Save the current annotations and load the next time-series file.
+    reset_annotations(event)
+        Reset the current annotations for the current time-series file.
+
+    Notes
+    -----
+    In an interactive Jupyter notebook, add the following before running the TimeSeriesAnnotator:
+
+    ```
+    matplotlib.use('nbagg')
+    %matplotlib notebook
+    ```
+    """
+
+    def __init__(self, file_list, output_dir, figsize=(11, 3.5)):
+        self.file_list = [f for f in file_list if not f.endswith('_contr.txt')]
+        self.output_dir = output_dir
+        self.current_file_index = 0
+        self.start_contraction = []
+        self.end_contraction = []
+        self.fig, self.ax = plt.subplots(figsize=figsize, constrained_layout=True)
+        self.load_next_file()
+
+        self.fig.canvas.mpl_connect('button_press_event', self.press_callback)
+        self.fig.canvas.mpl_connect('button_release_event', self.release_callback)
+
+        save_ax = self.fig.add_axes([0.85, 0.01, 0.1, 0.075])
+        self.save_button = Button(save_ax, 'Save & Next', color='lightgoldenrodyellow', hovercolor='0.975')
+        self.save_button.on_clicked(self.save_and_load_next)
+
+        reset_ax = self.fig.add_axes([0.75, 0.01, 0.1, 0.075])
+        self.reset_button = Button(reset_ax, 'Reset', color='lightcoral', hovercolor='0.975')
+        self.reset_button.on_clicked(self.reset_annotations)
+
+        plt.show()
+
+    def load_next_file(self):
+        """Load the next time-series file for annotation."""
+        if self.current_file_index < len(self.file_list):
+            file_path = self.file_list[self.current_file_index]
+            self.data = np.loadtxt(file_path)
+            self.ax.clear()
+            self.ax.plot(self.data, c='k')
+            self.ax.set_title(f'Annotating {os.path.basename(file_path)}')
+            self.start_contraction = []
+            self.end_contraction = []
+            plt.draw()
+        else:
+            plt.close(self.fig)
+            print('All files annotated.')
+
+    def press_callback(self, event):
+        """Record the x-coordinate (time) where the mouse button is pressed."""
+        self.ax.axvline(event.xdata, c='r', linestyle=':')
+        self.start_contraction.append(event.xdata)
+
+    def release_callback(self, event):
+        """Record the x-coordinate (time) where the mouse button is released."""
+        self.ax.axvline(event.xdata, c='r', linestyle=':')
+        self.ax.axvspan(self.start_contraction[-1], event.xdata, alpha=0.5, color='red')
+        self.end_contraction.append(event.xdata)
+
+    def save_and_load_next(self, event):
+        """Save the current annotations and load the next time-series file."""
+        if self.current_file_index < len(self.file_list):
+            file_path = self.file_list[self.current_file_index]
+            output_file = os.path.join(self.output_dir, f'{os.path.basename(file_path).split(".")[0]}_contr.txt')
+            start_end_contraction = np.asarray([self.start_contraction, self.end_contraction])
+            np.savetxt(output_file, start_end_contraction.T, header='start end', comments='', fmt='%f')
+            print(f'Saved contractions for {os.path.basename(file_path)} to {output_file}')
+
+            self.current_file_index += 1
+            self.load_next_file()
+
+    def reset_annotations(self, event):
+        """Reset the current annotations for the current time-series file."""
+        self.start_contraction = []
+        self.end_contraction = []
+        self.ax.clear()
+        self.ax.plot(self.data, c='k')
+        self.ax.set_title(f'Annotating {os.path.basename(self.file_list[self.current_file_index])}')
+        plt.draw()
+
+
+def split_trace(filepath, output_dir, chunk_size=512, p=1):
+    """
+    Split a time-series trace into smaller chunks and save them to separate files.
+
+    Parameters
+    ----------
+    filepath : str
+        Path to the input time-series data file.
+    output_dir : str
+        Directory where the output chunk files will be saved.
+    chunk_size : int, optional
+        The size of each chunk, by default 512.
+    p : float, optional
+        Probability of saving each chunk, by default 1 (save all chunks).
+    """
+    # Read the time-series data
+    data = np.loadtxt(filepath)
+
+    # Create the output directory if it doesn't exist
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
+    # Calculate the number of chunks
+    num_chunks = len(data) // chunk_size + (len(data) % chunk_size > 0)
+
+    # Split the data into chunks of size `chunk_size` and save each chunk to a separate file
+    for i in range(num_chunks):
+        if random.random() < p:
+            chunk = data[i * chunk_size:(i + 1) * chunk_size]
+            np.savetxt(os.path.join(output_dir, f'{os.path.basename(filepath).split(".")[0]}_{i}.txt'), chunk)

contraction_net/contraction_net.py

@@ -0,0 +1,81 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ContractionNet(nn.Module):
+    """
+        ContractionNet model for detecting contraction intervals from time-series data of individual Z-band positions
+        and sarcomere lengths of beating cardiomyocytes.
+
+        This neural network is designed to handle noisy data and distinguish between contracting and non-contracting intervals.
+        The network first extracts various features from a single input time-series by two convolutional layers with kernel size 5, followed by a dilated convolution in the third layer
+        to capture broader temporal patterns. Each convolution is followed by instance normalization and ReLU activation.
+        A self-attention layer enhances focus on salient features. The processed signal then undergoes two further
+        convolutions before being outputted through a sigmoid activation function.
+
+        Methods
+        -------
+        forward(x)
+            Forward pass through the network.
+        """
+    def __init__(self, n_filter=64, in_channels=1, out_channels=2, dropout_rate=0.5):
+        """
+        Parameters
+        ----------
+        n_filter : int, optional
+            Number of filters in the convolutional layers (default is 64).
+        in_channels : int, optional
+            Number of input channels (default is 1).
+        out_channels : int, optional
+            Number of output channels (default is 2).
+        dropout_rate : float, optional
+            Dropout rate (default is 0.5).
+        """
+        super(ContractionNet, self).__init__()
+        self.conv1 = nn.Conv1d(in_channels=in_channels, out_channels=n_filter, kernel_size=5, padding=2)
+        self.in1 = nn.InstanceNorm1d(n_filter)
+        self.conv2 = nn.Conv1d(in_channels=n_filter, out_channels=n_filter * 2, kernel_size=5, padding=2)
+        self.bn2 = nn.InstanceNorm1d(n_filter * 2)
+        self.conv3 = nn.Conv1d(in_channels=n_filter * 2, out_channels=n_filter * 4, kernel_size=5, padding=4,
+                               dilation=2)
+        self.bn3 = nn.InstanceNorm1d(n_filter * 4)
+        self.attention = nn.MultiheadAttention(embed_dim=n_filter * 4, num_heads=4, batch_first=True)
+        self.norm1 = nn.LayerNorm(n_filter * 4)
+        self.dropout_attention = nn.Dropout(dropout_rate)
+        self.conv4 = nn.Conv1d(in_channels=n_filter * 4, out_channels=n_filter * 2, kernel_size=5, padding=2,
+                               dilation=1)
+        self.bn4 = nn.InstanceNorm1d(n_filter * 2)
+        self.dropout_pre_output = nn.Dropout(dropout_rate)
+        self.conv_out = nn.Conv1d(in_channels=n_filter * 2, out_channels=out_channels, kernel_size=1)
+
+    def forward(self, x):
+        """
+        Forward pass through the network.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (batch_size, in_channels, sequence_length).
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor of shape (batch_size, out_channels, sequence_length) after sigmoid activation.
+        torch.Tensor
+            Raw output tensor of shape (batch_size, out_channels, sequence_length).
+        """
+        x = F.relu(self.in1(self.conv1(x)))
+        x = F.relu(self.bn2(self.conv2(x)))
+        x = F.relu(self.bn3(self.conv3(x)))
+        x = x.transpose(1, 2)
+        residual = x
+        x, attention_weights = self.attention(x, x, x)
+        x = x + residual
+        x = self.norm1(x)
+        x = self.dropout_attention(x)
+        x = x.transpose(1, 2)
+        x = F.relu(self.bn4(self.conv4(x)))
+        x = self.dropout_pre_output(x)
+        x = self.conv_out(x)
+        return torch.sigmoid(x), x

contraction_net/data.py

@@ -0,0 +1,158 @@
+import glob
+import os
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+from .utils import distance_transform
+
+
+class DataProcess(Dataset):
+    """
+    A Dataset class for creating training data objects for ContractionNet training.
+
+    Parameters
+    ----------
+    source_dir : Tuple[str, str]
+        Tuple containing paths to the directories of training data [images, labels]. Images should be in .tif format.
+    input_len : int, optional
+        Length of the input sequences (default is 512).
+    normalize : bool, optional
+        Whether to normalize each time-series (default is False).
+    aug_factor : int, optional
+        Factor for image augmentation (default is 10).
+    val_split : float, optional
+        Validation split for training (default is 0.2).
+    noise_amp : float, optional
+        Amplitude of Gaussian noise for image augmentation (default is 0.2).
+    aug_p : float, optional
+        Probability of applying augmentation (default is 0.5).
+    random_offset : float, optional
+        Amplitude of random offset applied to the input sequences (default is 0.25).
+    random_outlier : float, optional
+        Amplitude of random outliers added to the input sequences (default is 0.5).
+    random_drift : Tuple[float, float], optional
+        Parameters for random drift: (frequency, amplitude) (default is (0.01, 0.2)).
+    random_swap : float, optional
+        Probability of randomly swapping the sign of the input sequences (default is 0.5).
+    random_subsampling : Tuple[int, int], optional
+        Range for random subsampling intervals (default is None).
+
+    Methods
+    -------
+    __len__()
+        Returns the total number of samples.
+    __getitem__(idx)
+        Generates one sample of data.
+
+    """
+    def __init__(self, source_dir, input_len=512, normalize=False, val_split=0.2, aug_factor=10, aug_p=0.5,
+                 noise_amp=0.2, random_offset=0.25, random_outlier=0.5, random_drift=(0.01, 0.2), random_swap=0.5,
+                 random_subsampling=None):
+        self.source_dir = source_dir
+        self.data = []
+        self.is_real = []
+        self.input_len = input_len
+        self.val_split = val_split
+        self.normalize = normalize
+        self.aug_factor = aug_factor
+        self.aug_p = aug_p
+        self.noise_amp = noise_amp
+        self.random_offset = random_offset
+        self.random_drift = random_drift
+        self.random_outlier = random_outlier
+        self.random_subsampling = random_subsampling
+        self.random_swap = random_swap
+        self.mode = 'train'
+        self.__load_and_edit()
+        if self.aug_factor is not None:
+            self.__augment()
+
+    def __load_and_edit(self):
+        """
+        Loads and preprocesses the input data files from the source directory.
+        """
+        files = glob.glob(self.source_dir + '*.txt')
+        print(f'{len(files)} files found')
+        files_input = [f for f in files if 'peaks' not in f and 'contr' not in os.path.basename(f)]
+        for file_i in files_input:
+            input_i = np.loadtxt(file_i)[: self.input_len]
+            if self.normalize:
+                input_i -= np.mean(input_i)
+                input_i /= np.std(input_i)
+            start_end_systole_i = np.loadtxt(file_i[:-4] + '_contr.txt')[: self.input_len]
+            if len(start_end_systole_i) == len(input_i):
+                contr_i = start_end_systole_i
+            else:
+                contr_i = np.zeros_like(input_i)
+                start_end_systole_i = np.clip(start_end_systole_i, a_min=0, a_max=self.input_len)
+                if start_end_systole_i.size != 0:
+                    if len(start_end_systole_i.shape) == 1:
+                        contr_i[int(start_end_systole_i[0]): int(start_end_systole_i[1])] = 1
+                    elif len(start_end_systole_i.shape) == 2:
+                        for start, end in start_end_systole_i.T:
+                            contr_i[int(start): int(end)] = 1
+            if 'sim' in file_i or 'noise' in file_i:
+                is_real_i = False
+            else:
+                is_real_i = True
+            self.data.append([input_i, contr_i, distance_transform(input_i, contr_i)])
+            self.is_real.append(is_real_i)
+        self.data = np.asarray(self.data)
+
+    def __augment(self):
+        """
+        Applies data augmentation techniques to the loaded data.
+        """
+        _data = self.data.copy()
+        self.data = []
+        for i, d_i in enumerate(_data):
+            for j in range(self.aug_factor):
+                d_ij = d_i.copy()
+                std_ij = np.std(d_ij[0])
+                d_ij[0] = d_ij[0] / std_ij
+                d_ij[0] = np.random.choice([-1, 1], p=(1-self.random_swap, self.random_swap)) * d_ij[0]
+                if self.random_subsampling is not None and np.random.binomial(1, self.aug_p):
+                    d_ij = np.zeros_like(d_ij)
+                    d_ij_short = np.delete(d_i, np.arange(0, 512, np.random.randint(self.random_subsampling[0],
+                                                                                    self.random_subsampling[1])),
+                                           axis=1)
+                    d_ij[:, :d_ij_short.shape[1]] = d_ij_short
+                if not self.is_real[i] and np.random.binomial(1, self.aug_p):
+                    d_ij[0] += np.random.normal(0, self.noise_amp, size=d_i.shape[1])
+                if np.random.binomial(1, self.aug_p):
+                    d_ij[0] += np.random.normal(0, self.random_offset, size=1)
+                if np.random.binomial(1, self.aug_p):
+                    n_outliers = np.random.randint(0, 10)
+                    t_outliers = np.random.randint(0, self.input_len, n_outliers)
+                    d_ij[0, t_outliers] += np.random.normal(0, self.random_outlier, n_outliers)
+                if np.random.binomial(1, self.aug_p):
+                    d_ij[0] += self.random_drift[1] * np.cos(self.random_drift[0] * np.arange(self.input_len))
+                d_ij[0] = d_ij[0] * std_ij
+                self.data.append(d_ij)
+        self.data = np.asarray(self.data)
+
+    def __len__(self):
+        """
+        Returns the total number of samples.
+        """
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        """
+        Generates one sample of data.
+
+        Parameters
+        ----------
+        idx : int
+            Index of the sample to retrieve.
+
+        Returns
+        -------
+        sample : dict
+            Dictionary containing 'input', 'target', and 'distance' tensors.
+        """
+        data_torch = torch.from_numpy(self.data[idx]).float()
+        sample = {'input': data_torch[0], 'target': data_torch[1], 'distance': data_torch[2]}
+        return sample

contraction_net/losses.py

@@ -0,0 +1,97 @@
+import torch
+from torch import nn
+
+
+class BCEDiceLoss(nn.Module):
+    """
+    A combination of Binary Cross Entropy (BCE) and Dice Loss for binary segmentation tasks.
+
+    Parameters
+    ----------
+    loss_params : tuple, optional
+        A tuple containing the weights for BCE and Dice losses respectively. Default is (1, 1).
+
+    Methods
+    -------
+    dice_loss(inputs, targets, epsilon=1e-6)
+        Computes the Dice loss.
+
+    forward(inputs, targets)
+        Computes the combined BCE and Dice loss.
+    """
+
+    def __init__(self, loss_params=(1, 1)):
+        super(BCEDiceLoss, self).__init__()
+        self.bce_loss = nn.BCEWithLogitsLoss()
+        self.loss_params = loss_params
+
+    def dice_loss(self, inputs, targets, epsilon=1e-6):
+        inputs = torch.sigmoid(inputs)
+        intersection = (inputs * targets).sum()
+        dice_coeff = (2. * intersection + epsilon) / (inputs.sum() + targets.sum() + epsilon)
+        return 1 - dice_coeff
+
+    def forward(self, inputs, targets):
+        bce = self.bce_loss(inputs, targets)
+        dice = self.dice_loss(inputs, targets)
+        return self.loss_params[0] * bce + self.loss_params[1] * dice
+
+
+class SmoothnessLoss(nn.Module):
+    """
+    Computes the smoothness loss for a sequence of predictions.
+
+    Parameters
+    ----------
+    alpha : float, optional
+        Weight of the smoothness loss component. Default is 10.
+
+    Methods
+    -------
+    forward(predictions)
+        Computes the smoothness loss for a sequence of predictions.
+    """
+
+    def __init__(self, alpha=10):
+        super(SmoothnessLoss, self).__init__()
+        self.alpha = alpha
+
+    def forward(self, predictions):
+        if predictions.dim() < 3:
+            raise ValueError("The input tensor must be 3-dimensional.")
+        diffs = predictions[:, :, 1:] - predictions[:, :, :-1]
+        loss = torch.sum(diffs ** 2) / predictions.size(0)
+        return self.alpha * loss
+
+
+def f1_score(logits, true_labels, threshold=0.5, epsilon=1e-7):
+    """
+    Computes the F1 score for binary classification.
+
+    Parameters
+    ----------
+    logits : torch.Tensor
+        The raw output from the model (before applying sigmoid).
+    true_labels : torch.Tensor
+        The ground truth binary labels.
+    threshold : float, optional
+        The threshold to convert probabilities to binary predictions. Default is 0.5.
+    epsilon : float, optional
+        A small value to avoid division by zero. Default is 1e-7.
+
+    Returns
+    -------
+    float
+        The computed F1 score.
+    """
+    probabilities = torch.sigmoid(logits)
+    predictions = probabilities > threshold
+    predictions = predictions.float()
+    true_labels = true_labels.float()
+    tp = (predictions * true_labels).sum().item()
+    fp = ((1 - true_labels) * predictions).sum().item()
+    fn = (true_labels * (1 - predictions)).sum().item()
+    precision = tp / (tp + fp + epsilon)
+    recall = tp / (tp + fn + epsilon)
+    f1_score = 2 * (precision * recall) / (precision + recall + epsilon)
+    return f1_score

contraction_net/prediction.py

@@ -0,0 +1,41 @@
+import numpy as np
+import torch
+
+from .contraction_net import ContractionNet
+from .utils import get_device
+
+# select device
+device = get_device()
+
+
+def predict_contractions(data, model, network=ContractionNet):
+    """predict contraction intervals time-series with neural network
+
+    Parameters
+    ----------
+    data : ndarray
+        1D array with time-series of contraction
+    model : str
+        trained model weights (.pt file)
+    network : nn.Module
+        Network to predict contractions from time-series
+    standard_normalizer : bool, ndarray
+        If False, each data is normalized by its mean and std. If True, the mean and std from the training data set are
+        applied. If ndarray, the data is normalized with the entered args [[input_mean, input_std], [vel_mean, vel_std]]
+    """
+    # load model state
+    state_dict = torch.load(model, map_location=device)
+    # initiate model
+    model = network(state_dict['n_filter'], in_channels=state_dict['in_channels'],
+                    out_channels=state_dict['out_channels']).to(device)
+    model.load_state_dict(state_dict['state_dict'])
+    # resize data
+    len_data = data.shape[0]
+    contr = np.pad(data, (0, 32 - np.mod(len_data, 32)), mode='reflect')
+    # convert to torch
+    data = torch.from_numpy(contr.astype('float32')).view(1, 1, -1).to(device)
+    # predict data
+    res = model(data)[0][0]
+    # resize data and convert to numpy
+    res = res.detach().cpu().numpy()
+    return res[:, :len_data]