celldetective 1.1.1.post1__py3-none-any.whl → 1.1.1.post4__py3-none-any.whl

celldetective/io.py

@@ -26,6 +26,7 @@ from stardist import fill_label_holes
 from celldetective.utils import interpolate_nan
 from scipy.interpolate import griddata
 
+
 def get_experiment_wells(experiment):
 	
 	"""
@@ -801,6 +802,10 @@ def auto_load_number_of_frames(stack_path):
 	"""
 
 	# Try to estimate automatically # frames
+	
+	if stack_path is None:
+		return None
+
 	stack_path = stack_path.replace('\\','/')
 
 	with TiffFile(stack_path) as tif:
@@ -1370,6 +1375,14 @@ def control_segmentation_napari(position, prefix='Aligned', population="target",
 		exp_name = os.path.split(expfolder)[-1]
 		print(exp_name)
 
+		wells = get_experiment_wells(expfolder)
+		well_idx = list(wells).index(str(parent1)+os.sep)
+		ab = get_experiment_antibodies(expfolder)[well_idx]
+		conc = get_experiment_concentrations(expfolder)[well_idx]
+		ct = get_experiment_cell_types(expfolder)[well_idx]
+		pa = get_experiment_pharmaceutical_agents(expfolder)[well_idx]
+
+
 		spatial_calibration = float(ConfigSectionMap(config,"MovieSettings")["pxtoum"])
 		channel_names, channel_indices = extract_experiment_channels(config)
 
@@ -1424,7 +1437,7 @@ def control_segmentation_napari(position, prefix='Aligned', population="target",
 				multichannel = np.array(multichannel)        
 				save_tiff_imagej_compatible(annotation_folder + f"{exp_name}_{position.split(os.sep)[-2]}_{str(t).zfill(4)}_roi_{xmin}_{xmax}_{ymin}_{ymax}_labelled.tif", labels_layer[xmin:xmax,ymin:ymax].astype(np.int16), axes='YX')
 				save_tiff_imagej_compatible(annotation_folder + f"{exp_name}_{position.split(os.sep)[-2]}_{str(t).zfill(4)}_roi_{xmin}_{xmax}_{ymin}_{ymax}.tif", multichannel, axes='CYX')
-				info = {"spatial_calibration": spatial_calibration, "channels": list(channel_names)}
+				info = {"spatial_calibration": spatial_calibration, "channels": list(channel_names), 'cell_type': ct, 'antibody': ab, 'concentration': conc, 'pharmaceutical_agent': pa}
 				info_name = annotation_folder + f"{exp_name}_{position.split(os.sep)[-2]}_{str(t).zfill(4)}_roi_{xmin}_{xmax}_{ymin}_{ymax}.json"
 				with open(info_name, 'w') as f:
 					json.dump(info, f, indent=4)
@@ -1441,7 +1454,7 @@ def control_segmentation_napari(position, prefix='Aligned', population="target",
 			multichannel = np.array(multichannel)		
 			save_tiff_imagej_compatible(annotation_folder + f"{exp_name}_{position.split(os.sep)[-2]}_{str(t).zfill(4)}_labelled.tif", labels_layer, axes='YX')
 			save_tiff_imagej_compatible(annotation_folder + f"{exp_name}_{position.split(os.sep)[-2]}_{str(t).zfill(4)}.tif", multichannel, axes='CYX')
-			info = {"spatial_calibration": spatial_calibration, "channels": list(channel_names)}
+			info = {"spatial_calibration": spatial_calibration, "channels": list(channel_names), 'cell_type': ct, 'antibody': ab, 'concentration': conc, 'pharmaceutical_agent': pa}
 			info_name = annotation_folder + f"{exp_name}_{position.split(os.sep)[-2]}_{str(t).zfill(4)}.json"
 			with open(info_name, 'w') as f:
 				json.dump(info, f, indent=4)
@@ -1481,6 +1494,59 @@ def control_segmentation_napari(position, prefix='Aligned', population="target",
 		del labels
 		gc.collect()
 
+def correct_annotation(filename):
+
+	"""
+	New function to reannotate an annotation image in post, using napari and save update inplace.
+	"""
+	
+	def export_labels():
+		labels_layer = viewer.layers['segmentation'].data
+		for t,im in enumerate(tqdm(labels_layer)):
+
+			try:
+				im = auto_correct_masks(im)
+			except Exception as e:
+				print(e)
+
+			save_tiff_imagej_compatible(existing_lbl, im.astype(np.int16), axes='YX')
+		print("The labels have been successfully rewritten.")
+
+	@magicgui(call_button='Save the modified labels')
+	def save_widget():
+		return export_labels()
+	
+	img = imread(filename.replace('\\','/'))
+	if img.ndim==3:
+		img = np.moveaxis(img, 0, -1)
+	elif img.ndim==2:
+		img = img[:,:,np.newaxis]
+		
+	existing_lbl = filename.replace('.tif','_labelled.tif')
+	if os.path.exists(existing_lbl):
+		labels = imread(existing_lbl)[np.newaxis,:,:].astype(int)
+	else:
+		labels = np.zeros_like(img[:,:,0]).astype(int)[np.newaxis,:,:]
+		
+	stack = img[np.newaxis,:,:,:]
+	
+	viewer = napari.Viewer()
+	viewer.add_image(stack,channel_axis=-1,colormap=["gray"]*stack.shape[-1])
+	viewer.add_labels(labels, name='segmentation',opacity=0.4)
+	viewer.window.add_dock_widget(save_widget, area='right')
+	viewer.show(block=True)    
+
+	# temporary fix for slight napari memory leak
+	for i in range(100):
+		try:
+			viewer.layers.pop()
+		except:
+			pass
+	del viewer
+	del stack
+	del labels
+	gc.collect()
+
 
 def _view_on_napari(tracks=None, stack=None, labels=None):
 	
@@ -1628,7 +1694,7 @@ def locate_segmentation_model(name):
 	"""
 
 	main_dir = os.sep.join([os.path.split(os.path.dirname(os.path.realpath(__file__)))[0],"celldetective"])
-	modelpath = os.sep.join([main_dir, "models", "segmentation*", os.sep])
+	modelpath = os.sep.join([main_dir, "models", "segmentation*"]) + os.sep
 	print(f'Looking for {name} in {modelpath}')
 	models = glob(modelpath+f'*{os.sep}')
 

celldetective/neighborhood.py

@@ -48,14 +48,14 @@ def set_live_status(setA,setB,status, not_status_option):
 
 	"""
 
-
-	if status is None:
+	print(f"Provided statuses: {status}...")
+	if status is None or status==["live_status","live_status"] or status==[None,None]:
 		setA.loc[:,'live_status'] = 1
 		setB.loc[:,'live_status'] = 1
 		status = ['live_status', 'live_status']
 	elif isinstance(status,list):
 		assert len(status)==2,'Please provide only two columns to classify cells as alive or dead.'
-		if status[0] is None:
+		if status[0] is None or status[0]=='live_status':
 			setA.loc[:,'live_status'] = 1
 			status[0] = 'live_status'
 		elif status[0] is not None and isinstance(not_status_option,list):
@@ -63,7 +63,7 @@ def set_live_status(setA,setB,status, not_status_option):
 			if not_status_option[0]:
 				setA.loc[:,'not_'+status[0]] = [not a if a==0 or a==1 else np.nan for a in setA.loc[:,status[0]].values]
 				status[0] = 'not_'+status[0]
-		if status[1] is None:
+		if status[1] is None or status[1]=='live_status':
 			setB.loc[:,'live_status'] = 1
 			status[1] = 'live_status'
 		elif status[1] is not None and isinstance(not_status_option,list):
@@ -380,16 +380,39 @@ def compute_neighborhood_at_position(pos, distance, population=['targets','effec
 	if df_A_pkl is not None:
 		pkl_columns = np.array(df_A_pkl.columns)
 		neigh_columns = np.array([c.startswith('neighborhood') for c in pkl_columns])
-		cols = list(pkl_columns[neigh_columns]) + ['TRACK_ID','FRAME']
+		cols = list(pkl_columns[neigh_columns]) + ['FRAME']
+
+		if 'TRACK_ID' in list(pkl_columns):
+			cols.append('TRACK_ID')
+			on_cols = ['TRACK_ID','FRAME']
+		else:
+			cols.append('ID')
+			on_cols = ['ID','FRAME']
+
 		print(f'Recover {cols} from the pickle file...')
-		df_A = pd.merge(df_A, df_A_pkl.loc[:,cols], how="outer", on=['TRACK_ID','FRAME'])
-		print(df_A.columns)
+		try:
+			df_A = pd.merge(df_A, df_A_pkl.loc[:,cols], how="outer", on=on_cols)
+			print(df_A.columns)
+		except Exception as e:
+			print(f'Failure to merge pickle and csv files: {e}')
+
 	if df_B_pkl is not None and df_B is not None:
 		pkl_columns = np.array(df_B_pkl.columns)
 		neigh_columns = np.array([c.startswith('neighborhood') for c in pkl_columns])
-		cols = list(pkl_columns[neigh_columns]) + ['TRACK_ID','FRAME']
+		cols = list(pkl_columns[neigh_columns]) + ['FRAME']
+
+		if 'TRACK_ID' in list(pkl_columns):
+			cols.append('TRACK_ID')
+			on_cols = ['TRACK_ID','FRAME']
+		else:
+			cols.append('ID')
+			on_cols = ['ID','FRAME']
+
 		print(f'Recover {cols} from the pickle file...')
-		df_B = pd.merge(df_B, df_B_pkl.loc[:,cols], how="outer", on=['TRACK_ID','FRAME'])
+		try:
+			df_B = pd.merge(df_B, df_B_pkl.loc[:,cols], how="outer", on=on_cols)
+		except Exception as e:
+			print(f'Failure to merge pickle and csv files: {e}')
 
 	if clear_neigh:
 		unwanted = df_A.columns[df_A.columns.str.contains('neighborhood')]
@@ -408,17 +431,22 @@ def compute_neighborhood_at_position(pos, distance, population=['targets','effec
 		elif neighborhood_kwargs['mode']=='self':
 			neigh_col = f'neighborhood_self_circle_{d}_px'
 
-		edge_filter_A = (df_A['POSITION_X'] > td)&(df_A['POSITION_Y'] > td)&(df_A['POSITION_Y'] < (img_shape[0] - td))&(df_A['POSITION_X'] < (img_shape[1] - td))
-		edge_filter_B = (df_B['POSITION_X'] > td)&(df_B['POSITION_Y'] > td)&(df_B['POSITION_Y'] < (img_shape[0] - td))&(df_B['POSITION_X'] < (img_shape[1] - td))
-		df_A.loc[~edge_filter_A, neigh_col] = np.nan
-		df_B.loc[~edge_filter_B, neigh_col] = np.nan
+		# edge_filter_A = (df_A['POSITION_X'] > td)&(df_A['POSITION_Y'] > td)&(df_A['POSITION_Y'] < (img_shape[0] - td))&(df_A['POSITION_X'] < (img_shape[1] - td))
+		# edge_filter_B = (df_B['POSITION_X'] > td)&(df_B['POSITION_Y'] > td)&(df_B['POSITION_Y'] < (img_shape[0] - td))&(df_B['POSITION_X'] < (img_shape[1] - td))
+		# df_A.loc[~edge_filter_A, neigh_col] = np.nan
+		# df_B.loc[~edge_filter_B, neigh_col] = np.nan
 
+		print('Count neighborhood...')
 		df_A = compute_neighborhood_metrics(df_A, neigh_col, metrics=['inclusive','exclusive','intermediate'], decompose_by_status=True)
 		if neighborhood_kwargs['symmetrize']:
 			df_B = compute_neighborhood_metrics(df_B, neigh_col, metrics=['inclusive','exclusive','intermediate'], decompose_by_status=True)
+		print('Done...')
 		
-		df_A = mean_neighborhood_before_event(df_A, neigh_col, event_time_col)
-		df_A = mean_neighborhood_after_event(df_A, neigh_col, event_time_col)
+		if 'TRACK_ID' in list(df_A.columns):
+			print('Estimate average neighborhood before/after event...')
+			df_A = mean_neighborhood_before_event(df_A, neigh_col, event_time_col)
+			df_A = mean_neighborhood_after_event(df_A, neigh_col, event_time_col)
+			print('Done...')
 
 	df_A.to_pickle(path_A.replace('.csv','.pkl'))
 	if not population[0]==population[1]:
@@ -485,9 +513,14 @@ def compute_neighborhood_metrics(neigh_table, neigh_col, metrics=['inclusive','e
 
 	neigh_table = neigh_table.reset_index(drop=True)
 	if 'position' in list(neigh_table.columns):
-		groupbycols = ['position','TRACK_ID']
+		groupbycols = ['position']
 	else:
-		groupbycols = ['TRACK_ID']
+		groupbycols = []
+	if 'TRACK_ID' in list(neigh_table.columns):
+		groupbycols.append('TRACK_ID')
+	else:
+		groupbycols.append('ID')
+
 	neigh_table.sort_values(by=groupbycols+['FRAME'],inplace=True)
 
 	for tid,group in neigh_table.groupby(groupbycols):
@@ -607,10 +640,16 @@ def mean_neighborhood_before_event(neigh_table, neigh_col, event_time_col, metri
 	"""
 
 
+	neigh_table = neigh_table.reset_index(drop=True)
 	if 'position' in list(neigh_table.columns):
-		groupbycols = ['position','TRACK_ID']
+		groupbycols = ['position']
 	else:
-		groupbycols = ['TRACK_ID']
+		groupbycols = []
+	if 'TRACK_ID' in list(neigh_table.columns):
+		groupbycols.append('TRACK_ID')
+	else:
+		groupbycols.append('ID')
+
 	neigh_table.sort_values(by=groupbycols+['FRAME'],inplace=True)
 	suffix = '_before_event'
 	
@@ -681,10 +720,16 @@ def mean_neighborhood_after_event(neigh_table, neigh_col, event_time_col, metric
 	"""
 
 
+	neigh_table = neigh_table.reset_index(drop=True)
 	if 'position' in list(neigh_table.columns):
-		groupbycols = ['position','TRACK_ID']
+		groupbycols = ['position']
+	else:
+		groupbycols = []
+	if 'TRACK_ID' in list(neigh_table.columns):
+		groupbycols.append('TRACK_ID')
 	else:
-		groupbycols = ['TRACK_ID']
+		groupbycols.append('ID')
+
 	neigh_table.sort_values(by=groupbycols+['FRAME'],inplace=True)
 	suffix = '_after_event'
 	
@@ -1096,6 +1141,8 @@ def compute_contact_neighborhood_at_position(pos, distance, population=['targets
 
 	df_A, path_A = get_position_table(pos, population=population[0], return_path=True)
 	df_B, path_B = get_position_table(pos, population=population[1], return_path=True)
+	if df_A is None or df_B is None:
+		return None
 
 	df_A_pkl = get_position_pickle(pos, population=population[0], return_path=False)
 	df_B_pkl = get_position_pickle(pos, population=population[1], return_path=False)
@@ -1103,16 +1150,39 @@ def compute_contact_neighborhood_at_position(pos, distance, population=['targets
 	if df_A_pkl is not None:
 		pkl_columns = np.array(df_A_pkl.columns)
 		neigh_columns = np.array([c.startswith('neighborhood') for c in pkl_columns])
-		cols = list(pkl_columns[neigh_columns]) + ['TRACK_ID','FRAME']
+		cols = list(pkl_columns[neigh_columns]) + ['FRAME']
+		
+		if 'TRACK_ID' in list(pkl_columns):
+			cols.append('TRACK_ID')
+			on_cols = ['TRACK_ID','FRAME']
+		else:
+			cols.append('ID')
+			on_cols = ['ID','FRAME']
+
 		print(f'Recover {cols} from the pickle file...')
-		df_A = pd.merge(df_A, df_A_pkl.loc[:,cols], how="outer", on=['TRACK_ID','FRAME'])
-		print(df_A.columns)
+		try:
+			df_A = pd.merge(df_A, df_A_pkl.loc[:,cols], how="outer", on=on_cols)
+			print(df_A.columns)
+		except Exception as e:
+			print(f'Failure to merge pickle and csv files: {e}')
+
 	if df_B_pkl is not None and df_B is not None:
 		pkl_columns = np.array(df_B_pkl.columns)
 		neigh_columns = np.array([c.startswith('neighborhood') for c in pkl_columns])
-		cols = list(pkl_columns[neigh_columns]) + ['TRACK_ID','FRAME']
+		cols = list(pkl_columns[neigh_columns]) + ['FRAME']
+
+		if 'TRACK_ID' in list(pkl_columns):
+			cols.append('TRACK_ID')
+			on_cols = ['TRACK_ID','FRAME']
+		else:
+			cols.append('ID')
+			on_cols = ['ID','FRAME']
+
 		print(f'Recover {cols} from the pickle file...')
-		df_B = pd.merge(df_B, df_B_pkl.loc[:,cols], how="outer", on=['TRACK_ID','FRAME'])
+		try:
+			df_B = pd.merge(df_B, df_B_pkl.loc[:,cols], how="outer", on=on_cols)
+		except Exception as e:
+			print(f'Failure to merge pickle and csv files: {e}')
 
 	labelsA = locate_labels(pos, population=population[0])
 	if population[1]==population[0]:

celldetective/preprocessing.py

@@ -69,6 +69,7 @@ def estimate_background_per_condition(experiment, threshold_on_std=1, well_optio
 	...     print(bg["well"], bg["bg"].shape)
 	"""
 
+
 	config = get_config(experiment)
 	wells = get_experiment_wells(experiment)
 	len_movie = float(ConfigSectionMap(config,"MovieSettings")["len_movie"])
@@ -79,7 +80,7 @@ def estimate_background_per_condition(experiment, threshold_on_std=1, well_optio
 	channel_indices = _extract_channel_indices_from_config(config, [target_channel])
 	nbr_channels = _extract_nbr_channels_from_config(config)
 	img_num_channels = _get_img_num_per_channel(channel_indices, int(len_movie), nbr_channels)
-	
+
 	backgrounds = []
 
 	for k, well_path in enumerate(tqdm(wells[well_indices], disable=not show_progress_per_well)):
@@ -95,51 +96,65 @@ def estimate_background_per_condition(experiment, threshold_on_std=1, well_optio
 		for l,pos_path in enumerate(tqdm(positions, disable=not show_progress_per_pos)):
 			
 			stack_path = get_position_movie_path(pos_path, prefix=movie_prefix)
+			if stack_path is not None:
+				len_movie_auto = auto_load_number_of_frames(stack_path)
+				if len_movie_auto is not None:
+					len_movie = len_movie_auto
+					img_num_channels = _get_img_num_per_channel(channel_indices, int(len_movie), nbr_channels)
 
-			if mode=="timeseries":
-
-				frames = load_frames(img_num_channels[0,frame_range[0]:frame_range[1]], stack_path, normalize_input=False)
-				frames = np.moveaxis(frames, -1, 0).astype(float)
-
-				for i in range(len(frames)):
-					if np.all(frames[i].flatten()==0):
-						frames[i,:,:] = np.nan
-
-				frame_mean = np.nanmean(frames, axis=0)
-				
-				frame = frame_mean.copy().astype(float)
-				std_frame = filter_image(frame.copy(),filters=activation_protocol)
-				edge = estimate_unreliable_edge(activation_protocol)
-				mask = threshold_image(std_frame, threshold_on_std, 1.0E06, foreground_value=1, edge_exclusion=edge)
-				frame[np.where(mask.astype(int)==1)] = np.nan
-
-			elif mode=="tiles":
-
-				frames = load_frames(img_num_channels[0,:], stack_path, normalize_input=False).astype(float)
-				frames = np.moveaxis(frames, -1, 0).astype(float)
+				if mode=="timeseries":
 
-				for i in range(len(frames)):
+					frames = load_frames(img_num_channels[0,frame_range[0]:frame_range[1]], stack_path, normalize_input=False)
+					frames = np.moveaxis(frames, -1, 0).astype(float)
 
-					if np.all(frames[i].flatten()==0):
-						frames[i,:,:] = np.nan
-						continue
+					for i in range(len(frames)):
+						if np.all(frames[i].flatten()==0):
+							frames[i,:,:] = np.nan
 
-					f = frames[i].copy()
-					std_frame = filter_image(f.copy(),filters=activation_protocol)
+					frame_mean = np.nanmean(frames, axis=0)
+					
+					frame = frame_mean.copy().astype(float)
+					std_frame = filter_image(frame.copy(),filters=activation_protocol)
 					edge = estimate_unreliable_edge(activation_protocol)
-					mask = threshold_image(std_frame, threshold_on_std, 1.0E06, foreground_value=1, edge_exclusion=edge)
-					f[np.where(mask.astype(int)==1)] = np.nan
-
-					frames[i,:,:] = f
-				
-				frame = np.nanmedian(frames, axis=0)
+					mask = threshold_image(std_frame, threshold_on_std, np.inf, foreground_value=1, edge_exclusion=edge)
+					frame[np.where(mask.astype(int)==1)] = np.nan
+
+				elif mode=="tiles":
+
+					frames = load_frames(img_num_channels[0,:], stack_path, normalize_input=False).astype(float)
+					frames = np.moveaxis(frames, -1, 0).astype(float)
+
+					new_frames = []
+					for i in range(len(frames)):
+
+						if np.all(frames[i].flatten()==0):
+							empty_frame = np.zeros_like(frames[i])
+							empty_frame[:,:] = np.nan
+							new_frames.append(empty_frame)
+							continue
+
+						f = frames[i].copy()
+						std_frame = filter_image(f.copy(),filters=activation_protocol)
+						edge = estimate_unreliable_edge(activation_protocol)
+						mask = threshold_image(std_frame, threshold_on_std, np.inf, foreground_value=1, edge_exclusion=edge)
+						f[np.where(mask.astype(int)==1)] = np.nan
+						new_frames.append(f.copy())
+					
+					frame = np.nanmedian(new_frames, axis=0)
+			else:
+				print(f'Stack not found for position {pos_path}...')
+				frame = []
 
 			# store
 			frame_mean_per_position.append(frame)
 
-		background = np.nanmedian(frame_mean_per_position,axis=0)
-		backgrounds.append({"bg": background, "well": well_path})
-		print(f"Background successfully computed for well {well_name}...")
+		try:
+			background = np.nanmedian(frame_mean_per_position,axis=0)
+			backgrounds.append({"bg": background, "well": well_path})
+			print(f"Background successfully computed for well {well_name}...")
+		except Exception as e:
+			print(e)
+			backgrounds.append(None)		
 
 	return backgrounds
 
@@ -261,28 +276,35 @@ def correct_background_model_free(
 			
 			stack_path = get_position_movie_path(pos_path, prefix=movie_prefix)
 			print(f'Applying the correction to position {extract_position_name(pos_path)}...')
-
-			corrected_stack = apply_background_to_stack(stack_path, 
-														background,
-														target_channel_index=channel_indices[0],
-														nbr_channels=nbr_channels,
-														stack_length=len_movie,
-														threshold_on_std=threshold_on_std,
-														optimize_option=optimize_option,
-														opt_coef_range=opt_coef_range,
-														opt_coef_nbr=opt_coef_nbr,
-														operation=operation,
-														clip=clip,
-														export=export,
-														activation_protocol=activation_protocol,
-														prefix=export_prefix,
-													  )
-			print('Correction successful.')
-			if return_stacks:
-				stacks.append(corrected_stack)
+			if stack_path is not None:
+				len_movie_auto = auto_load_number_of_frames(stack_path)
+				if len_movie_auto is not None:
+					len_movie = len_movie_auto
+					img_num_channels = _get_img_num_per_channel(channel_indices, int(len_movie), nbr_channels)
+
+				corrected_stack = apply_background_to_stack(stack_path, 
+															background,
+															target_channel_index=channel_indices[0],
+															nbr_channels=nbr_channels,
+															stack_length=len_movie,
+															threshold_on_std=threshold_on_std,
+															optimize_option=optimize_option,
+															opt_coef_range=opt_coef_range,
+															opt_coef_nbr=opt_coef_nbr,
+															operation=operation,
+															clip=clip,
+															export=export,
+															activation_protocol=activation_protocol,
+															prefix=export_prefix,
+														  )
+				print('Correction successful.')
+				if return_stacks:
+					stacks.append(corrected_stack)
+				else:
+					del corrected_stack
+				collect()
 			else:
-				del corrected_stack
-			collect()
+				stacks.append(None)
 
 	if return_stacks:
 		return stacks
@@ -371,18 +393,25 @@ def apply_background_to_stack(stack_path, background, target_channel_index=0, nb
 
 			std_frame = filter_image(target_copy.copy(),filters=activation_protocol)
 			edge = estimate_unreliable_edge(activation_protocol)
-			mask = threshold_image(std_frame, threshold_on_std, 1.0E06, foreground_value=1, edge_exclusion=edge)
+			mask = threshold_image(std_frame, threshold_on_std, np.inf, foreground_value=1, edge_exclusion=edge)
 			target_copy[np.where(mask.astype(int)==1)] = np.nan
 
 			loss = []
 			
 			# brute-force regression, could do gradient descent instead
 			for c in coefficients:
+				
 				target_crop = unpad(target_copy,edge)
 				bg_crop = unpad(background, edge)
-				diff = np.subtract(target_crop, c*bg_crop, where=target_crop==target_crop)
-				s = np.sum(np.abs(diff, where=diff==diff), where=diff==diff)
+				
+				roi = np.zeros_like(target_crop).astype(int)
+				roi[target_crop!=target_crop] = 1
+				roi[bg_crop!=bg_crop] = 1
+
+				diff = np.subtract(target_crop, c*bg_crop, where=roi==0)
+				s = np.sum(np.abs(diff, where=roi==0), where=roi==0)
 				loss.append(s)
+
 			c = coefficients[np.argmin(loss)]
 			print(f"Frame: {i}; optimal coefficient: {c}...")
 			# if c==min(coefficients) or c==max(coefficients):
@@ -758,7 +787,10 @@ def correct_background_model(
 			
 			stack_path = get_position_movie_path(pos_path, prefix=movie_prefix)
 			print(f'Applying the correction to position {extract_position_name(pos_path)}...')
-			print(stack_path)
+			len_movie_auto = auto_load_number_of_frames(stack_path)
+			if len_movie_auto is not None:
+				len_movie = len_movie_auto
+				img_num_channels = _get_img_num_per_channel(channel_indices, int(len_movie), nbr_channels)
 
 			corrected_stack = fit_and_apply_model_background_to_stack(stack_path, 
 														target_channel_index=channel_indices[0],
@@ -954,7 +986,7 @@ def field_correction(img, threshold_on_std=1, operation='divide', model='parabol
 
 	std_frame = filter_image(target_copy,filters=activation_protocol)
 	edge = estimate_unreliable_edge(activation_protocol)
-	mask = threshold_image(std_frame, threshold_on_std, 1.0E06, foreground_value=1, edge_exclusion=edge).astype(int)
+	mask = threshold_image(std_frame, threshold_on_std, np.inf, foreground_value=1, edge_exclusion=edge).astype(int)
 	background = fit_background_model(img, cell_masks=mask, model=model, edge_exclusion=edge)
 
 	if operation=="divide":

celldetective/scripts/segment_cells.py

@@ -148,6 +148,7 @@ def segment_index(indices):
 		model.config.use_gpu = use_gpu
 		model.use_gpu = use_gpu
 		print(f"StarDist model {modelname} successfully loaded.")
+		scale_model = scale
 
 	elif model_type=='cellpose':
 

celldetective/scripts/train_segmentation_model.py

@@ -60,17 +60,10 @@ batch_size = training_instructions['batch_size']
 
 # Load dataset
 print(f'Datasets: {datasets}')
-X,Y = load_image_dataset(datasets, target_channels, train_spatial_calibration=spatial_calibration,
+X,Y,filenames = load_image_dataset(datasets, target_channels, train_spatial_calibration=spatial_calibration,
 						mask_suffix='labelled')
 print('Dataset loaded...')
 
-# Normalize images
-# X = normalize_per_channel(X,
-# 						  normalization_percentile_mode=normalization_percentile, 
-# 						  normalization_values=normalization_values, 
-# 						  normalization_clipping=normalization_clip
-# 						  )
-
 values = []
 percentiles = []
 for k in range(len(normalization_percentile)):
@@ -88,17 +81,6 @@ for k in range(len(X)):
 	x_interp = np.moveaxis([interpolate_nan(x[:,:,c].copy()) for c in range(x.shape[-1])],0,-1)
 	X[k] = x_interp
 
-# for x in X[:10]:
-# 	plt.imshow(x[:,:,0])
-# 	plt.colorbar()
-# 	plt.pause(2)
-# 	plt.close()
-
-# 	plt.imshow(x[:,:,1])
-# 	plt.colorbar()
-# 	plt.pause(2)
-# 	plt.close()
-
 Y = [fill_label_holes(y) for y in tqdm(Y)]
 
 assert len(X) > 1, "not enough training data"
@@ -107,7 +89,11 @@ ind = rng.permutation(len(X))
 n_val = max(1, int(round(validation_split * len(ind))))
 ind_train, ind_val = ind[:-n_val], ind[-n_val:]
 X_val, Y_val = [X[i] for i in ind_val]  , [Y[i] for i in ind_val]
-X_trn, Y_trn = [X[i] for i in ind_train], [Y[i] for i in ind_train] 
+X_trn, Y_trn = [X[i] for i in ind_train], [Y[i] for i in ind_train]
+
+files_train = [filenames[i] for i in ind_train]
+files_val = [filenames[i] for i in ind_val]
+
 print('number of images: %3d' % len(X))
 print('- training:       %3d' % len(X_trn))
 print('- validation:     %3d' % len(X_val))
@@ -134,7 +120,10 @@ if model_type=='cellpose':
 	import torch
 	
 	if not 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: ',pretrained)
@@ -163,7 +152,7 @@ if model_type=='cellpose':
 	config_inputs = {"channels": target_channels, "diameter": standard_diameter, 'cellprob_threshold': 0., 'flow_threshold': 0.4,
 	'normalization_percentile': normalization_percentile, 'normalization_clip': normalization_clip,
 	'normalization_values': normalization_values, 'model_type': 'cellpose',
-	'spatial_calibration': input_spatial_calibration}
+	'spatial_calibration': input_spatial_calibration, 'dataset': {'train': files_train, 'validation': files_val}}
 	json_input_config = json.dumps(config_inputs, indent=4)
 	with open(os.sep.join([target_directory, model_name, "config_input.json"]), "w") as outfile:
 		outfile.write(json_input_config)
@@ -234,7 +223,7 @@ elif model_type=='stardist':
 
 	config_inputs = {"channels": target_channels, 'normalization_percentile': normalization_percentile,
 	'normalization_clip': normalization_clip, 'normalization_values': normalization_values, 
-	'model_type': 'stardist', 'spatial_calibration': spatial_calibration}
+	'model_type': 'stardist', 'spatial_calibration': spatial_calibration, 'dataset': {'train': files_train, 'validation': files_val}}
 
 	json_input_config = json.dumps(config_inputs, indent=4)
 	with open(os.sep.join([target_directory, model_name, "config_input.json"]), "w") as outfile: