celldetective 1.1.1.post3__py3-none-any.whl → 1.2.0__py3-none-any.whl

celldetective/gui/retrain_segmentation_model_options.py

@@ -423,7 +423,7 @@ class ConfigSegmentationModelTraining(QMainWindow, Styles):
 						)
 		if self.dataset_folder is not None:
 
-			subfiles = glob(self.dataset_folder+"/*.tif")
+			subfiles = glob(self.dataset_folder+os.sep+"*.tif")
 			if len(subfiles)>0:
 				print(f'found {len(subfiles)} files in folder')
 				self.data_folder_label.setText(self.dataset_folder[:16]+'...')
@@ -459,14 +459,26 @@ class ConfigSegmentationModelTraining(QMainWindow, Styles):
 		self.data_folder_label.setToolTip('')
 		self.cancel_dataset.setVisible(False)
 
+	def load_stardist_train_config(self):
+		
+		config = os.sep.join([self.pretrained_model,"config.json"])
+		if os.path.exists(config):
+			with open(config, 'r') as f:
+				config = json.load(f)
+				if 'train_batch_size' in config:
+					bs = config['train_batch_size']
+					self.bs_le.setText(str(bs).replace('.',','))
+				if 'train_learning_rate' in config:
+					lr = config['train_learning_rate']
+					self.lr_le.setText(str(lr).replace('.',','))
 
 	def load_pretrained_config(self):
 
 		f = open(os.sep.join([self.pretrained_model,"config_input.json"]))
 		data = json.load(f)
 		channels = data["channels"]
-		self.seg_folder = self.pretrained_model.split('/')[-2]
-		self.model_name = self.pretrained_model.split('/')[-1]
+		self.seg_folder = self.pretrained_model.split(os.sep)[-2]
+		self.model_name = self.pretrained_model.split(os.sep)[-1]
 		if self.model_name.startswith('CP') and self.seg_folder=='segmentation_generic':
 			channels = ['brightfield_channel', 'live_nuclei_channel']
 			if self.model_name=="CP_nuclei":
@@ -484,6 +496,7 @@ class ConfigSegmentationModelTraining(QMainWindow, Styles):
 		if model_type=='stardist':
 			self.stardist_model.setChecked(True)
 			self.cellpose_model.setChecked(False)
+			self.load_stardist_train_config()
 		else:
 			self.stardist_model.setChecked(False)
 			self.cellpose_model.setChecked(True)			
@@ -593,7 +606,7 @@ class ConfigSegmentationModelTraining(QMainWindow, Styles):
 
 		print(training_instructions)
 
-		model_folder = '/'.join([self.software_models_dir,model_name, ''])
+		model_folder = os.sep.join([self.software_models_dir,model_name, ''])
 		print(model_folder)
 		if not os.path.exists(model_folder):
 			os.mkdir(model_folder)

celldetective/gui/retrain_signal_model_options.py

@@ -3,11 +3,11 @@ from PyQt5.QtCore import Qt, QSize
 from PyQt5.QtGui import QDoubleValidator, QIntValidator, QIcon
 from celldetective.gui.gui_utils import center_window, FeatureChoice, ListWidget, QHSeperationLine, FigureCanvas, GeometryChoice, OperationChoice
 from celldetective.gui.layouts import ChannelNormGenerator
-from superqt import QLabeledDoubleRangeSlider, QLabeledDoubleSlider,QLabeledSlider
+from superqt import QLabeledDoubleRangeSlider, QLabeledDoubleSlider, QLabeledSlider, QSearchableComboBox
 from superqt.fonticon import icon
 from fonticon_mdi6 import MDI6
 from celldetective.utils import extract_experiment_channels, get_software_location
-from celldetective.io import interpret_tracking_configuration, load_frames, locate_signal_dataset, get_signal_datasets_list
+from celldetective.io import interpret_tracking_configuration, load_frames, locate_signal_dataset, get_signal_datasets_list, load_experiment_tables
 from celldetective.measure import compute_haralick_features, contour_of_instance_segmentation
 from celldetective.signals import train_signal_model
 import numpy as np
