labdata 0.1.18__tar.gz → 0.1.20__tar.gz

{labdata-0.1.18 → labdata-0.1.20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: labdata
-Version: 0.1.18
+Version: 0.1.20
 Summary: Package to manage data in experimental neuroscience labs
 Author-email: Joao Couto <jpcouto@gmail.com>
 Project-URL: Homepage, https://github.com/jcouto/labdata

{labdata-0.1.18 → labdata-0.1.20}/labdata/__init__.py

@@ -1,4 +1,4 @@
-VERSION = "0.1.18"
+VERSION = "0.1.20"
 
 from .utils import *
 from .copy import copy_to_upload_server

{labdata-0.1.18 → labdata-0.1.20}/labdata/compute/__init__.py

@@ -6,6 +6,7 @@ Tasks can be scheduled and run on dedicated containers using job schedulers.
 Available compute tasks:
 - SpksCompute: Spike sorting using Kilosort/Phy via SPKS
 - DeeplabcutCompute: Animal pose estimation using DeepLabCut
+- LightningPoseCompute: Animal pose estimation using Lightning Pose
 
 Each compute task can be:
 - Scheduled to run on compute clusters via SLURM/PBS
@@ -24,6 +25,6 @@ The compute tasks handle:
 from .utils import *
 from .singularity import build_singularity_container
 from .ephys import SpksCompute
-from .pose import DeeplabcutCompute
+from .pose import DeeplabcutCompute, LightningPoseCompute
 from .caiman import CaimanCompute
 from .suite2p import Suite2pCompute

labdata-0.1.20/labdata/compute/caiman.py

@@ -0,0 +1,498 @@
+from ..utils import *
+from .utils import BaseCompute
+from .suite2p import Suite2pCompute
+class CaimanCompute(Suite2pCompute):
+    '''Cell segmentation compute class using CaImAn.
+
+    Works for both one-photon (miniscope) and two-photon datasets.  Runs
+    inside the ``labdata-caiman`` Apptainer container (CPU-only).
+
+    The compute pipeline:
+
+    1. Identify files and determine dataset type (``Miniscope`` or ``TwoPhoton``).
+    2. Copy the ``zarr.zip`` stack to the scratch directory.
+    3. Run CaImAn CNMF/CNMF-E with the specified parameters.
+    4. Delete memory-mapped intermediate files.
+    5. Upload the HDF5 results file to S3 and insert into ``CellSegmentation``.
+
+    Miniscope recordings run CNMF-E; two-photon recordings run CNMF, one
+    plane at a time on the channel given with ``--channel`` (0 by default),
+    and the ROIs of each plane are stored under its ``plane_num``. A job
+    submitted with ``--multisession`` segments several sessions of a subject
+    together, as ``Suite2pCompute`` does: the frames are concatenated in
+    session order, one set of ROIs per plane is found on all of them, and
+    the traces, projections and motion shifts are stored per session under
+    the same ROI numbers, so cells can be followed across sessions.
+
+    Parameters are stored in ``CellSegmentationParams`` (shared with
+    ``Suite2pCompute`` using the same part tables).
+    '''
+    container = 'labdata-caiman'
+    cuda = False
+    name = 'caiman'
+    url = 'http://github.com/flatironinstitute/CaImAn'
+    def __init__(self, job_id, project=None, allow_s3=None, **kwargs):
+        super(Suite2pCompute,self).__init__(job_id, project = project, allow_s3 = allow_s3) # takes BaseCompute init
+        self.file_filters = ['.zarr.zip'] # this only runs on zarr.zip for the moment.
+        # default parameters
+        self.parameters = dict(algorithm_name = 'caiman')
+        # will only store these in CellSegmentationParams
+        self.parameter_keys = ['pw_rigid','p','gSig','gSig_filt','max_shifts','gSiz','merge_thr','rf','stride','tsub',
+                               'ssub','nb','K','min_corr','min_pnr','ssub_B','ring_size_factor',
+                               'min_SNR', 'rval_thr', 'use_cnn', 'detrendWin','quantileMin','denoise_dff']
+        
+        self.parameter_set_num = None # identifier in CellSegmentationParams
+
+        self._init_job()
+        if not self.job_id is None:
+            self.add_parameter_key()
+      
+    def _secondary_parse(self,arguments,parameter_number = None):
+        '''
+        Handles parsing the command line interface
+        '''
+        if not parameter_number is None:
+            self.parameters = ((self.schema.CellSegmentationParams() & f'parameter_set_num = {parameter_number}')).fetch(as_dict = True)
+            if not len(self.parameters):
+                raise(f'Could not find parameter {parameter_number} in CellSegmentationParams.')
+            self.parameters = self.parameters[0]
+        else:
+            import argparse
+            parser = argparse.ArgumentParser(
+                description = 'Segmentation of imaging datasets using CaImAn.',
+                usage = 'caiman -a <SUBJECT> -s <SESSION> -- <PARAMETERS>')
+            parser.add_argument('-i','--parameter_set',default = None, type=int, help = 'Parameter set number')
+            parser.add_argument('-m','--pwrigid',
+                                action='store_true', default=False,
+                                help = "Piecewise-rigid registration")
+            parser.add_argument('--max-allowed-shifts',
+                                action='store', default=[20,20], type = int, nargs = 2,
+                                help = "maximum allowed shift for motion correction")
+            parser.add_argument('-p',
+                                action='store', default=1, type = int,
+                                help = "Order of the autoregressive system")
+            parser.add_argument('-g','--gsig',
+                                action='store', default=[6,6], type = int, nargs = 2,
+                                help = "Expected halfwidth of the neurons in pixels")
+            parser.add_argument('--gsig-filt',
+                                action='store', default=None, type = int, nargs = 2,
+                                help = "Filter for 1photon data, will highlight the somas for motion correction (defaults to [6, 6] for 1p data).")
+            parser.add_argument('--nb',
+                                action='store', default=None, type = int,
+                                help = "number of background components (rank); 0 runs CNMF-E. Default 0 for miniscope and 2 for two-photon")
+            parser.add_argument('-k', '--k',
+                                action='store', default=None, type = int,
+                                help = "[two-photon] number of components per patch for the initialisation; default 4 (unused by CNMF-E)")
+            parser.add_argument('--channel',
+                                action='store', default=0, type = int,
+                                help = "[two-photon] channel of the stack to segment; default 0")
+            parser.add_argument('-r','--rf',
+                                action='store', default=40, type = int,
+                                help = "half-size of the patches in pixels. e.g., if rf=40, patches are 80x80")
+            parser.add_argument('--stride',
+                                action='store', default=20, type = int,
+                                help = "size of the overlap in pixels")
+            parser.add_argument('-s','--ring-size',
+                                action='store', default=1.4, type = float,
+                                help = "radius of ring is gSiz*ring_size_factor")
+            parser.add_argument('--merge-thr',
+                                action='store', default=.7, type = float,
+                                help = "merging threshold, max correlation allowed")
+            parser.add_argument('--tsub',
+                                action='store', default=4, type = int,
+                                help = "downsampling factor in time for initialization, increase if you have memory problems")
+            parser.add_argument('--ssub',
+                                action='store', default=2, type = int,
+                                help = "downsampling factor in space for initialization, increase if you have memory problems")
+            parser.add_argument('--ssub_b',
+                                action='store', default=2, type = int,
+                                help = "additional downsampling factor for the background")
+            parser.add_argument('--min-corr',
+                                action='store', default=.8, type = float,
+                                help = "min peak value from correlation image")
+            parser.add_argument('--min-pnr',
+                                action='store', default=10, type = float,
+                                help = "min peak to noise ratio")
+
+            parser.add_argument('--snr_thr',
+                                action='store', default=7, type = float,
+                                help = "[cell selection] signal to noise ratio threshold")
+            parser.add_argument('--rval-thr',
+                                action='store', default=0.85, type = float,
+                                help = "[cell selection] spatial correlation threshold")
+            parser.add_argument('--use-cnn',
+                                action='store_true', default = False,
+                                help = "[cell selection] use a CNN to help deciding good from bad units")
+            parser.add_argument('--denoise-dff',
+                                action='store_true', default = False,
+                                help = "[df/f] compute dff on the denoised data")
+
+            parser.add_argument('--quantile_min',
+                                action='store', default=8, type = float,
+                                help = "[df/f] Minimum quantile for df/f detrending")
+            parser.add_argument('--detrend_win',
+                                action='store', default=250, type = int,
+                                help = "[df/f] Number of frames for detrending")
+            parser.add_argument('--roi',
+                                action='store', default=None, type = int, nargs = 4,
+                                help = "ROI")
+
+            args = parser.parse_args(arguments[1:])
+            
+            params = dict(pw_rigid = args.pwrigid,
+                          max_shifts = args.max_allowed_shifts,
+                          p = int(args.p),
+                          gSig = [int(a) for a in args.gsig], 
+                          gSiz = [int(a) for a in 2*np.array(args.gsig) + 1],
+                          gSig_filt = args.gsig_filt,
+                          merge_thr = float(args.merge_thr),      
+                          rf = int(args.rf),                    
+                          stride = int(args.stride),              
+                          tsub = int(args.tsub),                  
+                          ssub = int(args.ssub),                  
+                          nb = None if args.nb is None else int(args.nb),
+                          K = None if args.k is None else int(args.k),
+                          channel = int(args.channel),
+                          min_corr = float(args.min_corr),             
+                          min_pnr = float(args.min_pnr),               
+                          ssub_B = int(args.ssub_b),                 
+                          ring_size_factor = float(args.ring_size),
+                          min_SNR = float(args.snr_thr),
+                          rval_thr = float(args.rval_thr),
+                          use_cnn = bool(args.use_cnn),
+                          detrendWin = int(args.detrend_win),
+                          denoise_dff = bool(args.denoise_dff),
+                          quantileMin = float(args.quantile_min),
+                          roi = args.roi)
+            if not args.gsig_filt is None:
+                params['gSig_filt'] = [int(a) for a in args.gsig_filt], 
+
+            self.parameters = params
+        
+    def _recordings(self):
+        """Return the recordings of the job in session order.
+
+        Returns ``(kind, keys, recordings, frame_rate, n_planes, n_channels)``
+        with ``kind`` 'miniscope' or 'two-photon', ``keys`` the dataset dicts
+        sorted by session name, which is the order the frames are concatenated
+        in, and ``recordings`` the matching table queries. Raises ValueError
+        when a dataset has no recording, the datasets mix kinds, or the
+        sessions differ in frame rate, planes or channels.
+        """
+        seskeys = sorted((dict(k) for k in self.dataset_key), key = lambda k: (k['subject_name'], k['session_name']))
+        kinds = []
+        for k in seskeys:
+            if len(self.schema.Miniscope() & k):
+                kinds.append('miniscope')
+            elif len(self.schema.TwoPhoton() & k):
+                kinds.append('two-photon')
+            else:
+                raise ValueError(f'No miniscope or two-photon recording on record for {k}.')
+        if len(set(kinds)) > 1:
+            raise ValueError(f'The datasets mix miniscope and two-photon recordings: {list(zip(kinds, seskeys))}.')
+        kind = kinds[0]
+        table = self.schema.Miniscope if kind == 'miniscope' else self.schema.TwoPhoton
+        recordings = [table() & k for k in seskeys]
+        rows = [r.fetch1() for r in recordings]
+        rates = [float(r['frame_rate']) for r in rows]
+        if max(rates) - min(rates) > 1e-3:
+            raise ValueError(f'The sessions have different frame rates and cannot be segmented together: {rates}.')
+        if kind == 'two-photon':
+            planes = {int(r['n_planes']) for r in rows}
+            channels = {int(r['n_channels']) for r in rows}
+            if len(planes) > 1 or len(channels) > 1:
+                raise ValueError(f'The sessions differ in planes {sorted(planes)} or channels {sorted(channels)}.')
+            n_planes, n_channels = planes.pop(), channels.pop()
+        else:
+            n_planes, n_channels = 1, 1
+        return kind, seskeys, recordings, rates[0], n_planes, n_channels
+
+    def _cnmf_parameters(self, kind, params):
+        """Return the CNMF parameters for a kind of recording, the stored parameter set on top of the defaults of the kind.
+
+        Miniscope recordings run CNMF-E (``corr_pnr`` initialisation, no
+        background rank, ring model); two-photon recordings run CNMF
+        (``greedy_roi`` initialisation, two background components). A stored
+        value of None for ``nb``, ``K`` or ``gSig_filt`` keeps the default of
+        the kind.
+        """
+        cnmfparams = dict(motion_correct = True,
+                          del_duplicates = True,                # whether to remove duplicates from initialization
+                          update_background_components = True,  # sometimes setting to False improve the results
+                          method_deconvolution = 'oasis')       # could use 'cvxpy' alternatively
+        if kind == 'miniscope':
+            cnmfparams.update(gSig_filt = [6,6],                # 1p needs this for motion correction (spatial filter to register somas)
+                              method_init = 'corr_pnr',          # use this for 1 photon
+                              K = None,                          # for 1p use None
+                              nb = 0,                            # number of background components (rank), 0 for CNMFE
+                              nb_patch = 0,
+                              low_rank_background = None,        # for 1p
+                              normalize_init = False,            # just leave as is
+                              center_psf = True,                 # True for 1p
+                              only_init = True)                  # set it to True to run CNMF-E
+        else:
+            cnmfparams.update(gSig_filt = None,                  # no spatial filter for 2p motion correction
+                              method_init = 'greedy_roi',        # 2p initialisation
+                              K = 4,                             # components per patch
+                              nb = 2,                            # background components
+                              nb_patch = 1,
+                              low_rank_background = True,
+                              normalize_init = True,
+                              center_psf = False,
+                              only_init = False)                 # run the full CNMF after the initialisation
+        excluded = ['detrendWin','denoise_dff','quantileMin','roi','channel']
+        for k, v in params.items():
+            if k in excluded:
+                continue
+            if v is None and k in ('nb', 'K', 'gSig_filt'):
+                continue                                         # keep the default of the kind
+            cnmfparams[k] = v
+        if kind == 'miniscope' and cnmfparams.get('gSig_filt') is None:
+            cnmfparams['gSig_filt'] = [6,6]
+            print('[caiman] Using default gSig_filt: (6,6) because this is 1photon.')
+        return cnmfparams
+
+    def _segment_plane(self, stacks, plane, channel, seskeys, cnmfparams, params, frame_rate, cluster, folder):
+        """Motion correct and segment one plane of the sessions' stacks, concatenated in order, and return the rows to store.
+
+        ``plane`` is None for a miniscope stack. The result is a dict with
+        ``n_rois`` and the lists of rows for the plane, motion correction,
+        projection, ROI, trace and selection tables, one entry per session
+        with that session's frames of the traces and shifts.
+        """
+        import time
+        from ..stacks import export_to_tiff, compute_projections, get_roi_pixels
+        from caiman.source_extraction.cnmf.params import CNMFParams
+        from caiman.motion_correction import MotionCorrect
+        from caiman import save_memmap, load_memmap
+        from caiman.source_extraction import cnmf
+        folder = Path(folder)
+        folder.mkdir(parents = True, exist_ok = True)
+        os.environ['CAIMAN_DATA'] = f'{folder}'
+        plane_num = 0 if plane is None else int(plane)
+        # one tiff sequence per session, in session order; caiman treats the files as one movie
+        paths, n_frames = [], []
+        for isession, dat in enumerate(stacks):
+            paths += export_to_tiff(dat, folder/f'session_{isession:02d}', crop_region = self.parameters.get('roi'),
+                                    plane = plane, channel = channel if plane is not None else None)
+            n_frames.append(int(dat.shape[0]))
+        offsets = np.cumsum([0, *n_frames])
+        kk = ['motion_correct','pw_rigid','max_shifts','gSig_filt']
+        parameters = CNMFParams(params_dict = dict({k:cnmfparams[k] for k in kk if k in cnmfparams},
+                                                   fr = frame_rate))
+        pmotion = parameters.get_group('motion')
+
+        tstart = time.time()
+        mot_correct = MotionCorrect(paths, dview=cluster, **pmotion)
+        mot_correct.motion_correct(save_movie=True)
+        print(f'[caiman] Plane {plane_num}: done with motion correction in {(time.time() - tstart)/60.} min.')
+        
+        [os.unlink(f) for f in paths] # clear raw converted tiff
+        fname_mc = mot_correct.fname_tot_els if cnmfparams['pw_rigid'] else mot_correct.fname_tot_rig
+        if pmotion['pw_rigid']:
+            bord_px = np.ceil(np.maximum(np.max(np.abs(mot_correct.x_shifts_els)),
+                                         np.max(np.abs(mot_correct.y_shifts_els)))).astype(int)
+        else:
+            bord_px = np.ceil(np.max(np.abs(mot_correct.shifts_rig))).astype(int)
+
+        bord_px = 0 if pmotion['border_nan'] == 'copy' else bord_px
+        fname_new = save_memmap(fname_mc, base_name='memmap_', order='C',
+                                border_to_0 = bord_px)
+        [os.unlink(f) for f in fname_mc]
+        [os.unlink(f) for f in folder.rglob("*chunk_*order_F*")]; # clear motion correction temporary
+        
+        Yr, dims, T = load_memmap(fname_new)
+        images = Yr.T.reshape((T,) + dims, order='F')
+        if T != offsets[-1]:
+            raise ValueError(f'The motion corrected movie has {T} frames; the sessions have {offsets[-1]}.')
+
+        tstart = time.time()
+        parameters.change_params(cnmfparams)
+        cnmfe_model = cnmf.CNMF(n_processes = DEFAULT_N_JOBS, 
+                                dview = cluster, 
+                                params = parameters)
+        cnmfe_model.fit(images);
+        print(f'[caiman] Plane {plane_num}: done with CNMF in {(time.time() - tstart)/60.} min.')
+        
+        quality_params = {'min_SNR': params['min_SNR'],
+                          'rval_thr': params['rval_thr'],
+                          'use_cnn': params['use_cnn']}
+        cnmfe_model.params.change_params(params_dict=quality_params)
+        print(f"Evaluating components.")
+        cnmfe_model.estimates.evaluate_components(images, cnmfe_model.params, dview=cluster)
+        print(f"Computing df/f.")
+        cnmfe_model.estimates.detrend_df_f(quantileMin = params['quantileMin'], 
+                                           frames_window = int(params['detrendWin']),
+                                           flag_auto = False,
+                                           use_residuals = not params['denoise_dff'])
+        print('*****')
+        print(f"Plane {plane_num}: total number of components: {len(cnmfe_model.estimates.C)}")
+        print(f"Number accepted: {len(cnmfe_model.estimates.idx_components)}")
+        print(f"Number rejected: {len(cnmfe_model.estimates.idx_components_bad)}")
+        roi_masks = np.array(cnmfe_model.estimates.A.todense()).reshape((*cnmfe_model.dims,-1))
+        roi_masks = roi_masks.transpose(2,1,0)
+        # the ROIs are the same in every session; each session gets its own frames of the traces and shifts
+        rois = []
+        for icell,roi in enumerate(roi_masks):
+            roi_pixels, roi_pixels_values = get_roi_pixels(roi)
+            if len(roi_pixels)>5000:
+                print(f"[CaimanCompute] {icell} had an roi with {len(roi_pixels)} pixels. Skipping.")
+                continue
+            rois.append(dict(roi_num = icell,
+                             roi_pixels = roi_pixels.astype(int),
+                             roi_pixels_values = roi_pixels_values.astype(np.float32),
+                             selection = 1 if icell in cnmfe_model.estimates.idx_components else 0))
+        motion = None if params['pw_rigid'] else np.array(mot_correct.shifts_rig).astype(np.float32)
+        out = dict(n_rois = len(rois), planedicts = [], motiondicts = [], projdicts = [], roidict = [],
+                   tracesdict = [], rawtracesdict = [], deconvtracesdict = [], selectiondict = [])
+        for key, on, off in zip(seskeys, offsets[:-1], offsets[1:]):
+            mean_proj,std_proj,max_proj,corr_proj = compute_projections(images[on:off])
+            planekey = dict(key,
+                            parameter_set_num  = self.parameter_set_num,
+                            plane_num = plane_num)
+            out['planedicts'].append(dict(planekey,
+                                          plane_n_rois = len(rois),
+                                          dims = [a for a in mean_proj.shape]))
+            if motion is not None:
+                out['motiondicts'].append(dict(planekey,
+                                               motion_block_size = 0,
+                                               displacement = motion[on:off]))
+            out['projdicts'] += [dict(planekey,
+                                      proj_name = n,
+                                      proj_im = pi) for n,pi in zip(
+                                          ['mean','std','max','lcorr'],
+                                          [mean_proj,std_proj,max_proj,corr_proj])]
+            for roi in rois:
+                icell = roi['roi_num']
+                out['roidict'].append(dict(planekey,
+                                           roi_num = icell,
+                                           roi_pixels = roi['roi_pixels'],
+                                           roi_pixels_values = roi['roi_pixels_values']))
+                out['tracesdict'].append(dict(planekey,
+                                              roi_num = icell,
+                                              dff = cnmfe_model.estimates.F_dff[icell].astype(np.float32)[on:off]))
+                out['rawtracesdict'].append(dict(planekey,
+                                                 roi_num = icell,
+                                                 f_trace = cnmfe_model.estimates.C[icell].astype(np.float32)[on:off]))
+                out['deconvtracesdict'].append(dict(planekey,
+                                                    roi_num = icell,
+                                                    deconv = cnmfe_model.estimates.S[icell].astype(np.float32)[on:off]))
+                out['selectiondict'].append(dict(planekey,
+                                                 roi_num = icell,
+                                                 selection_method = 'auto',
+                                                 selection = roi['selection']))
+        print(f'[caiman] Plane {plane_num}: projections computed.')
+        del images, Yr
+        if not self.keep_intermediate:
+            os.unlink(fname_new)   # the motion corrected memmap of this plane
+        return out
+
+    def _compute(self):
+        # one dataset, or the sessions of a multisession job segmented together as suite2p does:
+        # the frames are concatenated in session order, one set of ROIs per plane is found on all
+        # of them, and the traces, projections and motion shifts are stored per session
+        import string
+        rand = ''.join(np.random.choice([s for s in string.ascii_lowercase + string.digits],9))
+        temporary_folder = Path(prefs['scratch_path'])/f'caiman_temporary_{rand}'
+        temporary_folder.mkdir(parents = True, exist_ok = True)
+        kind, seskeys, recordings, frame_rate, n_planes, n_channels = self._recordings()
+        channel = int(self.parameters.get('channel') or 0)
+        if not 0 <= channel < n_channels:
+            raise ValueError(f'The recording has {n_channels} channel(s); channel {channel} is out of range.')
+        parameters = (self.schema.CellSegmentationParams & f'parameter_set_num = {self.parameter_set_num}').fetch1()
+        params = json.loads(parameters['parameters_dict'])
+        cnmfparams = self._cnmf_parameters(kind, params)
+        print(f'[caiman] {kind}: {len(seskeys)} session(s), {n_planes} plane(s), channel {channel} of {n_channels}.', flush = True)
+        # downloaded files are removed at the end; files that were local stay
+        self.files_existed = all(len((self.schema.File & r).check_if_files_local()[1]) == 0 for r in recordings)
+        stacks = [r.open() for r in recordings]
+        if len(seskeys) > 1:
+            print(f'[caiman] Segmenting {len(seskeys)} sessions together, {sum(int(d.shape[0]) for d in stacks)} frames: ' +
+                  ', '.join(f"{k['session_name']} ({d.shape[0]})" for k, d in zip(seskeys, stacks)), flush = True)
+            
+        import logging
+        logger = logging.getLogger('caiman')
+        # Set to logging.INFO if you want more, not very useful output..
+        logger.setLevel(logging.WARNING)
+        handler = logging.StreamHandler()
+        logger.addHandler(handler)
+        
+        n_cpus = DEFAULT_N_JOBS
+        cluster = setup_cluster(n_cpus)
+        totals = [0 for k in seskeys]
+        collected = dict(planedicts = [], motiondicts = [], projdicts = [], roidict = [],
+                         tracesdict = [], rawtracesdict = [], deconvtracesdict = [], selectiondict = [])
+        try:
+            for iplane in range(n_planes):
+                plane = iplane if kind == 'two-photon' else None
+                out = self._segment_plane(stacks, plane, channel, seskeys, cnmfparams, params, frame_rate, cluster,
+                                          temporary_folder/f'plane_{iplane}')
+                totals = [t + out['n_rois'] for t in totals]
+                for name in collected:
+                    collected[name] += out[name]
+        finally:
+            cluster.terminate()
+        if not self.keep_intermediate:
+            if not self.files_existed:
+                print('[caiman] Freeing up space by deleting raw data that was downloaded for this compute task.')
+                for r in recordings:
+                    for f in (self.schema.File & r).get():
+                        os.unlink(f)
+        import caiman
+        segmentation_datetime = datetime.now()
+        cellsegdicts = [dict(key,
+                             parameter_set_num = self.parameter_set_num,
+                             algorithm_version = f'caiman {caiman.__version__}',
+                             n_rois = n_rois,
+                             crop_region = self.parameters.get('roi'),
+                             segmentation_datetime = segmentation_datetime) for key, n_rois in zip(seskeys, totals)]
+        self.schema.CellSegmentation.insert(cellsegdicts, allow_direct_insert = True)
+        self.schema.CellSegmentation.Plane.insert(collected['planedicts'],allow_direct_insert = True)
+        if collected['motiondicts']:
+            self.schema.CellSegmentation.MotionCorrection.insert(collected['motiondicts'], allow_direct_insert = True)
+        self.schema.CellSegmentation.Projection.insert(collected['projdicts'],allow_direct_insert = True)
+        self.schema.CellSegmentation.ROI.insert(collected['roidict'] ,allow_direct_insert = True)
+        # Insert traces in parallel to prevent timeout errors from mysql
+        parallel_insert(self.schema.schema_project,'CellSegmentation.Traces',collected['tracesdict'], n_jobs = DEFAULT_N_JOBS,
+                        skip_duplicates = True, ignore_extra_fields = True)
+        parallel_insert(self.schema.schema_project,'CellSegmentation.RawTraces',collected['rawtracesdict'], n_jobs = DEFAULT_N_JOBS,
+                        skip_duplicates = True, ignore_extra_fields = True)
+        parallel_insert(self.schema.schema_project,'CellSegmentation.Deconvolved',collected['deconvtracesdict'], n_jobs = DEFAULT_N_JOBS,
+                        skip_duplicates = True, ignore_extra_fields = True)
+        parallel_insert(self.schema.schema_project,'CellSegmentation.Selection',collected['selectiondict'], n_jobs = DEFAULT_N_JOBS,
+                        skip_duplicates = True, ignore_extra_fields = True)                        
+        if not self.keep_intermediate:
+            print(f'[{self.name} job] Removing the temporary folder.')
+            import shutil
+            shutil.rmtree(temporary_folder)
+        else:
+            print(f'[{self.name} job] Kept the temporary folder {temporary_folder}.')
+        for key in seskeys:
+            self.schema.CellSegmentationMetrics().populate(key, display_progress = True)
+        
+def setup_cluster(n_cpus, cluster=None):
+    '''Start (or restart) a CaImAn multiprocessing cluster.
+
+    If ``cluster`` is not ``None``, the existing cluster is stopped first.
+
+    Parameters
+    ----------
+    n_cpus : int
+        Number of parallel worker processes.
+    cluster : object or None, optional
+        Existing CaImAn cluster object to stop before starting a new one.
+
+    Returns
+    -------
+    object
+        The newly created CaImAn cluster (``dview``).
+    '''
+    from caiman import stop_server, cluster
+    #%% start a cluster for parallel processing (if a cluster already exists it will be closed and a new session will be opened)
+    if not cluster is None:  # 'locals' contains list of current local variables
+        stop_server(dview=cluster)
+    _, cluster, n_processes = cluster.setup_cluster(backend='multiprocessing',
+                                                    n_processes = n_cpus,
+                                                    ignore_preexisting = False)
+    print(f"Set up parallelization with {n_processes} processes.")
+    return cluster