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