def suns_online(filename_video, filename_CNN, Params_pre, Params_post, dims, \ frames_init, merge_every, batch_size_init=1, useSF=True, useTF=True, useSNR=True, \ med_subtract=False, update_baseline=False, \ useWT=False, show_intermediate=True, prealloc=True, display=True, useMP=True, p=None): '''The complete SUNS online procedure. It uses the trained CNN model from "filename_CNN" and the optimized hyper-parameters in "Params_post" to process the video "Exp_ID" in "dir_video" Inputs: filename_video (str): The path of the file of the input raw video. The file must be a ".h5" file, with dataset "mov" being the input video (shape = (T0,Lx0,Ly0)). filename_CNN (str): The path of the trained CNN model. Params_pre (dict): Parameters for pre-processing. Params_pre['gauss_filt_size'] (float): The standard deviation of the spatial Gaussian filter in pixels Params_pre['Poisson_filt'] (1D numpy.ndarray of float32): The temporal filter kernel Params_pre['num_median_approx'] (int): Number of frames used to compute the median and median-based standard deviation Params_pre['nn'] (int): Number of frames at the beginning of the video to be processed. The remaining video is not considered a part of the input video. Params_post (dict): Parameters for post-processing. Params_post['minArea']: Minimum area of a valid neuron mask (unit: pixels). Params_post['avgArea']: The typical neuron area (unit: pixels). Params_post['thresh_pmap']: The probablity threshold. Values higher than thresh_pmap are active pixels. It is stored in uint8, so it should be converted to float32 before using. Params_post['thresh_mask']: Threashold to binarize the real-number mask. Params_post['thresh_COM0']: Threshold of COM distance (unit: pixels) used for the first COM-based merging. Params_post['thresh_COM']: Threshold of COM distance (unit: pixels) used for the second COM-based merging. Params_post['thresh_IOU']: Threshold of IOU used for merging neurons. Params_post['thresh_consume']: Threshold of consume ratio used for merging neurons. Params_post['cons']: Minimum number of consecutive frames that a neuron should be active for. dims (tuplel of int, shape = (2,)): lateral dimension of the raw video. frames_init (int): Number of frames used for initialization. merge_every (int): SUNS online merge the newly segmented frames every "merge_every" frames. batch_size_init (int, default to 1): batch size of CNN inference for initialization frames. useSF (bool, default to True): True if spatial filtering is used. useTF (bool, default to True): True if temporal filtering is used. useSNR (bool, default to True): True if pixel-by-pixel SNR normalization filtering is used. med_subtract (bool, default to False): True if the spatial median of every frame is subtracted before temporal filtering. Can only be used when spatial filtering is not used. update_baseline (bool, default to False): True if the median and median-based std is updated every "frames_init" frames. useWT (bool, default to False): Indicator of whether watershed is used. show_intermediate (bool, default to True): Indicator of whether consecutive frame requirement is applied to screen neurons after every update. prealloc (bool, default to True): True if pre-allocate memory space for large variables. Achieve faster speed at the cost of higher memory occupation. display (bool, default to True): Indicator of whether to show intermediate information useMP (bool, defaut to True): indicator of whether multiprocessing is used to speed up. p (multiprocessing.Pool, default to None): Outputs: Masks (3D numpy.ndarray of bool, shape = (n,Lx0,Ly0)): the final segmented masks. Masks_2 (scipy.csr_matrix of bool, shape = (n,Lx0*Ly0)): the final segmented masks in the form of sparse matrix. time_total (list of float, shape = (3,)): the total time spent for initalization, online processing, and total processing time_frame (list of float, shape = (3,)): the average time spent on every frame for initalization, online processing, and total processing ''' if display: start = time.time() (Lx, Ly) = dims # zero-pad the lateral dimensions to multiples of 8, suitable for CNN rowspad = math.ceil(Lx / 8) * 8 colspad = math.ceil(Ly / 8) * 8 dimspad = (rowspad, colspad) Poisson_filt = Params_pre['Poisson_filt'] gauss_filt_size = Params_pre['gauss_filt_size'] nn = Params_pre['nn'] leng_tf = Poisson_filt.size leng_past = 2 * leng_tf # number of past frames stored for temporal filtering list_time_per = np.zeros(nn) # Load CNN model fff = get_shallow_unet() fff.load_weights(filename_CNN) # run CNN inference once to warm up init_imgs = np.zeros((batch_size_init, rowspad, colspad, 1), dtype='float32') init_masks = np.zeros((batch_size_init, rowspad, colspad, 1), dtype='uint8') fff.evaluate(init_imgs, init_masks, batch_size=batch_size_init) del init_imgs, init_masks # load optimal post-processing parameters minArea = Params_post['minArea'] avgArea = Params_post['avgArea'] # thresh_pmap = Params_post['thresh_pmap'] thresh_mask = Params_post['thresh_mask'] thresh_COM0 = Params_post['thresh_COM0'] thresh_COM = Params_post['thresh_COM'] thresh_IOU = Params_post['thresh_IOU'] thresh_consume = Params_post['thresh_consume'] cons = Params_post['cons'] # thresh_pmap_float = (Params_post['thresh_pmap']+1.5)/256 thresh_pmap_float = (Params_post['thresh_pmap'] + 1) / 256 # for published version # Spatial filtering preparation if useSF == True: # lateral dimensions slightly larger than the raw video but faster for FFT rows1 = cv2.getOptimalDFTSize(rowspad) cols1 = cv2.getOptimalDFTSize(colspad) # if the learned 2D and 3D wisdom files have been saved, load them. # Otherwise, learn wisdom later Length_data2 = str((rows1, cols1)) cc2 = load_wisdom_txt('wisdom\\' + Length_data2) Length_data3 = str((frames_init, rows1, cols1)) cc3 = load_wisdom_txt('wisdom\\' + Length_data3) if cc3: pyfftw.import_wisdom(cc3) # mask for spatial filter mask2 = plan_mask2(dims, (rows1, cols1), gauss_filt_size) # FFT planning (bb, bf, fft_object_b, fft_object_c) = plan_fft(frames_init, (rows1, cols1), prealloc) else: (mask2, bf, fft_object_b, fft_object_c) = (None, None, None, None) bb = np.zeros((frames_init, rowspad, colspad), dtype='float32') # Temporal filtering preparation frames_initf = frames_init - leng_tf + 1 if useTF == True: if prealloc: # past frames stored for temporal filtering past_frames = np.ones((leng_past, rowspad, colspad), dtype='float32') else: past_frames = np.zeros((leng_past, rowspad, colspad), dtype='float32') else: past_frames = None if prealloc: # Pre-allocate memory for some future variables med_frame2 = np.ones((rowspad, colspad, 2), dtype='float32') video_input = np.ones((frames_initf, rowspad, colspad), dtype='float32') pmaps_b_init = np.ones((frames_initf, Lx, Ly), dtype='uint8') frame_SNR = np.ones(dimspad, dtype='float32') pmaps_b = np.ones(dims, dtype='uint8') if update_baseline: video_tf_past = np.ones((frames_init, rowspad, colspad), dtype='float32') else: med_frame2 = np.zeros((rowspad, colspad, 2), dtype='float32') video_input = np.zeros((frames_initf, rowspad, colspad), dtype='float32') pmaps_b_init = np.zeros((frames_initf, Lx, Ly), dtype='uint8') frame_SNR = np.zeros(dimspad, dtype='float32') pmaps_b = np.zeros(dims, dtype='uint8') if update_baseline: video_tf_past = np.zeros((frames_init, rowspad, colspad), dtype='float32') if display: time_init = time.time() print('Parameter initialization time: {} s'.format(time_init - start)) # %% Load raw video h5_img = h5py.File(filename_video, 'r') video_raw = np.array(h5_img['mov']) h5_img.close() nframes = video_raw.shape[0] nframesf = nframes - leng_tf + 1 bb[:, :Lx, :Ly] = video_raw[:frames_init] if display: time_load = time.time() print('Load data: {} s'.format(time_load - time_init)) # %% Actual processing starts after the video is loaded into memory # Initialization using the first "frames_init" frames print('Initialization of algorithms using the first {} frames'.format( frames_init)) if display: start_init = time.time() med_frame3, segs_all, recent_frames = init_online( bb, dims, video_input, pmaps_b_init, fff, thresh_pmap_float, Params_post, \ med_frame2, mask2, bf, fft_object_b, fft_object_c, Poisson_filt, \ useSF=useSF, useTF=useTF, useSNR=useSNR, med_subtract=med_subtract, \ useWT=useWT, batch_size_init=batch_size_init, p=p) if useTF == True: past_frames[:leng_tf] = recent_frames tuple_temp = merge_complete(segs_all[:frames_initf], dims, Params_post) if show_intermediate: Masks_2 = select_cons(tuple_temp) if display: end_init = time.time() time_init = end_init - start_init time_frame_init = time_init / (frames_initf) * 1000 print('Initialization time: {:6f} s, {:6f} ms/frame'.format( time_init, time_frame_init)) if display: start_online = time.time() # Spatial filtering preparation for online processing. # Attention: this part counts to the total time if useSF: if cc2: pyfftw.import_wisdom(cc2) (bb, bf, fft_object_b, fft_object_c) = plan_fft2((rows1, cols1)) else: (bf, fft_object_b, fft_object_c) = (None, None, None) bb = np.zeros(dimspad, dtype='float32') print('Start frame by frame processing') # %% Online processing for the following frames current_frame = leng_tf + 1 t_merge = frames_initf for t in range(frames_initf, nframesf): if display: start_frame = time.time() # load the current frame bb[:Lx, :Ly] = video_raw[t + leng_tf - 1] bb[Lx:, :] = 0 bb[:, Ly:] = 0 # PreProcessing frame_SNR, frame_tf = preprocess_online(bb, dimspad, med_frame3, frame_SNR, \ past_frames[current_frame-leng_tf:current_frame], mask2, bf, fft_object_b, \ fft_object_c, Poisson_filt, useSF=useSF, useTF=useTF, useSNR=useSNR, \ med_subtract=med_subtract, update_baseline=update_baseline) if update_baseline: t_past = (t - frames_initf) % frames_init video_tf_past[t_past] = frame_tf if t_past == frames_init - 1: # update median and median-based standard deviation every "frames_init" frames if useSNR: med_frame3 = SNR_normalization(video_tf_past, med_frame2, (rowspad, colspad), 1, display=False) else: med_frame3 = median_normalization(video_tf_past, med_frame2, (rowspad, colspad), 1, display=False) # CNN inference frame_prob = CNN_online(frame_SNR, fff, dims) # first step of post-processing segs = separate_neuron_online(frame_prob, pmaps_b, thresh_pmap_float, minArea, avgArea, useWT) segs_all.append(segs) # temporal merging 1: combine neurons with COM distance smaller than thresh_COM0 if ((t + 1 - t_merge) == merge_every) or (t == nframesf - 1): # uniques, times_uniques = unique_neurons1_simp(segs_all[t_merge:], thresh_COM0) # minArea, totalmasks, neuronstate, COMs, areas, probmapID = segs_results( segs_all[t_merge:]) uniques, times_uniques = unique_neurons2_simp(totalmasks, neuronstate, COMs, \ areas, probmapID, minArea=0, thresh_COM0=thresh_COM0, useMP=useMP) # temporal merging 2: combine neurons with COM distance smaller than thresh_COM if ((t - 0 - t_merge) == merge_every) or (t == nframesf - 1): if uniques.size: groupedneurons, times_groupedneurons = \ group_neurons(uniques, thresh_COM, thresh_mask, dims, times_uniques, useMP=useMP) # temporal merging 3: combine neurons with IoU larger than thresh_IOU if ((t - 1 - t_merge) == merge_every) or (t == nframesf - 1): if uniques.size: piecedneurons_1, times_piecedneurons_1 = \ piece_neurons_IOU(groupedneurons, thresh_mask, thresh_IOU, times_groupedneurons) # temporal merging 4: combine neurons with conumse ratio larger than thresh_consume if ((t - 2 - t_merge) == merge_every) or (t == nframesf - 1): if uniques.size: piecedneurons, times_piecedneurons = \ piece_neurons_consume(piecedneurons_1, avgArea, thresh_mask, thresh_consume, times_piecedneurons_1) # masks of new neurons masks_add = piecedneurons # indices of frames when the neurons are active times_add = [ np.unique(x) + t_merge for x in times_piecedneurons ] # Refine neurons using consecutive occurence if masks_add.size: # new real-number masks masks_add = [x for x in masks_add] # new binary masks Masksb_add = [(x >= x.max() * thresh_mask).astype('float') for x in masks_add] # areas of new masks area_add = np.array([x.nnz for x in Masksb_add]) # indicators of whether the new masks satisfy consecutive frame requirement have_cons_add = refine_seperate_cons_online( times_add, cons) else: Masksb_add = [] area_add = np.array([]) have_cons_add = np.array([]) else: # does not find any active neuron Masksb_add = [] masks_add = [] times_add = times_uniques area_add = np.array([]) have_cons_add = np.array([]) tuple_add = (Masksb_add, masks_add, times_add, area_add, have_cons_add) # temporal merging 5: merge newly found neurons within the recent "merge_every" frames with existing neurons if ((t - 3 - t_merge) == merge_every) or (t == nframesf - 1): tuple_temp = merge_2(tuple_temp, tuple_add, dims, Params_post) t_merge += merge_every if show_intermediate: Masks_2 = select_cons(tuple_temp) current_frame += 1 # Update the stored latest frames when it runs out: move them "leng_tf" ahead if current_frame > leng_past: current_frame = leng_tf + 1 past_frames[:leng_tf] = past_frames[-leng_tf:] if display: end_frame = time.time() list_time_per[t] = end_frame - start_frame if t % 1000 == 0: print('{} frames has been processed'.format(t)) if not show_intermediate: Masks_2 = select_cons(tuple_temp) # final result. Masks_2 is a 2D sparse matrix of the segmented neurons if len(Masks_2): Masks_2 = sparse.vstack(Masks_2) else: Masks_2 = sparse.csc_matrix((0, dims[0] * dims[1])) if display: end_online = time.time() time_online = end_online - start_online time_frame_online = time_online / (nframesf - frames_initf) * 1000 print('Online time: {:6f} s, {:6f} ms/frame'.format( time_online, time_frame_online)) # Save total processing time, and average processing time per frame if display: end_final = time.time() time_all = end_final - start_init time_frame_all = time_all / nframes * 1000 print('Total time: {:6f} s, {:6f} ms/frame'.format( time_all, time_frame_all)) time_total = np.array([time_init, time_online, time_all]) time_frame = np.array( [time_frame_init, time_frame_online, time_frame_all]) else: time_total = np.zeros((3, )) time_frame = np.zeros((3, )) # convert to a 3D array of the segmented neurons Masks = np.reshape(Masks_2.toarray(), (Masks_2.shape[0], Lx, Ly)).astype('bool') return Masks, Masks_2, time_total, time_frame
def suns_online_track(filename_video, filename_CNN, Params_pre, Params_post, dims, \ frames_init, merge_every, batch_size_init=1, useSF=True, useTF=True, useSNR=True, \ med_subtract=False, update_baseline=False, \ useWT=False, prealloc=True, display=True, useMP=True, p=None): '''The complete SUNS online procedure with tracking. It uses the trained CNN model from "filename_CNN" and the optimized hyper-parameters in "Params_post" to process the video "Exp_ID" in "dir_video" Inputs: filename_video (str): The path of the file of the input raw video. The file must be a ".h5" file, with dataset "mov" being the input video (shape = (T0,Lx0,Ly0)). filename_CNN (str): The path of the trained CNN model. Params_pre (dict): Parameters for pre-processing. Params_pre['gauss_filt_size'] (float): The standard deviation of the spatial Gaussian filter in pixels Params_pre['Poisson_filt'] (1D numpy.ndarray of float32): The temporal filter kernel Params_pre['num_median_approx'] (int): Number of frames used to compute the median and median-based standard deviation Params_pre['nn'] (int): Number of frames at the beginning of the video to be processed. The remaining video is not considered a part of the input video. Params_post (dict): Parameters for post-processing. Params_post['minArea']: Minimum area of a valid neuron mask (unit: pixels). Params_post['avgArea']: The typical neuron area (unit: pixels). Params_post['thresh_pmap']: The probablity threshold. Values higher than thresh_pmap are active pixels. It is stored in uint8, so it should be converted to float32 before using. Params_post['thresh_mask']: Threashold to binarize the real-number mask. Params_post['thresh_COM0']: Threshold of COM distance (unit: pixels) used for the first COM-based merging. Params_post['thresh_COM']: Threshold of COM distance (unit: pixels) used for the second COM-based merging. Params_post['thresh_IOU']: Threshold of IOU used for merging neurons. Params_post['thresh_consume']: Threshold of consume ratio used for merging neurons. Params_post['cons']: Minimum number of consecutive frames that a neuron should be active for. dims (tuplel of int, shape = (2,)): lateral dimension of the raw video. frames_init (int): Number of frames used for initialization. merge_every (int): SUNS online merge the newly segmented frames every "merge_every" frames. batch_size_init (int, default to 1): batch size of CNN inference for initialization frames. useSF (bool, default to True): True if spatial filtering is used. useTF (bool, default to True): True if temporal filtering is used. useSNR (bool, default to True): True if pixel-by-pixel SNR normalization filtering is used. med_subtract (bool, default to False): True if the spatial median of every frame is subtracted before temporal filtering. Can only be used when spatial filtering is not used. update_baseline (bool, default to False): True if the median and median-based std is updated every "frames_init" frames. useWT (bool, default to False): Indicator of whether watershed is used. prealloc (bool, default to True): True if pre-allocate memory space for large variables. Achieve faster speed at the cost of higher memory occupation. display (bool, default to True): Indicator of whether to show intermediate information useMP (bool, defaut to True): indicator of whether multiprocessing is used to speed up. p (multiprocessing.Pool, default to None): Outputs: Masks (3D numpy.ndarray of bool, shape = (n,Lx0,Ly0)): the final segmented masks. Masks_2 (scipy.csr_matrix of bool, shape = (n,Lx0*Ly0)): the final segmented masks in the form of sparse matrix. time_total (list of float, shape = (3,)): the total time spent for initalization, online processing, and total processing time_frame (list of float, shape = (3,)): the average time spent on every frame for initalization, online processing, and total processing ''' if display: start = time.time() (Lx, Ly) = dims # zero-pad the lateral dimensions to multiples of 8, suitable for CNN rowspad = math.ceil(Lx / 8) * 8 colspad = math.ceil(Ly / 8) * 8 dimspad = (rowspad, colspad) Poisson_filt = Params_pre['Poisson_filt'] gauss_filt_size = Params_pre['gauss_filt_size'] nn = Params_pre['nn'] leng_tf = Poisson_filt.size leng_past = 2 * leng_tf # number of past frames stored for temporal filtering list_time_per = np.zeros(nn) # Load CNN model fff = get_shallow_unet() fff.load_weights(filename_CNN) # run CNN inference once to warm up init_imgs = np.zeros((batch_size_init, rowspad, colspad, 1), dtype='float32') init_masks = np.zeros((batch_size_init, rowspad, colspad, 1), dtype='uint8') fff.evaluate(init_imgs, init_masks, batch_size=batch_size_init) del init_imgs, init_masks # load optimal post-processing parameters minArea = Params_post['minArea'] avgArea = Params_post['avgArea'] # thresh_pmap = Params_post['thresh_pmap'] thresh_mask = Params_post['thresh_mask'] # thresh_COM0 = Params_post['thresh_COM0'] # thresh_COM = Params_post['thresh_COM'] thresh_IOU = Params_post['thresh_IOU'] thresh_consume = Params_post['thresh_consume'] # cons = Params_post['cons'] # thresh_pmap_float = (Params_post['thresh_pmap']+1.5)/256 thresh_pmap_float = (Params_post['thresh_pmap'] + 1) / 256 # for published version # Spatial filtering preparation if useSF == True: # lateral dimensions slightly larger than the raw video but faster for FFT rows1 = cv2.getOptimalDFTSize(rowspad) cols1 = cv2.getOptimalDFTSize(colspad) # if the learned 2D and 3D wisdom files have been saved, load them. # Otherwise, learn wisdom later Length_data2 = str((rows1, cols1)) cc2 = load_wisdom_txt('wisdom\\' + Length_data2) Length_data3 = str((frames_init, rows1, cols1)) cc3 = load_wisdom_txt('wisdom\\' + Length_data3) if cc3: pyfftw.import_wisdom(cc3) # mask for spatial filter mask2 = plan_mask2(dims, (rows1, cols1), gauss_filt_size) # FFT planning (bb, bf, fft_object_b, fft_object_c) = plan_fft(frames_init, (rows1, cols1), prealloc) else: (mask2, bf, fft_object_b, fft_object_c) = (None, None, None, None) bb = np.zeros((frames_init, rowspad, colspad), dtype='float32') # Temporal filtering preparation frames_initf = frames_init - leng_tf + 1 if useTF == True: if prealloc: # past frames stored for temporal filtering past_frames = np.ones((leng_past, rowspad, colspad), dtype='float32') else: past_frames = np.zeros((leng_past, rowspad, colspad), dtype='float32') else: past_frames = None if prealloc: # Pre-allocate memory for some future variables med_frame2 = np.ones((rowspad, colspad, 2), dtype='float32') video_input = np.ones((frames_initf, rowspad, colspad), dtype='float32') pmaps_b_init = np.ones((frames_initf, Lx, Ly), dtype='uint8') frame_SNR = np.ones(dimspad, dtype='float32') pmaps_b = np.ones(dims, dtype='uint8') if update_baseline: video_tf_past = np.ones((frames_init, rowspad, colspad), dtype='float32') else: med_frame2 = np.zeros((rowspad, colspad, 2), dtype='float32') video_input = np.zeros((frames_initf, rowspad, colspad), dtype='float32') pmaps_b_init = np.zeros((frames_initf, Lx, Ly), dtype='uint8') frame_SNR = np.zeros(dimspad, dtype='float32') pmaps_b = np.zeros(dims, dtype='uint8') if update_baseline: video_tf_past = np.zeros((frames_init, rowspad, colspad), dtype='float32') if display: time_init = time.time() print('Parameter initialization time: {} s'.format(time_init - start)) # %% Load raw video h5_img = h5py.File(filename_video, 'r') video_raw = np.array(h5_img['mov']) h5_img.close() nframes = video_raw.shape[0] nframesf = nframes - leng_tf + 1 bb[:, :Lx, :Ly] = video_raw[:frames_init] if display: time_load = time.time() print('Load data: {} s'.format(time_load - time_init)) # %% Actual processing starts after the video is loaded into memory # Initialization using the first "frames_init" frames print('Initialization of algorithms using the first {} frames'.format( frames_init)) if display: start_init = time.time() med_frame3, segs_all, recent_frames = init_online( bb, dims, video_input, pmaps_b_init, fff, thresh_pmap_float, Params_post, \ med_frame2, mask2, bf, fft_object_b, fft_object_c, Poisson_filt, \ useSF=useSF, useTF=useTF, useSNR=useSNR, med_subtract=med_subtract, \ useWT=useWT, batch_size_init=batch_size_init, p=p) if useTF == True: past_frames[:leng_tf] = recent_frames tuple_temp = merge_complete(segs_all[:frames_initf], dims, Params_post) # Initialize Online track variables (Masksb_temp, masks_temp, times_temp, area_temp, have_cons_temp) = tuple_temp # list of previously found neurons that satisfy consecutive frame requirement Masks_cons = select_cons(tuple_temp) # sparse matrix of previously found neurons that satisfy consecutive frame requirement Masks_cons_2D = sparse.vstack(Masks_cons) # indices of previously found neurons that satisfy consecutive frame requirement ind_cons = have_cons_temp.nonzero()[0] segs0 = segs_all[0] # segs of initialization frames # segs if no neurons are found segs_empty = (segs0[0][0:0], segs0[1][0:0], segs0[2][0:0], segs0[3][0:0]) # Number of previously found neurons that satisfy consecutive frame requirement N1 = len(Masks_cons) # list of "segs" for neurons that are not previously found list_segs_new = [] # list of newly segmented masks for old neurons (segmented in previous frames) list_masks_old = [[] for _ in range(N1)] # list of the newly active indices of frames of old neurons times_active_old = [[] for _ in range(N1)] # True if the old neurons are active in the previous frame active_old_previous = np.zeros(N1, dtype='bool') if display: end_init = time.time() time_init = end_init - start_init time_frame_init = time_init / (frames_initf) * 1000 print('Initialization time: {:6f} s, {:6f} ms/frame'.format( time_init, time_frame_init)) if display: start_online = time.time() # Spatial filtering preparation for online processing. # Attention: this part counts to the total time if useSF: if cc2: pyfftw.import_wisdom(cc2) (bb, bf, fft_object_b, fft_object_c) = plan_fft2((rows1, cols1)) else: (bf, fft_object_b, fft_object_c) = (None, None, None) bb = np.zeros(dimspad, dtype='float32') print('Start frame by frame processing') # %% Online processing for the following frames current_frame = leng_tf + 1 t_merge = frames_initf for t in range(frames_initf, nframesf): if display: start_frame = time.time() # load the current frame bb[:Lx, :Ly] = video_raw[t + leng_tf - 1] bb[Lx:, :] = 0 bb[:, Ly:] = 0 # PreProcessing frame_SNR, frame_tf = preprocess_online(bb, dimspad, med_frame3, frame_SNR, \ past_frames[current_frame-leng_tf:current_frame], mask2, bf, fft_object_b, fft_object_c, \ Poisson_filt, useSF=useSF, useTF=useTF, useSNR=useSNR, \ med_subtract=med_subtract, update_baseline=update_baseline) if update_baseline: t_past = (t - frames_initf) % frames_init video_tf_past[t_past] = frame_tf if t_past == frames_init - 1: # update median and median-based standard deviation every "frames_init" frames if useSNR: med_frame3 = SNR_normalization(video_tf_past, med_frame2, (rowspad, colspad), 1, display=False) else: med_frame3 = median_normalization(video_tf_past, med_frame2, (rowspad, colspad), 1, display=False) # CNN inference frame_prob = CNN_online(frame_SNR, fff, dims) # first step of post-processing segs = separate_neuron_online(frame_prob, pmaps_b, thresh_pmap_float, minArea, avgArea, useWT) active_old = np.zeros( N1, dtype='bool' ) # True if the old neurons are active in the current frame masks_t, neuronstate_t, cents_t, areas_t = segs N2 = neuronstate_t.size if N2: # Try to merge the new masks to old neurons new_found = np.zeros(N2, dtype='bool') for n2 in range(N2): masks_t2 = masks_t[n2] cents_t2 = np.round(cents_t[n2, 1]) * Ly + np.round(cents_t[n2, 0]) # If a new masks belongs to an old neuron, the COM of the new mask must be inside the old neuron area. # Select possible old neurons that the new mask can merge to possible_masks1 = Masks_cons_2D[:, cents_t2].nonzero()[0] IOUs = np.zeros(len(possible_masks1)) areas_t2 = areas_t[n2] for (ind, n1) in enumerate(possible_masks1): # Calculate IoU and consume ratio to determine merged neurons area_i = Masks_cons[n1].multiply(masks_t2).nnz area_temp1 = area_temp[n1] area_u = area_temp1 + areas_t2 - area_i IOU = area_i / area_u consume = area_i / min(area_temp1, areas_t2) contain = (IOU >= thresh_IOU) or (consume >= thresh_consume) if contain: # merging criterion satisfied IOUs[ind] = IOU num_contains = IOUs.nonzero()[0].size if num_contains: # The new mask can merge to one of the old neurons. # If there are multiple candicates, choose the one with the highest IoU belongs = possible_masks1[IOUs.argmax()] # merge the mask and active frame index list_masks_old[belongs].append(masks_t2) times_active_old[belongs].append(t + frames_initf) # This old neurons is active in the current frame active_old[belongs] = True else: # The new mask can not merge to any old neuron. new_found[n2] = True if np.any( new_found ): # There are some new masks that can not merge to old neurons segs_new = (masks_t[new_found], neuronstate_t[new_found], cents_t[new_found], areas_t[new_found]) else: # All masks already merged to old neurons segs_new = segs_empty else: # No neurons fould in the current frame segs_new = segs list_segs_new.append(segs_new) if (t + 1 - t_merge) != merge_every or t == (nframesf - 1): # Update the old neurons with new appearances in the current frame. if t < (nframesf - 1): # True if the neurons are active in the previous frame but not active in the current frame inactive = np.logical_and( active_old_previous, np.logical_not(active_old)).nonzero()[0] else: # last frame # All active neurons should be updated, so they are treated as inactive in the next frame inactive = active_old_previous.nonzero()[0] # Update the indicators of the previous frame using the current frame active_old_previous = active_old.copy() for n1 in inactive: # merge new active frames to existing active frames for already found neurons # n1 is the index in the old neurons that satisfy consecutive frame requirement. # n10 is the index in all old neurons. n10 = ind_cons[n1] # Add all the new masks to the overall real-number masks mask_update = masks_temp[n10] + sum(list_masks_old[n1]) masks_temp[n10] = mask_update # Add indices of active frames times_add = np.unique(np.array(times_active_old[n1])) times_temp[n10] = np.hstack([times_temp[n10], times_add]) # reset lists used to store the information from new frames related to old neurons list_masks_old[n1] = [] times_active_old[n1] = [] # update the binary masks and areas Maskb_update = mask_update >= mask_update.max() * thresh_mask Masksb_temp[n10] = Maskb_update Masks_cons[n1] = Maskb_update area_temp[n10] = Maskb_update.nnz if inactive.size: Masks_cons_2D = sparse.vstack(Masks_cons) if (t + 1 - t_merge) == merge_every or t == (nframesf - 1): if t < (nframesf - 1): # delay merging new frame to next frame by assuming all the neurons active in the previous frame # are still active in the current frame, to reserve merging time for new neurons active_old_previous = np.logical_or(active_old_previous, active_old) # merge new neurons with old masks that do not satisfy consecutive frame requirement tuple_temp = (Masksb_temp, masks_temp, times_temp, area_temp, have_cons_temp) # merge the remaining new masks from the most recent "merge_every" frames tuple_add = merge_complete(list_segs_new, dims, Params_post) (Masksb_add, masks_add, times_add, area_add, have_cons_add) = tuple_add # update the indices of active frames times_add = [x + merge_every for x in times_add] tuple_add = (Masksb_add, masks_add, times_add, area_add, have_cons_add) # merge the remaining new masks with the existing masks that do not satisfy consecutive frame requirement tuple_temp = merge_2_nocons(tuple_temp, tuple_add, dims, Params_post) (Masksb_temp, masks_temp, times_temp, area_temp, have_cons_temp) = tuple_temp # Update the indices of old neurons that satisfy consecutive frame requirement ind_cons_new = have_cons_temp.nonzero()[0] for (ind, ind_cons_0) in enumerate(ind_cons_new): if ind_cons_0 not in ind_cons: # update lists used to store the information from new frames related to old neurons if ind_cons_0 > ind_cons.max(): list_masks_old.append([]) times_active_old.append([]) else: list_masks_old.insert(ind, []) times_active_old.insert(ind, []) # Update the list of previously found neurons that satisfy consecutive frame requirement Masks_cons = select_cons(tuple_temp) Masks_cons_2D = sparse.vstack(Masks_cons) N1 = len(Masks_cons) list_segs_new = [] # Update whether the old neurons are active in the previous frame active_old_previous = np.zeros_like(have_cons_temp) active_old_previous[ind_cons] = active_old active_old_previous = active_old_previous[ind_cons_new] ind_cons = ind_cons_new t_merge += merge_every current_frame += 1 # Update the stored latest frames when it runs out: move them "leng_tf" ahead if current_frame > leng_past: current_frame = leng_tf + 1 past_frames[:leng_tf] = past_frames[-leng_tf:] if display: end_frame = time.time() list_time_per[t] = end_frame - start_frame if t % 1000 == 0: print('{} frames has been processed'.format(t)) Masks_cons = select_cons(tuple_temp) # final result. Masks_2 is a 2D sparse matrix of the segmented neurons if len(Masks_cons): Masks_2 = sparse.vstack(Masks_cons) else: Masks_2 = sparse.csc_matrix((0, dims[0] * dims[1])) if display: end_online = time.time() time_online = end_online - start_online time_frame_online = time_online / (nframesf - frames_initf) * 1000 print('Online time: {:6f} s, {:6f} ms/frame'.format( time_online, time_frame_online)) # Save total processing time, and average processing time per frame if display: end_final = time.time() time_all = end_final - start_init time_frame_all = time_all / nframes * 1000 print('Total time: {:6f} s, {:6f} ms/frame'.format( time_all, time_frame_all)) time_total = np.array([time_init, time_online, time_all]) time_frame = np.array( [time_frame_init, time_frame_online, time_frame_all]) else: time_total = np.zeros((3, )) time_frame = np.zeros((3, )) # convert to a 3D array of the segmented neurons Masks = np.reshape(Masks_2.toarray(), (Masks_2.shape[0], Lx, Ly)).astype('bool') return Masks, Masks_2, time_total, time_frame
def preprocess_video_online(dir_video:str, Exp_ID:str, Params:dict, frames_init=10**6, dir_network_input:str = None, useSF=False, useTF=True, useSNR=True, \ med_subtract=False, update_baseline=False, prealloc=True, display=True): '''Pre-process the registered video "Exp_ID" in "dir_video" into an SNR video "network_input". The process includes spatial filter, temporal filter, and SNR normalization. Each step is optional. The function does some preparations, including pre-allocating memory space (optional), FFT planing, and setting filter kernels, before loading the video. Inputs: dir_video (str): The folder containing the input video. Each file must be a ".h5" file, with dataset "mov" being the input video (shape = (T0,Lx0,Ly0)). Exp_ID (str): The filer name of the input video. Params (dict): Parameters for pre-processing. Params['gauss_filt_size'] (float): The standard deviation of the spatial Gaussian filter in pixels Params['Poisson_filt'] (1D numpy.ndarray of float32): The temporal filter kernel Params['num_median_approx'] (int): Number of frames used to compute the median and median-based standard deviation Params['nn'] (int): Number of frames at the beginning of the video to be processed. The remaining video is not considered a part of the input video. frames_init (int, default to a very large number): Median and median-based standard deviation are calculate from the first "frames_init" frames (before temporal filtering) of the video dir_network_input (str, default to None): The folder to save the SNR video (network_input) in hard drive. If dir_network_input == None, then the SNR video is not stored in hard drive useSF (bool, default to True): True if spatial filtering is used. useTF (bool, default to True): True if temporal filtering is used. useSNR (bool, default to True): True if pixel-by-pixel SNR normalization filtering is used. med_subtract (bool, default to False): True if the spatial median of every frame is subtracted before temporal filtering. Can only be used when spatial filtering is not used. update_baseline (bool, default to False): True if the median and median-based std is updated every "frames_init" frames. prealloc (bool, default to True): True if pre-allocate memory space for large variables. Achieve faster speed at the cost of higher memory occupation. Outputs: network_input (numpy.ndarray of float32, shape = (T,Lx,Ly)): the SNR video obtained after pre-processing. The shape is (T,Lx,Ly), where T is shorter than T0 due to temporal filtering, and Lx and Ly are Lx0 and Ly0 padded to multiples of 8, so that the images can be process by the shallow U-Net. start (float): the starting time of the pipline (after the data is loaded into memory) In addition, the SNR video is saved in "dir_network_input" as a "(Exp_ID).h5" file containing a dataset "network_input". ''' nn = Params['nn'] if useSF: gauss_filt_size = Params['gauss_filt_size'] Poisson_filt = Params['Poisson_filt'] # num_median_approx = Params['num_median_approx'] h5_video = os.path.join(dir_video, Exp_ID + '.h5') h5_file = h5py.File(h5_video, 'r') (nframes, rows, cols) = h5_file['mov'].shape # Make the lateral number of pixels a multiple of 8, so that the CNN can process them rowspad = math.ceil(rows / 8) * 8 colspad = math.ceil(cols / 8) * 8 # Only keep the first "nn" frames to process nframes = min(nframes, nn) if useSF: # lateral dimensions slightly larger than the raw video but faster for FFT rows1 = cv2.getOptimalDFTSize(rows) cols1 = cv2.getOptimalDFTSize(cols) if display: print(rows, cols, nn, '->', rows1, cols1, nn) start_plan = time.time() # if the learned wisdom files have been saved, load them. Otherwise, learn wisdom later Length_data = str((nn, rows1, cols1)) cc = load_wisdom_txt('wisdom\\' + Length_data) if cc: pyfftw.import_wisdom(cc) # FFT planning bb, bf, fft_object_b, fft_object_c = plan_fft(nn, (rows1, cols1), prealloc) if display: end_plan = time.time() print('FFT planning: {} s'.format(end_plan - start_plan)) # %% Initialization: Calculate the spatial filter and set variables. mask2 = plan_mask2((rows, cols), (rows1, cols1), gauss_filt_size) else: bb = np.zeros((nframes, rowspad, colspad), dtype='float32') fft_object_b = None fft_object_c = None bf = None end_plan = time.time() mask2 = None # %% Initialization: Set variables. if useTF: leng_tf = Poisson_filt.size nframesf = nframes - leng_tf + 1 # Number of frames after temporal filtering else: nframesf = nframes frames_initf = frames_init - leng_tf + 1 if prealloc: network_input = np.ones((nframesf, rowspad, colspad), dtype='float32') med_frame2 = np.ones((rowspad, colspad, 2), dtype='float32') else: network_input = np.zeros((nframesf, rowspad, colspad), dtype='float32') med_frame2 = np.zeros((rowspad, colspad, 2), dtype='float32') # median_decimate = max(1,nframes//num_median_approx) if display: end_init = time.time() print('Initialization: {} s'.format(end_init - end_plan)) # %% Load the raw video into "bb" for t in range(nframes): # use this one to save memory bb[t, :rows, :cols] = np.array(h5_file['mov'][t]) h5_file.close() if display: end_load = time.time() print('data loading: {} s'.format(end_load - end_init)) start = time.time( ) # The pipline starts after the video is loaded into memory network_input = preprocess_complete_online(bb, (rowspad,colspad), network_input, med_frame2, \ Poisson_filt, mask2, bf, fft_object_b, fft_object_c, useSF, useTF, useSNR, \ med_subtract, update_baseline, frames_initf, frames_init, prealloc, display) if display: end = time.time() print('total per frame: {} ms'.format((end - start) / nframes * 1000)) # if "dir_network_input" is not None, save network_input to an ".h5" file if dir_network_input: f = h5py.File(os.path.join(dir_network_input, Exp_ID + ".h5"), "w") f.create_dataset("network_input", data=network_input) f.close() if display: end_saving = time.time() print('Network_input saving: {} s'.format(end_saving - end)) return network_input, start