adding cellpose2.0 options to GUI, and fixing some text warnings in GUI

Author: carsen-stringer

suite2p/default_ops.py

@@ -74,7 +74,6 @@ def default_ops():
         '1Preg': False,  # whether to perform high-pass filtering and tapering
         'spatial_hp': 42,  # window for spatial high-pass filtering before registration
         'spatial_hp_reg': 42,  # window for spatial high-pass filtering before registration
-        'spatial_hp_detect': 25,  # window for spatial high-pass filtering for neuropil subtraction before detection
         'pre_smooth': 0,  # whether to smooth before high-pass filtering before registration
         'spatial_taper': 40,  # how much to ignore on edges (important for vignetted windows, for FFT padding do not set BELOW 3*ops['smooth_sigma'])
 
@@ -89,14 +88,16 @@ def default_ops():
         'threshold_scaling': 1.0,  # adjust the automatically determined threshold by this scalar multiplier
         'max_overlap': 0.75,  # cells with more overlap than this get removed during triage, before refinement
         'high_pass': 100,  # running mean subtraction with window of size 'high_pass' (use low values for 1P)
+        'spatial_hp_detect': 25,  # window for spatial high-pass filtering for neuropil subtraction before detection
         'denoise': False, # denoise binned movie for cell detection in sparse_mode
 
-        # cell detection settings with cellpose
-        'anatomical_only': 0, # use cellpose masks from mean image (no functional segmentation)
-        'cellprob_threshold': 0.0, # cellprob_threshold for cellpose (if anatomical_only > 1)
-        'flow_threshold': 1.5, # flow_threshold for cellpose (if anatomical_only > 1)
-        'diameter': 0,  # if anatomical_only, use diameter for cellpose, if 0 estimate diameter
-        'pretrained_model': 'cyto', # path to pretrained model or model type string in Cellpose
+        # cell detection settings with cellpose (used if anatomical_only > 0)
+        'anatomical_only': 0, # run cellpose to get masks on 1: max_proj / mean_img; 2: mean_img; 3: mean_img enhanced, 4: max_proj
+        'diameter': 0,  # use diameter for cellpose, if 0 estimate diameter
+        'cellprob_threshold': 0.0, # cellprob_threshold for cellpose
+        'flow_threshold': 1.5, # flow_threshold for cellpose
+        'spatial_hp_cp': 0, # high-pass image spatially by a multiple of the diameter
+        'pretrained_model': 'cyto', # path to pretrained model or model type string in Cellpose (can be user model)
 
         # classification parameters
         'soma_crop': True, # crop dendrites for cell classification stats like compactness

suite2p/detection/anatomical.py

@@ -1,9 +1,10 @@
 import numpy as np
 from typing import Any, Dict
-from scipy.ndimage import find_objects
+from scipy.ndimage import find_objects, gaussian_filter
 from cellpose.models import CellposeModel, Cellpose
 from cellpose import transforms, dynamics
 from cellpose.utils import fill_holes_and_remove_small_masks
+from cellpose.transforms import normalize99
 import time
 import cv2
 import os
@@ -150,25 +151,29 @@ def select_rois(ops: Dict[str, Any], mov: np.ndarray,
     max_proj = mov.max(axis=0)
     #max_proj = np.percentile(mov, 90, axis=0) #.mean(axis=0)
     if ops['anatomical_only'] == 1:
-        mproj = np.log(np.maximum(1e-3, max_proj / np.maximum(1e-3, mean_img)))
+        img = np.log(np.maximum(1e-3, max_proj / np.maximum(1e-3, mean_img)))
         weights = max_proj
     elif ops['anatomical_only']==2:
-        mproj = mean_img
+        img = mean_img
         weights = 0.1 + np.clip((mean_img - np.percentile(mean_img,1)) / 
                                 (np.percentile(mean_img,99) - np.percentile(mean_img,1)), 0, 1)
-    else:
+    elif ops['anatomical_only']==3:
         if 'meanImgE' in ops:
-            mproj = ops['meanImgE'][ops['yrange'][0]:ops['yrange'][1], ops['xrange'][0]:ops['xrange'][1]]
+            img = ops['meanImgE'][ops['yrange'][0]:ops['yrange'][1], ops['xrange'][0]:ops['xrange'][1]]
         else:
-            mproj = mean_img
+            img = mean_img
             print('no enhanced mean image, using mean image instead')
         weights = 0.1 + np.clip((mean_img - np.percentile(mean_img,1)) / 
                                 (np.percentile(mean_img,99) - np.percentile(mean_img,1)), 0, 1)
+    else:
+        img = max_proj.copy()
+        weights = max_proj
+
     t0 = time.time()
     if diameter is not None:
         if isinstance(diameter, (list, np.ndarray)) and len(ops['diameter'])>1:
             rescale = diameter[1] / diameter[0]
-            mproj = cv2.resize(mproj, (Lxc, int(Lyc*rescale)))
+            img = cv2.resize(img, (Lxc, int(Lyc*rescale)))
         else:
             rescale = 1.0
             diameter = [diameter, diameter]
@@ -178,13 +183,18 @@ def select_rois(ops: Dict[str, Any], mov: np.ndarray,
             print("!NOTE! diameter set to 0 or None, diameter will be estimated by cellpose")
     else:
         print("!NOTE! diameter set to 0 or None, diameter will be estimated by cellpose")
-    masks, centers, median_diam, mask_diams = roi_detect(mproj, diameter=diameter[1],
+
+    if ops.get('spatial_hp_cp', 0):
+        img = np.clip(normalize99(img), 0, 1)
+        img -= gaussian_filter(img, diameter[1]*ops['spatial_hp_cp'])
+
+    masks, centers, median_diam, mask_diams = roi_detect(img, diameter=diameter[1],
                                                          flow_threshold=ops['flow_threshold'],
                                                          cellprob_threshold=ops['cellprob_threshold'],
                                                          pretrained_model=ops['pretrained_model'])
     if rescale != 1.0:
         masks = cv2.resize(masks, (Lxc, Lyc), interpolation=cv2.INTER_NEAREST)
-        mproj = cv2.resize(mproj, (Lxc, Lyc))
+        img = cv2.resize(img, (Lxc, Lyc))
     stats = masks_to_stats(masks, weights)
     print('Detected %d ROIs, %0.2f sec' % (len(stats), time.time() - t0))
 
@@ -194,7 +204,7 @@ def select_rois(ops: Dict[str, Any], mov: np.ndarray,
             'Vmax': 0,
             'ihop': 0,
             'Vsplit': 0,
-            'Vcorr': mproj,
+            'Vcorr': img,
             'Vmap': 0,
             'spatscale_pix': 0
         }

suite2p/detection/detect.py

@@ -139,7 +139,7 @@ def detection_wrapper(f_reg, mov=None, yrange=None, xrange=None,
 			print('~~~ tried to import cellpose to run anatomical but failed, install with: ~~~')
 			print('$ pip install cellpose')
 		else:
-			print('>>>> CELLPOSE finding masks in ' + ['max_proj / mean_img', 'mean_img', 'enhanced_mean_img'][int(ops['anatomical_only'])-1])
+			print('>>>> CELLPOSE finding masks in ' + ['max_proj / mean_img', 'mean_img', 'enhanced_mean_img', 'max_proj'][int(ops['anatomical_only'])-1])
 			stat = anatomical.select_rois(
 						ops=ops,
 						mov=mov,

suite2p/gui/masks.py

@@ -230,17 +230,9 @@ def init_masks(parent):
     parent.rois['Lam']    = np.zeros((2,3,Ly,Lx), np.float32)
     parent.rois['iROI']   = -1 * np.ones((2,3,Ly,Lx), np.int32)
 
-    for n in range(len(parent.roi_text_labels)):
+    if parent.checkBoxN.isChecked():
         parent.checkBoxN.setChecked(False)
-        try:
-            parent.p1.removeItem(parent.roi_text_labels[n])
-        except:
-            pass 
-        try:
-            parent.p2.removeItem(parent.roi_text_labels[n])
-        except:
-            pass 
-        
+            
     # ignore merged cells
     iignore = np.zeros(ncells, 'bool')
     parent.roi_text_labels = []
@@ -536,14 +528,13 @@ def flip_roi(parent):
     n = parent.ichosen
     i = int(1-parent.iscell[n])
     i0 = 1-i
-    if parent.roitext:
-        if i0==1:
+    if parent.checkBoxN.isChecked():
+        if i0==0:
             parent.p1.removeItem(parent.roi_text_labels[n])
             parent.p2.addItem(parent.roi_text_labels[n])
         else:
             parent.p2.removeItem(parent.roi_text_labels[n])
             parent.p1.addItem(parent.roi_text_labels[n])
-        
 
     # remove ROI
     remove_roi(parent, n, i0)

suite2p/gui/rungui.py

@@ -81,19 +81,21 @@ def reset_ops(self):
     def create_buttons(self):
         self.intkeys = ['nplanes', 'nchannels', 'functional_chan', 'align_by_chan', 'nimg_init',
                    'batch_size', 'max_iterations', 'nbinned','inner_neuropil_radius',
-                   'min_neuropil_pixels', 'spatial_scale', 'do_registration']
+                   'min_neuropil_pixels', 'spatial_scale', 'do_registration', 'anatomical_only']
         self.boolkeys = ['delete_bin', 'move_bin','do_bidiphase', 'reg_tif', 'reg_tif_chan2',
                      'save_mat', 'save_NWB' 'combined', '1Preg', 'nonrigid', 
                     'connected', 'roidetect', 'neuropil_extract', 
                     'spikedetect', 'keep_movie_raw', 'allow_overlap', 'sparse_mode']
+        self.stringkeys = ['pretrained_model']
         tifkeys = ['nplanes','nchannels','functional_chan','tau','fs','do_bidiphase','bidiphase', 'multiplane_parallel', 'ignore_flyback']
         outkeys = ['preclassify','save_mat','save_NWB','combined','reg_tif','reg_tif_chan2','aspect','delete_bin','move_bin']
         regkeys = ['do_registration','align_by_chan','nimg_init','batch_size','smooth_sigma', 'smooth_sigma_time','maxregshift','th_badframes','keep_movie_raw','two_step_registration']
         nrkeys = [['nonrigid','block_size','snr_thresh','maxregshiftNR'], ['1Preg','spatial_hp_reg','pre_smooth','spatial_taper']]
-        cellkeys = ['roidetect', 'denoise', 'anatomical_only', 'diameter', 'spatial_scale', 'threshold_scaling', 'max_overlap','max_iterations','high_pass']
+        cellkeys = ['roidetect', 'denoise', 'spatial_scale', 'threshold_scaling', 'max_overlap','max_iterations','high_pass','spatial_hp_detect']
+        anatkeys = ['anatomical_only', 'diameter', 'cellprob_threshold', 'flow_threshold', 'pretrained_model', 'spatial_hp_cp']
         neudeconvkeys = [['neuropil_extract', 'allow_overlap','inner_neuropil_radius','min_neuropil_pixels'], ['soma_crop','spikedetect','win_baseline','sig_baseline','neucoeff']]
-        keys = [tifkeys, outkeys, regkeys, nrkeys, cellkeys, neudeconvkeys]
-        labels = ['Main settings','Output settings','Registration',['Nonrigid','1P'],'ROI detection',['Extraction/Neuropil','Classification/Deconvolution']]
+        keys = [tifkeys, outkeys, regkeys, nrkeys, cellkeys, anatkeys, neudeconvkeys]
+        labels = ['Main settings','Output settings','Registration',['Nonrigid','1P'],'Functional detect', 'Anat detect', ['Extraction/Neuropil','Classify/Deconv']]
         tooltips = ['each tiff has this many planes in sequence',
                     'each tiff has this many channels per plane',
                     'this channel is used to extract functional ROIs (1-based)',
@@ -130,15 +132,20 @@ def create_buttons(self):
                     'window for spatial high-pass filtering before registration',
                     'whether to smooth before high-pass filtering before registration',
                     "how much to ignore on edges (important for vignetted windows, for FFT padding do not set BELOW 3*smooth_sigma)",
-                    'if 1, run cell (ROI) detection',
+                    'if 1, run cell (ROI) detection (either functional or anatomical if anatomical_only > 0)',
                     'if 1, run PCA denoising on binned movie to improve cell detection',
-                    'run cellpose to get masks on 1: max_proj / mean_img; 2: mean_img; or 3: mean_img enhanced',
-                    'if anatomical_only>0, input average diameter of ROIs in recording (can give a list e.g. 6,9)',
-                    'if anatomical_only=0, choose size of ROIs: 0 = multi-scale; 1 = 6 pixels, 2 = 12, 3 = 24, 4 = 48',
+                    'choose size of ROIs: 0 = multi-scale; 1 = 6 pixels, 2 = 12, 3 = 24, 4 = 48',
                     'adjust the automatically determined threshold for finding ROIs by this scalar multiplier',
                     'ROIs with greater than this overlap as a fraction of total pixels will be discarded',
                     'maximum number of iterations for ROI detection',
-                    'running mean subtraction with window of size "high_pass" (use low values for 1P)',
+                    'temporal running mean subtraction with window of size "high_pass" (use low values for 1P)',
+                    'spatial high-pass filter size (used to remove spatially-correlated neuropil)',
+                    'run cellpose to get masks on 1: max_proj / mean_img; 2: mean_img; 3: mean_img enhanced, 4: max_proj',
+                    'input average diameter of ROIs in recording (can give a list e.g. 6,9 if aspect not equal), if set to 0 auto-determination run by Cellpose',
+                    'cellprob_threshold for cellpose',
+                    'flow_threshold for cellpose (throws out masks, if getting too few masks, set to 0)',
+                    'high-pass image spatially by a multiple of the diameter (if field is non-uniform, a value of ~2 is recommended',
+                    'model type string from Cellpose (can be a built-in model or a user model that is added to the Cellpose GUI)',
                     'whether or not to extract neuropil; if 0, Fneu is set to 0',
                     'allow shared pixels to be used for fluorescence extraction from overlapping ROIs (otherwise excluded from both ROIs)',
                     'number of pixels between ROI and neuropil donut',
@@ -147,7 +154,8 @@ def create_buttons(self):
                     'if 1, run spike detection (deconvolution)',
                     'window for maximin',
                     'smoothing constant for gaussian filter',
-                    'neuropil coefficient']
+                    'neuropil coefficient',
+                    ]
 
         bigfont = QtGui.QFont("Arial", 10, QtGui.QFont.Bold)
         qlabel = QLabel('File paths')
@@ -263,7 +271,7 @@ def create_buttons(self):
         self.layout.addWidget(self.binlabel,19,0,1,cw)
         self.runButton = QPushButton('RUN SUITE2P')
         self.runButton.clicked.connect(self.run_S2P)
-        n0 = 21
+        n0 = 22
         self.layout.addWidget(self.runButton,n0,0,1,1)
         self.runButton.setEnabled(False)
         self.textEdit = QTextEdit()
@@ -289,7 +297,7 @@ def create_buttons(self):
         self.cleanButton.clicked.connect(self.clean_script)
         self.cleanLabel = QLabel('')
         self.layout.addWidget(self.cleanLabel,n0,4,1,12)
-        n0+=1
+        #n0+=1
         self.listOps = QPushButton('save settings and\n add more (batch)')
         self.listOps.clicked.connect(self.add_batch)
         self.layout.addWidget(self.listOps,n0,12,1,2)
@@ -357,7 +365,7 @@ def add_ops(self):
 
     def compile_ops_db(self):
         for k,key in enumerate(self.keylist):
-            self.ops[key] = self.editlist[k].get_text(self.intkeys, self.boolkeys)
+            self.ops[key] = self.editlist[k].get_text(self.intkeys, self.boolkeys, self.stringkeys)
         self.db = {}
         self.db['data_path'] = self.data_path
         self.db['subfolders'] = []
@@ -606,7 +614,7 @@ def __init__(self,k,key,parent=None):
         self.key = key
         #self.textEdited.connect(lambda: self.edit_changed(parent.ops, k))
 
-    def get_text(self,intkeys,boolkeys):
+    def get_text(self,intkeys,boolkeys,stringkeys):
         key = self.key
         if key=='diameter' or key=='block_size':
             diams = self.text().replace(' ','').split(',')
@@ -625,6 +633,8 @@ def get_text(self,intkeys,boolkeys):
                 okey = int(float(self.text()))
             elif key in boolkeys:
                 okey = bool(int(self.text()))
+            elif key in stringkeys:
+                okey = self.text()
             else:
                 okey = float(self.text())
         return okey