celldetective 1.3.7__py3-none-any.whl → 1.3.7.post2__py3-none-any.whl

celldetective/gui/processes/train_segmentation_model.py

@@ -0,0 +1,270 @@
+from multiprocessing import Process
+import time
+import os
+import shutil
+from glob import glob
+import json
+from tqdm import tqdm
+import numpy as np
+import random
+
+from celldetective.utils import load_image_dataset, augmenter, interpolate_nan
+from celldetective.io import normalize_multichannel
+from stardist import fill_label_holes
+from art import tprint
+from distutils.dir_util import copy_tree
+from csbdeep.utils import save_json
+
+
+class TrainSegModelProcess(Process):
+
+	def __init__(self, queue=None, process_args=None, *args, **kwargs):
+		
+		super().__init__(*args, **kwargs)
+		
+		self.queue = queue
+
+		if process_args is not None:
+			for key, value in process_args.items():
+				setattr(self, key, value)
+
+		tprint("Train segmentation")
+		self.read_instructions()
+		self.extract_training_params()
+		self.load_dataset()
+		self.split_test_train()
+
+		self.sum_done = 0
+		self.t0 = time.time()
+
+	def read_instructions(self):
+
+		if os.path.exists(self.instructions):
+			with open(self.instructions, 'r') as f:
+				self.training_instructions = json.load(f)
+		else:
+			print('Training instructions could not be found. Abort.')
+			self.abort_process()
+
+	def run(self):
+
+		if self.model_type=="cellpose":
+			self.train_cellpose_model()
+		elif self.model_type=="stardist":
+			self.train_stardist_model()
+
+		self.queue.put("finished")
+		self.queue.close()
+
+	def train_stardist_model(self):
+
+		from stardist import calculate_extents, gputools_available
+		from stardist.models import Config2D, StarDist2D
+		
+		n_rays = 32
+		print(gputools_available())
+		
+		n_channel = self.X_trn[0].shape[-1]
+		
+		# Predict on subsampled grid for increased efficiency and larger field of view
+		grid = (2,2)
+		conf = Config2D(
+			n_rays       = n_rays,
+			grid         = grid,
+			use_gpu      = self.use_gpu,
+			n_channel_in = n_channel,
+			train_learning_rate = self.learning_rate,
+			train_patch_size = (256,256),
+			train_epochs = self.epochs,
+			train_reduce_lr = {'factor': 0.1, 'patience': 30, 'min_delta': 0},
+			train_batch_size = self.batch_size,
+			train_steps_per_epoch = int(self.augmentation_factor*len(self.X_trn)),
+		)
+		
+		if self.use_gpu:
+			from csbdeep.utils.tf import limit_gpu_memory
+			limit_gpu_memory(None, allow_growth=True)
+
+		if self.pretrained is None:
+			model = StarDist2D(conf, name=self.model_name, basedir=self.target_directory)
+		else:
+			os.rename(self.instructions, os.sep.join([self.target_directory, self.model_name, 'temp.json']))
+			copy_tree(self.pretrained, os.sep.join([self.target_directory, self.model_name]))
+			
+			if os.path.exists(os.sep.join([self.target_directory, self.model_name, 'training_instructions.json'])):
+				os.remove(os.sep.join([self.target_directory, self.model_name, 'training_instructions.json']))
+			if os.path.exists(os.sep.join([self.target_directory, self.model_name, 'config_input.json'])):
+				os.remove(os.sep.join([self.target_directory, self.model_name, 'config_input.json']))
+			if os.path.exists(os.sep.join([self.target_directory, self.model_name, 'logs'+os.sep])):
+				shutil.rmtree(os.sep.join([self.target_directory, self.model_name, 'logs']))
+			os.rename(os.sep.join([self.target_directory, self.model_name, 'temp.json']),os.sep.join([self.target_directory, self.model_name, 'training_instructions.json']))
+
+			#shutil.copytree(pretrained, os.sep.join([target_directory, model_name]))
+			model = StarDist2D(None, name=self.model_name, basedir=self.target_directory)
+			model.config.train_epochs = self.epochs
+			model.config.train_batch_size = min(len(self.X_trn),self.batch_size)
+			model.config.train_learning_rate = self.learning_rate # perf seems bad if lr is changed in transfer
+			model.config.use_gpu = self.use_gpu
+			model.config.train_reduce_lr = {'factor': 0.1, 'patience': 10, 'min_delta': 0}
+			print(f'{model.config=}')
+
+			save_json(vars(model.config), os.sep.join([self.target_directory, self.model_name, 'config.json']))
+
+		median_size = calculate_extents(list(self.Y_trn), np.mean)
+		fov = np.array(model._axes_tile_overlap('YX'))
+		print(f"median object size:      {median_size}")
+		print(f"network field of view :  {fov}")
+		if any(median_size > fov):
+			print("WARNING: median object size larger than field of view of the neural network.")
+
+		if self.augmentation_factor==1.0:
+			model.train(self.X_trn, self.Y_trn, validation_data=(self.X_val,self.Y_val))
+		else:
+			model.train(self.X_trn, self.Y_trn, validation_data=(self.X_val,self.Y_val), augmenter=augmenter)
+		model.optimize_thresholds(self.X_val,self.Y_val)
+
+		config_inputs = {"channels": self.target_channels, 'normalization_percentile': self.normalization_percentile,
+		'normalization_clip': self.normalization_clip, 'normalization_values': self.normalization_values, 
+		'model_type': 'stardist', 'spatial_calibration': self.spatial_calibration, 'dataset': {'train': self.files_train, 'validation': self.files_val}}
+
+		json_input_config = json.dumps(config_inputs, indent=4)
+		with open(os.sep.join([self.target_directory, self.model_name, "config_input.json"]), "w") as outfile:
+			outfile.write(json_input_config)
+
+	def train_cellpose_model(self):
+
+		# do augmentation in place
+		X_aug = []; Y_aug = [];
+		n_val = max(1, int(round(self.augmentation_factor * len(self.X_trn))))
+		indices = random.choices(list(np.arange(len(self.X_trn))), k=n_val)
+		print('Performing image augmentation pre-training...')
+		for i in tqdm(indices):
+			x_aug,y_aug = augmenter(self.X_trn[i], self.Y_trn[i])
+			X_aug.append(x_aug)
+			Y_aug.append(y_aug)
+		
+		# Channel axis in front for cellpose
+		X_aug = [np.moveaxis(x,-1,0) for x in X_aug]
+		self.X_val = [np.moveaxis(x,-1,0) for x in self.X_val]
+		print('number of augmented images: %3d' % len(X_aug))
+
+		from cellpose.models import CellposeModel
+		from cellpose.io import logger_setup
+		import torch
+		
+		if not self.use_gpu:
+			print('Using CPU for training...')
+			device = torch.device("cpu")
+		else:
+			print('Using GPU for training...')
+			
+		logger, log_file = logger_setup()
+		print(f'Pretrained model: ', self.pretrained)
+		if self.pretrained is not None:
+			pretrained_path = os.sep.join([self.pretrained,os.path.split(self.pretrained)[-1]])
+		else:
+			pretrained_path = self.pretrained
+		
+		model = CellposeModel(gpu=self.use_gpu, model_type=None, pretrained_model=pretrained_path, diam_mean=30.0, nchan=X_aug[0].shape[0],) 
+		model.train(train_data=X_aug, train_labels=Y_aug, normalize=False, channels=None, batch_size=self.batch_size,
+					min_train_masks=1,save_path=self.target_directory+os.sep+self.model_name,n_epochs=self.epochs, model_name=self.model_name, learning_rate=self.learning_rate, test_data = self.X_val, test_labels=self.Y_val)
+
+		file_to_move = glob(os.sep.join([self.target_directory, self.model_name, 'models','*']))[0]
+		shutil.move(file_to_move, os.sep.join([self.target_directory, self.model_name,''])+os.path.split(file_to_move)[-1])
+		os.rmdir(os.sep.join([self.target_directory, self.model_name, 'models']))
+
+		diameter = model.diam_labels
+
+		if self.pretrained is not None and os.path.split(self.pretrained)[-1]=='CP_nuclei':
+			standard_diameter = 17.0
+		else:
+			standard_diameter = 30.0
+
+		input_spatial_calibration = self.spatial_calibration #*diameter / standard_diameter
+
+		config_inputs = {"channels": self.target_channels, "diameter": standard_diameter, 'cellprob_threshold': 0., 'flow_threshold': 0.4,
+		'normalization_percentile': self.normalization_percentile, 'normalization_clip': self.normalization_clip,
+		'normalization_values': self.normalization_values, 'model_type': 'cellpose',
+		'spatial_calibration': input_spatial_calibration, 'dataset': {'train': self.files_train, 'validation': self.files_val}}
+		json_input_config = json.dumps(config_inputs, indent=4)
+		with open(os.sep.join([self.target_directory, self.model_name, "config_input.json"]), "w") as outfile:
+			outfile.write(json_input_config)
+
+
+	def split_test_train(self):
+
+		if not len(self.X) > 1:
+			print("Not enough training data")
+			self.abort_process()
+
+		rng = np.random.RandomState()
+		ind = rng.permutation(len(self.X))
+		n_val = max(1, int(round(self.validation_split * len(ind))))
+		ind_train, ind_val = ind[:-n_val], ind[-n_val:]
+		self.X_val, self.Y_val = [self.X[i] for i in ind_val]  , [self.Y[i] for i in ind_val]
+		self.X_trn, self.Y_trn = [self.X[i] for i in ind_train], [self.Y[i] for i in ind_train]
+
+		self.files_train = [self.filenames[i] for i in ind_train]
+		self.files_val = [self.filenames[i] for i in ind_val]
+
+		print('number of images: %3d' % len(self.X))
+		print('- training:       %3d' % len(self.X_trn))
+		print('- validation:     %3d' % len(self.X_val))
+
+	def extract_training_params(self):
+
+		self.model_name = self.training_instructions['model_name']
+		self.target_directory = self.training_instructions['target_directory']
+		self.model_type = self.training_instructions['model_type']
+		self.pretrained = self.training_instructions['pretrained']
+
+		self.datasets = self.training_instructions['ds']
+
+		self.target_channels = self.training_instructions['channel_option']
+		self.normalization_percentile = self.training_instructions['normalization_percentile']
+		self.normalization_clip = self.training_instructions['normalization_clip']
+		self.normalization_values = self.training_instructions['normalization_values']
+		self.spatial_calibration = self.training_instructions['spatial_calibration']
+
+		self.validation_split = self.training_instructions['validation_split']
+		self.augmentation_factor = self.training_instructions['augmentation_factor']
+
+		self.learning_rate = self.training_instructions['learning_rate']
+		self.epochs = self.training_instructions['epochs']
+		self.batch_size = self.training_instructions['batch_size']
+
+	def load_dataset(self):
+
+		print(f'Datasets: {self.datasets}')
+		self.X,self.Y,self.filenames = load_image_dataset(self.datasets, self.target_channels, train_spatial_calibration=self.spatial_calibration,
+								mask_suffix='labelled')
+		print('Dataset loaded...')
+
+		self.values = []
+		self.percentiles = []
+		for k in range(len(self.normalization_percentile)):
+			if self.normalization_percentile[k]:
+				self.percentiles.append(self.normalization_values[k])
+				self.values.append(None)
+			else:
+				self.percentiles.append(None)
+				self.values.append(self.normalization_values[k])
+
+		self.X = [normalize_multichannel(x, **{"percentiles": self.percentiles, 'values': self.values, 'clip': self.normalization_clip}) for x in self.X]
+
+		for k in range(len(self.X)):
+			x = self.X[k].copy()
+			x_interp = np.moveaxis([interpolate_nan(x[:,:,c].copy()) for c in range(x.shape[-1])],0,-1)
+			self.X[k] = x_interp
+
+		self.Y = [fill_label_holes(y) for y in tqdm(self.Y)]
+
+	def end_process(self):
+
+		self.terminate()
+		self.queue.put("finished")
+
+	def abort_process(self):
+		
+		self.terminate()
+		self.queue.put("error")

celldetective/gui/processes/train_signal_model.py

@@ -0,0 +1,108 @@
+from multiprocessing import Process
+import time
+import os
+import json
+from glob import glob
+import numpy as np
+from art import tprint
+from celldetective.signals import SignalDetectionModel
+from celldetective.io import locate_signal_model
+
+
+class TrainSignalModelProcess(Process):
+
+	def __init__(self, queue=None, process_args=None, *args, **kwargs):
+		
+		super().__init__(*args, **kwargs)
+		
+		self.queue = queue
+
+		if process_args is not None:
+			for key, value in process_args.items():
+				setattr(self, key, value)
+
+		tprint("Train segmentation")
+		self.read_instructions()
+		self.extract_training_params()
+
+
+		self.sum_done = 0
+		self.t0 = time.time()
+
+	def read_instructions(self):
+
+		if os.path.exists(self.instructions):
+			with open(self.instructions, 'r') as f:
+				self.training_instructions = json.load(f)
+		else:
+			print('Training instructions could not be found. Abort.')
+			self.abort_process()
+
+		all_classes = []
+		for d in self.training_instructions["ds"]:
+			datasets = glob(d+os.sep+"*.npy")
+			for dd in datasets:
+				data = np.load(dd, allow_pickle=True)
+				classes = np.unique([ddd["class"] for ddd in data])
+				all_classes.extend(classes)
+		all_classes = np.unique(all_classes)
+		n_classes = len(all_classes)
+
+		self.model_params = {k:self.training_instructions[k] for k in ('pretrained', 'model_signal_length', 'channel_option', 'n_channels', 'label') if k in self.training_instructions}
+		self.model_params.update({'n_classes': n_classes})
+		self.train_params = {k:self.training_instructions[k] for k in ('model_name', 'target_directory', 'channel_option','recompile_pretrained', 'test_split', 'augment', 'epochs', 'learning_rate', 'batch_size', 'validation_split','normalization_percentile','normalization_values','normalization_clip') if k in self.training_instructions}
+
+	def neighborhood_postprocessing(self):
+
+		# if neighborhood of interest in training instructions, write it in config!
+		if 'neighborhood_of_interest' in self.training_instructions:
+			if self.training_instructions['neighborhood_of_interest'] is not None:
+				
+				model_path = locate_signal_model(self.training_instructions['model_name'], path=None, pairs=True)
+				complete_path = model_path #+model
+				complete_path = rf"{complete_path}"
+				model_config_path = os.sep.join([complete_path,'config_input.json'])
+				model_config_path = rf"{model_config_path}"
+
+				f = open(model_config_path)
+				config = json.load(f)
+				config.update({'neighborhood_of_interest': self.training_instructions['neighborhood_of_interest'], 'reference_population': self.training_instructions['reference_population'], 'neighbor_population': self.training_instructions['neighbor_population']})
+				json_string = json.dumps(config)
+				with open(model_config_path, 'w') as outfile:
+					outfile.write(json_string)
+
+	def run(self):
+
+		model = SignalDetectionModel(**self.model_params)
+		model.fit_from_directory(self.training_instructions['ds'], **self.train_params)
+		self.neighborhood_postprocessing()
+		self.queue.put("finished")
+		self.queue.close()
+
+
+	def extract_training_params(self):
+
+		self.training_instructions.update({'n_channels': len(self.training_instructions['channel_option'])})
+		if self.training_instructions['augmentation_factor']>1.0:
+			self.training_instructions.update({'augment': True})
+		else:
+			self.training_instructions.update({'augment': False})
+		self.training_instructions.update({'test_split': 0.})
+
+
+	def end_process(self):
+		
+		# self.terminate()
+
+		# if self.model_type=="stardist":
+		# 	from stardist.models import StarDist2D
+		# 	self.model = StarDist2D(None, name=self.model_name, basedir=self.target_directory)
+		# 	self.model.optimize_thresholds(self.X_val,self.Y_val)
+
+		self.terminate()
+		self.queue.put("finished")
+
+	def abort_process(self):
+		
+		self.terminate()
+		self.queue.put("error")

celldetective/io.py

@@ -58,6 +58,7 @@ def extract_experiment_from_well(well_path):
 	>>> well_path = "/path/to/experiment/plate/well"
 	>>> extract_experiment_from_well(well_path)
 	'/path/to/experiment'
+	
 	"""
 
 	if not well_path.endswith(os.sep):
@@ -94,6 +95,7 @@ def extract_well_from_position(pos_path):
 	>>> pos_path = "/path/to/experiment/plate/well/position"
 	>>> extract_well_from_position(pos_path)
 	'/path/to/experiment/plate/well/'
+
 	"""
 
 	if not pos_path.endswith(os.sep):
@@ -129,6 +131,7 @@ def extract_experiment_from_position(pos_path):
 	>>> pos_path = "/path/to/experiment/plate/well/position"
 	>>> extract_experiment_from_position(pos_path)
 	'/path/to/experiment'
+
 	"""
 
 	pos_path = pos_path.replace(os.sep, '/')
@@ -187,6 +190,7 @@ def collect_experiment_metadata(pos_path=None, well_path=None):
 	>>> metadata = collect_experiment_metadata(well_path=well_path)
 	>>> metadata["concentration"]
 	10.0
+
 	"""
 
 	if pos_path is not None:
@@ -289,6 +293,7 @@ def get_config(experiment):
 	>>> config_path = get_config(experiment)
 	>>> print(config_path)
 	'/path/to/experiment/config.ini'
+
 	"""
 
 	if not experiment.endswith(os.sep):
@@ -336,6 +341,7 @@ def get_spatial_calibration(experiment):
 	>>> calibration = get_spatial_calibration(experiment)
 	>>> print(calibration)
 	0.325  # pixels-to-micrometers conversion factor
+
 	"""
 
 	config = get_config(experiment)
@@ -380,6 +386,7 @@ def get_temporal_calibration(experiment):
 	>>> calibration = get_temporal_calibration(experiment)
 	>>> print(calibration)
 	0.5  # frames-to-minutes conversion factor
+
 	"""
 
 	config = get_config(experiment)
@@ -435,6 +442,7 @@ def get_experiment_concentrations(experiment, dtype=str):
 	>>> concentrations = get_experiment_concentrations(experiment, dtype=float)
 	>>> print(concentrations)
 	[0.1, 0.2, 0.5, 1.0]
+
 	"""
 	
 	config = get_config(experiment)
@@ -489,6 +497,7 @@ def get_experiment_cell_types(experiment, dtype=str):
 	>>> cell_types = get_experiment_cell_types(experiment, dtype=str)
 	>>> print(cell_types)
 	['TypeA', 'TypeB', 'TypeC', 'TypeD']
+
 	"""
 
 	config = get_config(experiment)
@@ -540,6 +549,7 @@ def get_experiment_antibodies(experiment, dtype=str):
 
 	>>> get_experiment_antibodies("path/to/experiment2", dtype=int)
 	array([0, 1, 2])
+
 	"""
 
 	config = get_config(experiment)