@@ -24,6 +24,7 @@ from datetime import datetime
 import pandas as pd
 from functools import partial
 from celldetective.gui import Styles
+from pandas.api.types import is_numeric_dtype
 
 class ConfigSignalModelTraining(QMainWindow, Styles):
 	
@@ -32,7 +33,7 @@ class ConfigSignalModelTraining(QMainWindow, Styles):
 
 	"""
 
-	def __init__(self, parent_window=None):
+	def __init__(self, parent_window=None, signal_mode='single-cells'):
 		
 		super().__init__()
 		self.parent_window = parent_window
@@ -43,7 +44,16 @@ class ConfigSignalModelTraining(QMainWindow, Styles):
 		self.soft_path = get_software_location()
 		self.pretrained_model = None 
 		self.dataset_folder = None
-		self.signal_models_dir = self.soft_path+os.sep+os.sep.join(['celldetective','models','signal_detection'])
+		self.current_neighborhood = None
+		self.reference_population = None
+		self.neighbor_population = None
+		self.signal_mode = signal_mode
+
+		if self.signal_mode=='single-cells':
+			self.signal_models_dir = self.soft_path+os.sep+os.sep.join(['celldetective','models','signal_detection'])
+		elif self.signal_mode=='pairs':
+			self.signal_models_dir = self.soft_path+os.sep+os.sep.join(['celldetective','models','pair_signal_detection'])
+			self.mode = 'pairs'
 
 		self.onlyFloat = QDoubleValidator()
 		self.onlyInt = QIntValidator()
@@ -272,6 +282,14 @@ class ConfigSignalModelTraining(QMainWindow, Styles):
 		modelname_layout.addWidget(self.modelname_le, 70)
 		layout.addLayout(modelname_layout)
 
+		if self.signal_mode=='pairs':
+			neighborhood_layout = QHBoxLayout()
+			neighborhood_layout.addWidget(QLabel('neighborhood of interest: '), 30)
+			self.neighborhood_choice_cb = QSearchableComboBox()
+			self.fill_available_neighborhoods()
+			neighborhood_layout.addWidget(self.neighborhood_choice_cb, 70)
+			layout.addLayout(neighborhood_layout)			
+
 		classname_layout = QHBoxLayout()
 		classname_layout.addWidget(QLabel('event name: '), 30)
 		self.class_name_le = QLineEdit()
@@ -311,6 +329,10 @@ class ConfigSignalModelTraining(QMainWindow, Styles):
 		self.ch_norm = ChannelNormGenerator(self, mode='signals')
 		layout.addLayout(self.ch_norm)
 
+		if self.signal_mode=='pairs':
+			self.neighborhood_choice_cb.currentIndexChanged.connect(self.neighborhood_changed)
+			self.neighborhood_changed()
+
 		model_length_layout = QHBoxLayout()
 		model_length_layout.addWidget(QLabel('Max signal length: '), 30)
 		self.model_length_slider = QLabeledSlider()
@@ -323,6 +345,84 @@ class ConfigSignalModelTraining(QMainWindow, Styles):
 		model_length_layout.addWidget(self.model_length_slider, 70)
 		layout.addLayout(model_length_layout)
 
+	def neighborhood_changed(self):
+
+		neigh = self.neighborhood_choice_cb.currentText()
+		self.current_neighborhood = neigh.replace('target_ref_','').replace('effector_ref_','')
+		self.reference_population = ['targets' if 'target' in neigh else 'effectors'][0]
+		if 'target' in neigh:
+			if 'self' in neigh:
+				self.neighbor_population = 'targets'
+			else:
+				self.neighbor_population = 'effectors'
+		else:
+			if 'self' in neigh:
+				self.neighbor_population = 'effectors'
+			else:
+				self.neighbor_population = 'targets'
+		
+		print(f'Current neighborhood: {self.current_neighborhood}')
+		print(f'New reference population: {self.reference_population}')
+		print(f'New neighbor population: {self.neighbor_population}')
+
+		# reload reference signals / neighbor signals / pair signals
+		# fill the channel cbs
+		self.df_reference = self.dataframes[self.reference_population]
+		self.df_neighbor = self.dataframes[self.neighbor_population]
+		self.df_pairs = load_experiment_tables(self.parent_window.exp_dir, population='pairs', load_pickle=False)
+
+		self.df_reference = self.df_reference.rename(columns=lambda x: 'reference_' + x)
+		num_cols_reference = [c for c in list(self.df_reference.columns) if is_numeric_dtype(self.df_reference[c])]
+		self.df_neighbor = self.df_neighbor.rename(columns=lambda x: 'neighbor_' + x)
+		num_cols_neighbor = [c for c in list(self.df_neighbor.columns) if is_numeric_dtype(self.df_neighbor[c])]
+		self.df_pairs = self.df_pairs.rename(columns=lambda x: 'pair_' + x)
+		num_cols_pairs = [c for c in list(self.df_pairs.columns) if is_numeric_dtype(self.df_pairs[c])]
+
+		self.signals = ['--'] + num_cols_pairs + num_cols_reference + num_cols_neighbor
+
+		for cb in self.ch_norm.channel_cbs:
+			# try:
+			# 	cb.disconnect()
+			# except:
+			# 	pass
+			cb.clear()
+			cb.addItems(self.signals)
+
+	def fill_available_neighborhoods(self):
+		
+		df_targets = load_experiment_tables(self.parent_window.exp_dir, population='targets', load_pickle=True)
+		df_effectors = load_experiment_tables(self.parent_window.exp_dir, population='effectors', load_pickle=True)
+
+		self.dataframes = {
+			'targets': df_targets,
+			'effectors': df_effectors,
+		}
+
+		self.neighborhood_cols = []
+		self.reference_populations = []
+		self.neighbor_populations = []
+		if df_targets is not None:
+			self.neighborhood_cols.extend(['target_ref_'+c for c in list(df_targets.columns) if c.startswith('neighborhood')])
+			self.reference_populations.extend(['targets' for c in list(df_targets.columns) if c.startswith('neighborhood')])
+			for c in list(df_targets.columns):
+				if c.startswith('neighborhood') and '_2_' in c:
+					self.neighbor_populations.append('effectors')
+				elif c.startswith('neighborhood') and 'self' in c:
+					self.neighbor_populations.append('targets')
+		
+		if df_effectors is not None:
+			self.neighborhood_cols.extend(['effector_ref_'+c for c in list(df_effectors.columns) if c.startswith('neighborhood')])
+			self.reference_populations.extend(['effectors' for c in list(df_effectors.columns) if c.startswith('neighborhood')])
+			for c in list(df_effectors.columns):
+				if c.startswith('neighborhood') and '_2_' in c:
+					self.neighbor_populations.append('targets')
+				elif c.startswith('neighborhood') and 'self' in c:
+					self.neighbor_populations.append('effectors')
+
+		print(f"The following neighborhoods were detected: {self.neighborhood_cols=} {self.reference_populations=} {self.neighbor_populations=}")
+
+		self.neighborhood_choice_cb.addItems(self.neighborhood_cols)
+
 	def showDialog_pretrained(self):
 
 		self.pretrained_model = QFileDialog.getExistingDirectory(
@@ -481,9 +581,9 @@ class ConfigSignalModelTraining(QMainWindow, Styles):
 		training_instructions = {'model_name': model_name,'pretrained': pretrained_model, 'channel_option': channels, 'normalization_percentile': normalization_mode,
 		'normalization_clip': clip_values,'normalization_values': norm_values, 'model_signal_length': signal_length,
 		'recompile_pretrained': recompile_op, 'ds': data_folders, 'augmentation_factor': aug_factor, 'validation_split': val_split,
-		'learning_rate': lr, 'batch_size': bs, 'epochs': epochs, 'label': self.class_name_le.text()}
+		'learning_rate': lr, 'batch_size': bs, 'epochs': epochs, 'label': self.class_name_le.text(), 'neighborhood_of_interest': self.current_neighborhood, 'reference_population': self.reference_population, 'neighbor_population': self.neighbor_population}
 
-		model_folder = self.signal_models_dir + model_name + os.sep
+		model_folder = self.signal_models_dir +os.sep+ model_name + os.sep
 		if not os.path.exists(model_folder):
 			os.mkdir(model_folder)
 

celldetective/gui/signal_annotator.py

@@ -27,6 +27,7 @@ from matplotlib.cm import tab10
 import pandas as pd
 from sklearn.preprocessing import MinMaxScaler
 from celldetective.gui import Styles
+from celldetective.measure import contour_of_instance_segmentation
 
 class SignalAnnotator(QMainWindow, Styles):
 	"""
