  1# -*- coding: utf-8 -*-
     2"""
     3Created on Mon Jul 31 16:46:21 2023
     4
     5@author: Bart Steemans. Govers Lab.
     6"""
     7import pickle
     8import pandas as pd
     9import numpy as np
    10
    11
    12class Curation:
    13
    14    def __init__(self, svm_df):
    15
    16        self.svm_df = svm_df
    17        self.features = None
    18        self.curated_df = None
    19        self.svm_model = None
    20
    21    def compiled_curation(self, path_to_model, cols):
    22        """
    23        Perform compiled cell curation using a trained SVM model.
    24
    25        This method curates the dataset by making predictions using a trained Support Vector Machine (SVM) model.
    26        Cells that do not meet the criteria specified by the model are flagged for potential removal.
    27
    28        Parameters
    29        ----------
    30        path_to_model : str
    31            The path to the trained SVM model file.
    32
    33        cols : int
    34            The number of columns to exclude in the dataset during curation.
    35
    36        Returns
    37        -------
    38        pd.DataFrame
    39            A DataFrame containing the curated dataset.
    40
    41        """
    42        self.load_model(path_to_model)
    43        self.prepare_dataframe(cols)
    44        self.curated_df = self.make_predictions()
    45        return self.curated_df
    46
    47    def load_model(self, path_to_model):
    48        """
    49        Load a trained SVM model from a file.
    50
    51        This method loads a trained Support Vector Machine (SVM) model from a file located at the given path.
    52
    53        Parameters
    54        ----------
    55        path_to_model : str
    56            The path to the trained SVM model file.
    57
    58        """
    59        with open(path_to_model, "rb") as f:
    60            self.svm_model = pickle.load(f)
    61
    62    def prepare_dataframe(self, cols):
    63        """
    64        Prepare the feature dataset for SVM curation.
    65    
    66        This method prepares the feature dataset for cell curation using a Support Vector Machine (SVM) model.
    67        It extracts the relevant feature columns while ensuring that `cell_id` and `frame` are preserved.
    68    
    69        Parameters
    70        ----------
    71        cols : int
    72            Retained for backward compatibility; feature columns are selected
    73            by name and this argument is currently unused.
    74        """
    75
    76        # Identify non-feature columns (but keep cell_id & frame for tracking)
    77        exclude_columns = ["image_name", "contour", "label"]  # Keep cell_id & frame
    78        cols_to_include = [col for col in self.svm_df.columns if col not in exclude_columns]
    79        cols_to_include.sort()
    80
    81        # Preserve cell_id and frame separately for re-mapping
    82        self.cell_id_frame = self.svm_df[["cell_id", "frame"]].copy()
    83
    84        # Select only feature columns for training
    85        self.features = self.svm_df[cols_to_include].drop(columns=["cell_id", "frame"])
    86
    87        # Convert numeric columns to float
    88        self.features[self.features.select_dtypes(include=[np.number]).columns] = (
    89            self.features.select_dtypes(include=[np.number]).astype(float)
    90        )
    91
    92        # Check for infinite and NaN values
    93        mask = ~(
    94            np.isinf(self.features).any(axis=1) | np.isnan(self.features).any(axis=1)
    95        )
    96
    97        # Store discarded rows separately
    98        self.svm_df_nan = self.svm_df[~mask].copy()
    99        self.svm_df_nan["label"] = 0  # Assign label 0 to discarded cells
   100
   101        # Keep only valid rows for training
   102        self.features = self.features[mask]
   103        self.cell_id_frame = self.cell_id_frame[mask]  # Ensure cell_id & frame match
   104        self.svm_df = self.svm_df[mask]
   105
   106        # Print the number of removed rows
   107        num_removed_rows = len(mask) - mask.sum()
   108        print(f"\nNumber of removed rows: {num_removed_rows}")
   109
   110    def get_label_proportions(self):
   111        """
   112        Calculate the proportions of labels in the curated dataset.
   113
   114        This method calculates and returns the proportions of label 1 and label 0 in the curated dataset.
   115        It counts the occurrences of each label and divides them by the total number of samples in the dataset to compute the proportions.
   116
   117        Returns
   118        -------
   119        float
   120            The proportion of label 1 in the curated dataset.
   121        float
   122            The proportion of label 0 in the curated dataset.
   123
   124        """
   125        label_counts = self.curated_df["label"].value_counts()
   126        total_samples = len(self.curated_df)
   127
   128        proportion_label_1 = label_counts.get(1, 0) / total_samples
   129        proportion_label_0 = label_counts.get(0, 0) / total_samples
   130
   131        return proportion_label_1, proportion_label_0
   132
   133    def make_predictions(self):
   134        """
   135        Make predictions using the SVM model and reattach `cell_id` and `frame`.
   136    
   137        Returns
   138        -------
   139        pd.DataFrame
   140            A DataFrame containing the predicted labels with corresponding `cell_id` and `frame`.
   141        """
   142
   143        # Use the schema stored with the trained model. An unnamed model can
   144        # only check width and could silently interpret reordered columns.
   145        expected_columns = getattr(self.svm_model, "feature_names_in_", None)
   146        if expected_columns is None:
   147            raise ValueError(
   148                "The SVM model has no feature_names_in_ schema. Train it on a "
   149                "named pandas DataFrame before using it for curation."
   150            )
   151        expected_columns = list(expected_columns)
   152        actual_columns = list(self.features.columns)
   153        if actual_columns != expected_columns:
   154            raise ValueError(
   155                "SVM feature columns do not match the model's training schema. "
   156                f"Expected {expected_columns}; got {actual_columns}."
   157            )
   158
   159        # Predict labels
   160        predictions = self.svm_model.predict(self.features)
   161
   162        # Create curated DataFrame
   163        self.curated_df = self.svm_df.copy()
   164        self.curated_df["label"] = predictions
   165
   166        # Reattach `cell_id` and `frame`
   167        self.curated_df[["cell_id", "frame"]] = self.cell_id_frame.values
   168
   169        # Add back NaN/Inf-removed rows
   170        self.curated_df = pd.concat([self.curated_df, self.svm_df_nan], ignore_index=True)
   171
   172        return self.curated_df
   173
   174    def get_control(self, num_pos=5, num_neg=5):
   175        """
   176        Get control cell IDs for positive and negative classes.
   177
   178        This method retrieves cell IDs for positive and negative classes based on the curated dataset.
   179        It randomly selects a specified number of positive and negative cell IDs and returns them as lists, along with the corresponding 'frame' values.
   180
   181        Parameters
   182        ----------
   183        num_pos : int, optional
   184            The number of positive class cell IDs to retrieve. The default is 5.
   185        num_neg : int, optional
   186            The number of negative class cell IDs to retrieve. The default is 5.
   187
   188        Returns
   189        -------
   190        Tuple
   191            Two lists containing tuples of ('frame', 'cell_id') for positive and negative class control cell IDs.
   192            If there are not enough samples for one class, the corresponding list will be empty.
   193        """
   194
   195        positive_indices = self.curated_df[self.curated_df["label"] == 1].index.tolist()
   196        negative_indices = self.curated_df[self.curated_df["label"] == 0].index.tolist()
   197
   198        # Check if there are not enough positive or negative samples
   199        if len(positive_indices) < num_pos:
   200            print("No positive labels")
   201            return None, None  # Return None for positive and available negative indices
   202        elif len(negative_indices) < num_neg:
   203            print("No negative labels")
   204            return None, None  # Return available positive indices and None for negative
   205
   206        positive_indices_random = np.random.choice(
   207            positive_indices, size=num_pos, replace=False
   208        )
   209        negative_indices_random = np.random.choice(
   210            negative_indices, size=num_neg, replace=False
   211        )
   212
   213        # Retrieve cell_id and frame for selected indices
   214        positive_cell_frames = [
   215            (self.curated_df.loc[idx, "cell_id"], self.curated_df.loc[idx, "frame"])
   216            for idx in positive_indices_random
   217        ]
   218        negative_cell_frames = [
   219            (self.curated_df.loc[idx, "cell_id"], self.curated_df.loc[idx, "frame"])
   220            for idx in negative_indices_random
   221        ]
   222
   223        return positive_cell_frames, negative_cell_frames