@@ -594,6 +604,7 @@ def get_experiment_pharmaceutical_agents(experiment, dtype=str):
 	>>> antibodies = get_experiment_antibodies(experiment, dtype=str)
 	>>> print(antibodies)
 	['AntibodyA', 'AntibodyB', 'AntibodyC', 'AntibodyD']
+
 	"""
 
 	config = get_config(experiment)
@@ -702,6 +713,7 @@ def extract_well_name_and_number(well):
 	>>> well_path = "another/path/W1"
 	>>> extract_well_name_and_number(well_path)
 	('W1', 1)
+
 	"""
 
 	split_well_path = well.split(os.sep)
@@ -740,6 +752,7 @@ def extract_position_name(pos):
 	>>> pos_path = "another/path/positionA"
 	>>> extract_position_name(pos_path)
 	'positionA'
+
 	"""
 
 	split_pos_path = pos.split(os.sep)
@@ -890,6 +903,7 @@ def get_position_movie_path(pos, prefix=''):
 	>>> pos_path = "nonexistent/path"
 	>>> get_position_movie_path(pos_path)
 	None
+
 	"""
 	
 
@@ -961,6 +975,7 @@ def load_experiment_tables(experiment, population='targets', well_option='*', po
 	Use pickle files for faster loading:
 
 	>>> df = load_experiment_tables("experiment_01", load_pickle=True)
+
 	"""
 
 	config = get_config(experiment)
@@ -1171,6 +1186,7 @@ def locate_labels(position, population='target', frames=None):
 	Load multiple specific frames:
 
 	>>> labels = locate_labels("/path/to/position", population="target", frames=[0, 1, 2])
+
 	"""
 
 	if not position.endswith(os.sep):
@@ -1241,6 +1257,7 @@ def fix_missing_labels(position, population='target', prefix='Aligned'):
 	-------
 	None
 		The function creates new label files in the corresponding folder for any frames missing label files.
+
 	"""
 
 	if not position.endswith(os.sep):
@@ -1414,6 +1431,7 @@ def auto_load_number_of_frames(stack_path):
 	>>> frames = auto_load_number_of_frames(None)
 	>>> print(frames)
 	None
+
 	"""
 
 	if stack_path is None:
@@ -1510,6 +1528,7 @@ def parse_isotropic_radii(string):
 	- It identifies ranges using square brackets and assumes that ranges are always 
 	  two consecutive values.
 	- Non-integer sections of the string are ignored.
+
 	"""
 
 	sections = re.split(',| ', string)
@@ -1618,6 +1637,7 @@ def interpret_tracking_configuration(config):
 
 	>>> interpret_tracking_configuration(None)
 	'/path/to/default/config.json'
+
 	"""
 
 	if isinstance(config, str):
@@ -1792,6 +1812,7 @@ def locate_signal_model(name, path=None, pairs=False):
 
 	>>> locate_signal_model("remote_model")
 	'path/to/celldetective/models/signal_detection/remote_model/'
+
 	"""
 
 	main_dir = os.sep.join([os.path.split(os.path.dirname(os.path.realpath(__file__)))[0], "celldetective"])
@@ -1859,6 +1880,7 @@ def locate_pair_signal_model(name, path=None):
 
 	>>> locate_pair_signal_model("custom_model", path="/additional/models/")
 	'/additional/models/custom_model/'
+
 	"""
 
 
@@ -1937,6 +1959,7 @@ def relabel_segmentation(labels, df, exclude_nans=True, column_labels={'track':
 	... }
 	>>> new_labels = relabel_segmentation(labels, df, column_labels=column_labels, exclude_nans=True)
 	Done.
+
 	"""
 
 	n_threads = threads
@@ -2037,6 +2060,7 @@ def control_tracks(position, prefix="Aligned", population="target", relabel=True
 	Example
 	-------
 	>>> control_tracks("/path/to/data/position_1", prefix="Aligned", population="target", relabel=True, flush_memory=True, threads=4)
+	
 	"""
 	
 	if not position.endswith(os.sep):
@@ -2089,6 +2113,7 @@ def tracks_to_btrack(df, exclude_nans=False):
 	Example
 	-------
 	>>> data, properties, graph = tracks_to_btrack(df, exclude_nans=True)
+
 	"""
 
 	graph = {}
@@ -2268,7 +2293,7 @@ def view_tracks_in_napari(position, population, stack=None, labels=None, relabel
 				new_cell['TRACK_ID'] = value_under
 				df = pd.concat([df, new_cell], ignore_index=True)
 
-			relabel = np.amax(df['TRACK_ID'].unique()) + 1
+			relabel = np.amax(viewer.layers['segmentation'].data) + 1
 			for f in viewer.layers['segmentation'].data[int(frame):]:
 				if target_track_id!=0:
 					f[np.where(f==target_track_id)] = relabel