@@ -56,6 +57,7 @@ class SignalAnnotator(QMainWindow, Styles):
 
 		self.screen_height = self.parent_window.parent_window.parent_window.screen_height
 		self.screen_width = self.parent_window.parent_window.parent_window.screen_width
+		self.value_magnitude = 1
 
 		# default params
 		self.class_name = 'class'
@@ -76,9 +78,9 @@ class SignalAnnotator(QMainWindow, Styles):
 
 		self.populate_widget()
 
-		self.setMinimumWidth(int(0.8 * self.screen_width))
+		#self.setMinimumWidth(int(0.8 * self.screen_width))
 		# self.setMaximumHeight(int(0.8*self.screen_height))
-		self.setMinimumHeight(int(0.8 * self.screen_height))
+		#self.setMinimumHeight(int(0.8 * self.screen_height))
 		# self.setMaximumHeight(int(0.8*self.screen_height))
 
 		self.setAttribute(Qt.WA_DeleteOnClose)
@@ -119,7 +121,6 @@ class SignalAnnotator(QMainWindow, Styles):
 
 		self.class_choice_cb.addItems(self.class_cols)
 		self.class_choice_cb.currentIndexChanged.connect(self.compute_status_and_colors)
-		self.class_choice_cb.setCurrentIndex(0)
 
 		class_hbox.addWidget(self.class_choice_cb, 70)
 
@@ -346,6 +347,8 @@ class SignalAnnotator(QMainWindow, Styles):
 		main_layout.addLayout(self.right_panel, 65)
 		self.button_widget.adjustSize()
 
+		self.compute_status_and_colors(0)
+
 		self.setCentralWidget(self.button_widget)
 		self.show()
 
@@ -462,6 +465,7 @@ class SignalAnnotator(QMainWindow, Styles):
 		self.newClassWidget.close()
 
 	def compute_status_and_colors(self, i):
+
 		self.class_name = self.class_choice_cb.currentText()
 		self.expected_status = 'status'
 		suffix = self.class_name.replace('class', '').replace('_', '', 1)
@@ -474,11 +478,15 @@ class SignalAnnotator(QMainWindow, Styles):
 		self.status_name = self.expected_status
 
 		print('selection and expected names: ', self.class_name, self.expected_time, self.expected_status)
+		cols = list(self.df_tracks.columns)
 
-		if self.time_name in self.df_tracks.columns and self.class_name in self.df_tracks.columns and not self.status_name in self.df_tracks.columns:
+		if self.time_name in cols and self.class_name in cols and not self.status_name in cols:
 			# only create the status column if it does not exist to not erase static classification results
 			self.make_status_column()
-		elif self.time_name in self.df_tracks.columns and self.class_name in self.df_tracks.columns:
+		elif self.time_name in cols and self.class_name in cols and self.df_tracks[self.status_name].isnull().all():
+			print('this is the case!', )
+			self.make_status_column()
+		elif self.time_name in cols and self.class_name in cols:
 			# all good, do nothing
 			pass
 		else:
@@ -819,6 +827,8 @@ class SignalAnnotator(QMainWindow, Styles):
 
 	def plot_signals(self):
 
+		range_values = []
+
 		try:
 			yvalues = []
 			for i in range(len(self.signal_choice_cb)):
@@ -834,6 +844,8 @@ class SignalAnnotator(QMainWindow, Styles):
 					xdata = self.df_tracks.loc[self.df_tracks['TRACK_ID'] == self.track_of_interest, 'FRAME'].to_numpy()
 					ydata = self.df_tracks.loc[
 						self.df_tracks['TRACK_ID'] == self.track_of_interest, signal_choice].to_numpy()
+					
+					range_values.extend(ydata)
 
 					xdata = xdata[ydata == ydata]  # remove nan
 					ydata = ydata[ydata == ydata]
@@ -855,7 +867,21 @@ class SignalAnnotator(QMainWindow, Styles):
 			self.cell_ax.legend()
 			self.cell_fcanvas.canvas.draw()
 		except Exception as e:
-			print(f"{e=}")
+			print(f"Plot signals: {e=}")
+		
+		if len(range_values)>0:
+			range_values = np.array(range_values)
+			if len(range_values[range_values==range_values])>0:
+				if len(range_values[range_values>0])>0:
+					self.value_magnitude = np.nanpercentile(range_values, 1)
+				else:
+					self.value_magnitude = 1
+				self.non_log_ymin = 0.98*np.nanmin(range_values)
+				self.non_log_ymax = np.nanmax(range_values)*1.02
+				if self.cell_ax.get_yscale()=='linear':
+					self.cell_ax.set_ylim(self.non_log_ymin, self.non_log_ymax)
+				else:
+					self.cell_ax.set_ylim(self.value_magnitude, self.non_log_ymax)					
 
 	def extract_scatter_from_trajectories(self):
 
@@ -1109,7 +1135,7 @@ class SignalAnnotator(QMainWindow, Styles):
 			if len(min_values) > 0:
 				self.cell_ax.set_ylim(np.amin(min_values), np.amax(max_values))
 		except Exception as e:
-			print(e)
+			print('Ylim error:',e)
 
 	def draw_frame(self, framedata):
 
@@ -1287,21 +1313,24 @@ class SignalAnnotator(QMainWindow, Styles):
 		"""
 
 		try:
-			if self.cell_ax.get_yscale() == 'linear':
+			if self.cell_ax.get_yscale()=='linear':
+				ymin,ymax = self.cell_ax.get_ylim()
 				self.cell_ax.set_yscale('log')
-				self.log_btn.setIcon(icon(MDI6.math_log, color="#1565c0"))
+				self.log_btn.setIcon(icon(MDI6.math_log,color="#1565c0"))
+				self.cell_ax.set_ylim(self.value_magnitude, ymax)
 			else:
 				self.cell_ax.set_yscale('linear')
-				self.log_btn.setIcon(icon(MDI6.math_log, color="black"))
+				self.log_btn.setIcon(icon(MDI6.math_log,color="black"))
 		except Exception as e:
 			print(e)
 
-		# self.cell_ax.autoscale()
+		#self.cell_ax.autoscale()
 		self.cell_fcanvas.canvas.draw_idle()
 
-
 class MeasureAnnotator(SignalAnnotator):
+
 	def __init__(self, parent_window=None):
+
 		QMainWindow.__init__(self)
 		self.parent_window = parent_window
 		self.setWindowTitle("Signal annotator")
@@ -1315,11 +1344,11 @@ class MeasureAnnotator(SignalAnnotator):
 		self.int_validator = QIntValidator()
 		self.current_alpha=0.5
 		if self.mode == "targets":
-			self.instructions_path = self.exp_dir + "configs/signal_annotator_config_targets.json"
-			self.trajectories_path = self.pos + 'output/tables/trajectories_targets.csv'
+			self.instructions_path = self.exp_dir + os.sep.join(['configs','signal_annotator_config_targets.json'])
+			self.trajectories_path = self.pos + os.sep.join(['output','tables','trajectories_targets.csv'])
 		elif self.mode == "effectors":
-			self.instructions_path = self.exp_dir + "configs/signal_annotator_config_effectors.json"
-			self.trajectories_path = self.pos + 'output/tables/trajectories_effectors.csv'
+			self.instructions_path = self.exp_dir + os.sep.join(['configs','signal_annotator_config_effectors.json'])
+			self.trajectories_path = self.pos + os.sep.join(['output','tables','trajectories_effectors.csv'])
 
 		self.screen_height = self.parent_window.parent_window.parent_window.screen_height
 		self.screen_width = self.parent_window.parent_window.parent_window.screen_width
@@ -1330,8 +1359,13 @@ class MeasureAnnotator(SignalAnnotator):
 		center_window(self)
 
 		self.locate_stack()
+
 		data, properties, graph, labels, _ = load_napari_data(self.pos, prefix=None, population=self.mode,return_stack=False)
-		self.labels = relabel_segmentation(labels,data,properties)
+		if data is not None:
+			self.labels = relabel_segmentation(labels,data,properties)
+		else:
+			self.labels = labels
+
 		self.current_channel = 0
 
 		self.locate_tracks()
@@ -1577,8 +1611,10 @@ class MeasureAnnotator(SignalAnnotator):
 		self.class_choice_cb = QComboBox()
 
 		cols = np.array(self.df_tracks.columns)
-		self.class_cols = np.array([c.startswith('group') for c in list(self.df_tracks.columns)])
+		self.class_cols = np.array([c.startswith('group') or c.startswith('status') for c in list(self.df_tracks.columns)])
 		self.class_cols = list(cols[self.class_cols])
+		print(self.class_cols)
+
 		try:
 			self.class_cols.remove('group_id')
 		except Exception:
@@ -1711,6 +1747,14 @@ class MeasureAnnotator(SignalAnnotator):
 		btn_hbox.addWidget(self.save_btn, 90)
 		self.left_panel.addLayout(btn_hbox)
 
+		self.export_btn = QPushButton('')
+		self.export_btn.setStyleSheet(self.button_select_all)
+		self.export_btn.clicked.connect(self.export_measurements)
+		self.export_btn.setIcon(icon(MDI6.export, color="black"))
+		self.export_btn.setIconSize(QSize(25, 25))
+		btn_hbox.addWidget(self.export_btn, 10)
+		self.left_panel.addLayout(btn_hbox)
+
 		# Animation
 		animation_buttons_box = QHBoxLayout()
 
@@ -1810,6 +1854,28 @@ class MeasureAnnotator(SignalAnnotator):
 		# del self.img
 		gc.collect()
 
+
+	def export_measurements(self):
+
+		auto_dataset_name = self.pos.split(os.sep)[-4] + '_' + self.pos.split(os.sep)[-2] + f'_{str(self.current_frame).zfill(3)}' + f'_{self.status_name}.npy'
+
+		if self.normalized_signals:
+			self.normalize_features_btn.click()
+
+		subdf = self.df_tracks.loc[self.df_tracks['FRAME']==self.current_frame,:]
+		subdf['class'] = subdf[self.status_name]
+		dico = subdf.to_dict('records')
+
+		pathsave = QFileDialog.getSaveFileName(self, "Select file name", self.exp_dir + auto_dataset_name, ".npy")[0]
+		if pathsave != '':
+			if not pathsave.endswith(".npy"):
+				pathsave += ".npy"
+			try:
+				np.save(pathsave, dico)
+				print(f'File successfully written in {pathsave}.')
+			except Exception as e:
+				print(f"Error {e}...")
+
 	def set_next_frame(self):
 
 		self.current_frame = self.current_frame + 1
@@ -1887,13 +1953,17 @@ class MeasureAnnotator(SignalAnnotator):
 
 	def give_cell_information(self):
 
-		cell_selected = f"cell: {self.track_of_interest}\n"
-		if 'TRACK_ID' in self.df_tracks.columns:
-			cell_status = f"phenotype: {self.df_tracks.loc[self.df_tracks['TRACK_ID'] == self.track_of_interest, self.status_name].to_numpy()[0]}\n"
-		else:
-			cell_status = f"phenotype: {self.df_tracks.loc[self.df_tracks['ID'] == self.track_of_interest, self.status_name].to_numpy()[0]}\n"
-		self.cell_info.setText(cell_selected + cell_status)
-
+		try:
+			cell_selected = f"cell: {self.track_of_interest}\n"
+			if 'TRACK_ID' in self.df_tracks.columns:
+				cell_status = f"phenotype: {self.df_tracks.loc[self.df_tracks['TRACK_ID'] == self.track_of_interest, self.status_name].to_numpy()[0]}\n"
+			else:
+				cell_status = f"phenotype: {self.df_tracks.loc[self.df_tracks['ID'] == self.track_of_interest, self.status_name].to_numpy()[0]}\n"
+			self.cell_info.setText(cell_selected + cell_status)
+		except Exception as e:
+			print('Cell info:',e)
+			print(self.track_of_interest, self.status_name)
+			
 	def create_new_event_class(self):
 
 		# display qwidget to name the event
@@ -1978,8 +2048,14 @@ class MeasureAnnotator(SignalAnnotator):
 		self.frame_lbl.setText(f'position: {self.framedata}')
 		self.im.set_array(self.img)
 		self.status_scatter.set_offsets(self.positions[self.framedata])
-		self.status_scatter.set_edgecolors(self.colors[self.framedata][:, 0])
+		try:
+			self.status_scatter.set_edgecolors(self.colors[self.framedata][:, 0])
+		except Exception as e:
+			print('L1993: ',e)
+
 		self.current_label = self.labels[self.current_frame]
+		self.current_label = contour_of_instance_segmentation(self.current_label, 5)
+
 		self.im_mask.remove()
 		self.im_mask = self.ax.imshow(np.ma.masked_where(self.current_label == 0, self.current_label),
 											   cmap='viridis', interpolation='none',alpha=self.current_alpha,vmin=0,vmax=np.nanmax(self.labels.flatten()))
@@ -2108,9 +2184,9 @@ class MeasureAnnotator(SignalAnnotator):
 
 			self.extract_scatter_from_trajectories()
 			if 'TRACK_ID' in self.df_tracks.columns:
-				self.track_of_interest = self.df_tracks['TRACK_ID'].min()
+				self.track_of_interest = self.df_tracks.dropna(subset='TRACK_ID')['TRACK_ID'].min()
 			else:
-				self.track_of_interest = self.df_tracks['ID'].min()
+				self.track_of_interest = self.df_tracks.dropna(subset='ID')['ID'].min()
 
 			self.loc_t = []
 			self.loc_idx = []
@@ -2172,9 +2248,13 @@ class MeasureAnnotator(SignalAnnotator):
 		"""
 		self.current_frame = self.frame_slider.value()
 		self.reload_frame()
-		if 'ID' in self.df_tracks.columns:
+		if 'TRACK_ID' in list(self.df_tracks.columns):
+			pass
+		elif 'ID' in list(self.df_tracks.columns):
+			print('ID in cols... change class of interest... ')
 			self.track_of_interest = self.df_tracks[self.df_tracks['FRAME'] == self.current_frame]['ID'].min()
 			self.modify()
+
 		self.draw_frame(self.current_frame)
 		self.fcanvas.canvas.draw()
 		self.plot_signals()
@@ -2302,7 +2382,7 @@ class MeasureAnnotator(SignalAnnotator):
 		try:
 			self.selection.pop(0)
 		except Exception as e:
-			print(e)
+			print('Cancel selection: ',e)
 
 		try:
 			for k, (t, idx) in enumerate(zip(self.loc_t, self.loc_idx)):