epicure.tracking

   1from qtpy.QtWidgets import QVBoxLayout, QWidget # type: ignore
   2from epicure.laptrack_centroids import LaptrackCentroids
   3import epicure.Utils as ut
   4laptrack_over = False
   5try:    
   6    from epicure.laptrack_overlaps import LaptrackOverlaps
   7    laptrack_over = True
   8except ImportError:
   9    print("Laptrack overlap not available in your laptrack version. Only the centroid option will be proposed. Update laptrack to 0.16 to have it")
  10    pass
  11import laptrack
  12if ut.version_above(laptrack, "0.16"):
  13    try:    
  14        from laptrack.data_conversion import split_merge_df_to_napari_graph as to_napari_graph# type: ignore
  15    except ImportError:
  16        from laptrack.data_conversion import convert_split_merge_df_to_napari_graph as to_napari_graph # type: ignore
  17else:
  18    try:    
  19        from laptrack.data_conversion import convert_split_merge_df_to_napari_graph as to_napari_graph # type: ignore
  20    except ImportError:
  21        from laptrack.data_conversion import split_merge_df_to_napari_graph as to_napari_graph # type: ignore
  22from napari.utils import progress # type: ignore
  23from skimage.transform import warp
  24from skimage.registration import optical_flow_ilk
  25import pandas as pd
  26import numpy as np
  27import scipy.ndimage as ndi
  28import epicure.epiwidgets as wid
  29from joblib import Parallel, delayed
  30
  31class Tracking(QWidget):
  32    """
  33        Handles tracking of cells, track operations with the Tracks layer
  34    """
  35    def __init__(self, napari_viewer, epic):
  36        super().__init__()
  37        self.viewer = napari_viewer
  38        self.epicure = epic
  39        self.graph = None      ## init 
  40        self.tracklayer = None      ## track layer with information (centroids, labels, tree..)
  41        self.track_data = None ## keep the updated data, and update the layer only from time to time (slow to do)
  42        self.tracklayer_name = "Tracks"  ## name of the layer containing tracks
  43        self.nframes = self.epicure.nframes
  44        self.properties = ["label", "centroid"]
  45
  46        layout = QVBoxLayout()
  47        
  48        ## Add update track button 
  49        self.track_update = wid.add_button( "Update tracks display", self.update_track_layer, "Update the Track layer with the changements made since the last update" )
  50        layout.addWidget(self.track_update)
  51        
  52        ## Correct track button 
  53        #track_reset = wid.add_button( "Correct track data", self.reset_tracks, "Correct the track data after some track was lost" )
  54        #layout.addWidget(track_reset)
  55
  56        ## Method specific
  57        track_method, self.track_choice = wid.list_line( "Tracking method", "Choose the tracking method to use and display its parameter", func=None )
  58        layout.addWidget(self.track_choice)
  59        
  60        self.track_choice.addItem("Laptrack-Centroids")
  61        self.create_laptrack_centroids()
  62        layout.addWidget(self.gLapCentroids)
  63
  64        if laptrack_over: 
  65            self.track_choice.addItem("Laptrack-Overlaps")
  66            self.create_laptrack_overlap()
  67            layout.addWidget(self.gLapOverlap)
  68        else:
  69            self.min_iou = None
  70            self.split_cost = None
  71            self.merg_cost = None
  72
  73        drift_layout, self.drift_correction, self.drift_radius = wid.check_value( check="With drift correction", checked=False, value=str(50), descr="Taking into account local drift in tracking calculations") 
  74        layout.addLayout( drift_layout )
  75        
  76        self.track_go = wid.add_button( "Track", self.do_tracking, "Launch the tracking with the current parameter. Can take time" )
  77        layout.addWidget(self.track_go)
  78        self.setLayout(layout)
  79
  80        ## General tracking options
  81        frame_line, self.frame_range, self.range_group = wid.checkgroup_help( "Track only some frames", False, "Option to track only a given range of frames", None ) 
  82        self.frame_range.clicked.connect( self.show_frame_range )
  83        range_layout = QVBoxLayout()
  84        ntrack, self.start_frame = wid.ranged_value_line( "Track from frame:", 0, self.nframes-1, 1, 0, "Set first frame to begin tracking" )
  85        range_layout.addLayout(ntrack)
  86        
  87        entrack, self.end_frame = wid.ranged_value_line( "Until frame:", 1, self.nframes-1, 1, self.nframes-1, "Set the last frame unitl which to track" )
  88        range_layout.addLayout(entrack)
  89        self.start_frame.valueChanged.connect( self.changed_start )
  90        self.end_frame.valueChanged.connect( self.changed_end )
  91        
  92        self.range_group.setLayout( range_layout )
  93        layout.addWidget( self.frame_range )
  94        layout.addWidget( self.range_group )
  95        
  96        self.show_frame_range()
  97        self.show_trackoptions()
  98        self.track_choice.currentIndexChanged.connect(self.show_trackoptions)
  99        
 100
 101    def show_frame_range( self ):
 102        """ Show/Hide frame range options """
 103        self.range_group.setVisible( self.frame_range.isChecked() )
 104        
 105    #### settings
 106
 107    def get_current_settings( self ):
 108        """ Get current settings to save as preferences """
 109        settings = {}
 110        settings["Track method"] = self.track_choice.currentText() 
 111        settings["Add feat"] = self.check_penalties.isChecked()
 112        settings["Max distance"] = self.max_dist.text()
 113        settings["Splitting cost"] = self.splitting_cost.text()
 114        settings["Merging cutoff"] = self.merging_cost.text()
 115        settings["Min IOU"] = self.min_iou.text()
 116        settings["Over split"] = self.split_cost.text()
 117        settings["Over merge"] = self.merg_cost.text()
 118        return settings
 119
 120    def apply_settings( self, settings ):
 121        """ Set the parameters/current display from the prefered settings """
 122        for setty, val in settings.items():
 123            if setty == "Track method":
 124                self.track_choice.setCurrentText( val )
 125            if setty == "Add feat":
 126                self.check_penalties.setChecked( val )
 127            if setty == "Max distance":
 128                self.max_dist.setText( val )
 129            if setty == "Splitting cost":
 130                self.splitting_cost.setText( val )
 131            if setty == "Merging cutoff":
 132                self.merging_cost.setText( val )
 133            if laptrack_over:
 134                if setty == "Min IOU":
 135                    self.min_iou.setText( val )
 136                if setty == "Over split":
 137                    self.split_cost.setText( val )
 138                if setty == "Over merge":
 139                    self.merg_cost.setText( val )
 140            
 141    ##########################################
 142    #### Tracks layer and function
 143
 144    def reset( self ):
 145        """ Reset Tracks layer and data """
 146        self.graph = None
 147        self.track_data = None
 148        ut.remove_layer( self.viewer, "Tracks" )
 149
 150    def init_tracks(self):
 151        """ Add a track layer with the new tracks """
 152        track_table, track_prop = self.create_tracks()
 153        
 154        ## plot tracks
 155        if len(track_table) > 0:
 156            self.clear_graph()
 157            self.viewer.add_tracks(
 158                track_table,
 159                graph=self.graph, 
 160                name=self.tracklayer_name,
 161                properties = track_prop,
 162                scale = self.viewer.layers["Segmentation"].scale,
 163                )
 164            self.viewer.layers[self.tracklayer_name].visible=True
 165            self.viewer.layers[self.tracklayer_name].color_by="track_id"
 166            ut.set_active_layer(self.viewer, "Segmentation")
 167            self.tracklayer = self.viewer.layers[self.tracklayer_name]
 168            self.track_data = self.tracklayer.data
 169            #self.track.display_id = True
 170            self.color_tracks_as_labels()
 171
 172    def color_tracks_as_labels(self):
 173        """ Color the tracks the same as the label layer """
 174        ## must color it manually by getting the Label layer colors for each track_id
 175        cols = np.zeros((len(self.tracklayer.data[:,0]),4))
 176        for i, tr in enumerate(self.tracklayer.data[:,0]):
 177            cols[i] = (self.epicure.seglayer.get_color(tr))
 178        self.tracklayer._track_colors = cols
 179        self.tracklayer.events.color_by()
 180    
 181    def color_tracks_by_lineage(self):
 182        """ Color the tracks by their lineage (daughters same colors as parents) """
 183        ## must color it manually by getting the Label layer colors for each track_id
 184        cols = np.zeros((len(self.tracklayer.data[:,0]),4))
 185        for i, tr in enumerate(self.tracklayer.data[:,0]):
 186            ## find the parent cell,n going up the tree until no more parent
 187            while tr in self.graph.keys():
 188                tr = self.graph_parent( tr )
 189            cols[i] = (self.epicure.seglayer.get_color(tr))
 190        self.tracklayer._track_colors = cols
 191        self.tracklayer.events.color_by()
 192
 193    def graph_parent( self, ind ):
 194        """ Get the value of the parent from the graph """
 195        if ind not in self.graph.keys():
 196            return None
 197        if isinstance(self.graph[ind], list):
 198            return self.graph[ind][0]
 199        return self.graph[ind]
 200
 201    def replace_tracks(self, track_df):
 202        """ Replace all tracks based on the dataframe """
 203        if not self.undrifted and self.drift_correction.isChecked():
 204            ## recalculate the label centroids as it was corrected for drift
 205            track_table, track_prop = self.create_tracks()
 206        else:
 207            track_table, track_prop = self.build_tracks( track_df )
 208        self.tracklayer.data = track_table
 209        self.track_data = self.tracklayer.data
 210        self.tracklayer.properties = track_prop
 211        self.tracklayer.refresh()
 212        self.color_tracks_as_labels()
 213
 214    def reset_tracks(self):
 215        """ Reset tracks and reload them from labels """
 216        ut.remove_layer(self.viewer, "Tracks")
 217        self.init_tracks()
 218
 219    def update_track_layer(self):
 220        """ Update the track layer (slow) """
 221        self.viewer.window._status_bar._toggle_activity_dock(True)
 222        progress_bar = progress(total=1)
 223        progress_bar.set_description( "Updating track layer" )
 224        self.tracklayer.data = self.track_data
 225        progress_bar.close()
 226        self.color_tracks_as_labels()
 227        self.viewer.window._status_bar._toggle_activity_dock(False)
 228
 229    def measure_intensity_features( self, feat, intimg=None, frames=None ):
 230        """ Measure mean value of a feature in a track """
 231        if ( intimg is not None ):
 232            if frames is None:
 233                tracks = self.get_track_list()
 234                seg = self.epicure.seg
 235                iimg = intimg
 236            else:
 237                tracks = self.get_tracks_list_frames( frames )
 238                seg = self.epicure.seg[frames]
 239                iimg = intimg[frames]
 240        if feat == "intensity_mean":
 241            mean_intensities = ndi.mean( iimg, seg, tracks )
 242            return tracks, mean_intensities
 243        if feat == "intensity_sum":
 244            sum_intensities = ndi.sum( iimg, seg, tracks )
 245            return tracks, sum_intensities
 246        if feat == "intensity_max":
 247            sum_intensities = ndi.maximum( iimg, seg, tracks )
 248            return tracks, sum_intensities
 249        if feat == "intensity_min":
 250            sum_intensities = ndi.minimum( iimg, seg, tracks )
 251            return tracks, sum_intensities
 252        if feat == "intensity_median":
 253            sum_intensities = ndi.median( iimg, seg, tracks )
 254            return tracks, sum_intensities
 255        print( "Mean feature on track not implemented" )
 256        return None
 257
 258    def measure_track_features( self, track_id, scaling=False ):
 259        """ Measure features (length, total displacement...) of given track """
 260        features = {}
 261        track = self.get_track_data( track_id )
 262        if track.shape[0] == 0:
 263            return features
 264        track = track[track[:,1].argsort()]
 265        start = int(np.min(track[:,1]))
 266        end = int(np.max(track[:,1]))
 267        temp_unit = ""
 268        vel_unit = ""
 269        disp_unit = ""
 270        temp_scale = 1
 271        vel_scale = 1
 272        disp_scale = 1
 273        if scaling:
 274            temp_unit = "_"+self.epicure.epi_metadata["UnitT"]
 275            vel_unit = "_"+self.epicure.epi_metadata["UnitXY"]+"/"+self.epicure.epi_metadata["UnitT"]
 276            disp_unit = "_"+self.epicure.epi_metadata["UnitXY"]
 277            temp_scale = self.epicure.epi_metadata["ScaleT"]
 278            vel_scale = self.epicure.epi_metadata["ScaleXY"]/self.epicure.epi_metadata["ScaleT"]
 279            disp_scale = self.epicure.epi_metadata["ScaleXY"]
 280        features["Label"] = track_id
 281        features["TrackDuration"+temp_unit] = (end - start + 1)*temp_scale
 282        features["TrackStart"+temp_unit] = start * temp_scale
 283        features["TrackEnd"+temp_unit] = end * temp_scale
 284        features["NbGaps"] = end - start + 1 - len(track)
 285        if (end-start) == 0:
 286            ## only one frame
 287            features["TotalDisplacement"+disp_unit] = None
 288            features["NetDisplacement"+disp_unit] = None
 289            features["Straightness"] = None
 290            features["MeanVelocity"+vel_unit] = None
 291        else:
 292            features["TotalDisplacement"+disp_unit] = ut.total_distance( track[:,2:4] ) * disp_scale
 293            features["NetDisplacement"+disp_unit] = ut.net_distance( track[:,2:4] ) * disp_scale
 294            features["MeanVelocity"+vel_unit] = np.mean( ut.velocities( track[:,1:4] ) ) * vel_scale 
 295            if features["TotalDisplacement"+disp_unit] > 0:
 296                features["Straightness"] = features["NetDisplacement"+disp_unit]/features["TotalDisplacement"+disp_unit]
 297            else:
 298                features["Straightness"] = None
 299        return features
 300
 301    def measure_speed( self, track_id ):
 302        """ Returns the velocities of the track """
 303        track = self.get_track_data( track_id )
 304        if track.shape[0] == 0:
 305            return None 
 306        track = track[track[:,1].argsort()]
 307        return ut.velocities( track[:,1:4] )
 308
 309    def measure_features( self, track_id, features ):
 310        """ Measure features along all the track """
 311        mask = self.epicure.get_mask( track_id )
 312        res = {}
 313        for feat in features:
 314            res[feat] = []
 315        for frame in mask:
 316            props = ut.labels_properties( frame )
 317            if len(props) > 0:
 318                if "Area" in features:
 319                    res["Area"].append( props[0].area )
 320                if "Hull" in features:
 321                    res["Hull"].append( props[0].area_convex )
 322                if "Elongation" in features:
 323                    res["Elongation"].append( props[0].axis_major_length )
 324                if "Eccentricity" in features:
 325                    res["Eccentricity"].append( props[0].eccentricity )
 326                if "Perimeter" in features:
 327                    res["Perimeter"].append( props[0].perimeter )
 328                if "Solidity" in features:
 329                    res["Solidity"].append( props[0].solidity )
 330        return res
 331
 332    def measure_specific_feature( self, track_id, featureName ):
 333        """ Measure some specific feature """
 334        if featureName == "Similarity":
 335            import skimage.metrics as imetrics
 336            movie = self.epicure.get_label_movie( track_id, extend=1.5 )
 337            sim_scores = []
 338            for i in range(0, len(movie)-1):
 339                score = imetrics.normalized_mutual_information( movie[i], movie[i+1] ) 
 340                sim_scores.append(score)
 341            return sim_scores
 342
 343    def measure_labels(self, segimg):
 344        """ Get the dataframe of the labels in the segmented image """
 345        resdf = None
 346        for iframe, frame in progress(enumerate(segimg)):
 347            frame_table = ut.labels_to_table( frame, iframe )
 348            if resdf is None:
 349                resdf = pd.DataFrame(frame_table)
 350            else:
 351                resdf = pd.concat([resdf, pd.DataFrame(frame_table)])
 352        return resdf
 353
 354    def add_track_frame(self, label, frame, centroid, tree=None, group=None):
 355        """ Add one frame to the track """
 356        new_frame = np.array([label, frame, centroid[0], centroid[1]])
 357        self.track_data = np.vstack((self.track_data, new_frame))
 358            
 359    def get_track_list(self):
 360        """ Get list of unique track_ids """
 361        return np.unique( self.track_data[:,0] )
 362    
 363    def get_tracks_list_frames( self, frames ):
 364        """ Return list of tracks present on list of frames """
 365        return np.unique( self.track_data[ np.isin( self.track_data[:,1], frames), 0] ) 
 366    
 367    def get_tracks_on_frame( self, tframe ):
 368        """ Return list of tracks present on given frame """
 369        return np.unique( self.track_data[ self.track_data[:,1]==tframe, 0] ) 
 370
 371    def has_track(self, label):
 372        """ Test if track label is present """
 373        return label in self.track_data[:,0]
 374    
 375    def has_tracks(self, labels):
 376        """ Test if track labels are present """
 377        return np.isin( labels, self.track_data[:,0] )
 378    
 379    def nb_points(self):
 380        """ Number of points in the tracks """
 381        return self.track_data.shape[0]
 382
 383    def nb_tracks(self):
 384        """ Return number of tracks """
 385        #return self.track._manager.__len__()
 386        return len(self.get_track_list())
 387
 388    def gaped_track(self, track_id):
 389        """ Check if there is a gap (missing frame) in a track """
 390        indexes = self.get_track_indexes(track_id)
 391        if len(indexes) <= 0:
 392            return False
 393        track_frames = self.track_data[indexes,1]
 394        return ((np.max(track_frames)-np.min(track_frames)+1) > len(track_frames) )
 395
 396    def gap_frames(self, track_id):
 397        """ Returns the frame(s) at which the gap(s) are """
 398        track_frames = self.get_track_column( track_id, "frame" )
 399        gaps = []
 400        if len( track_frames ) > 0:
 401            min_frame = int( np.min(track_frames) )
 402            max_frame = int( np.max(track_frames) )
 403            gaps = np.setdiff1d( np.arange(min_frame+1, max_frame), track_frames ).tolist()
 404            if len(gaps) > 0:
 405                gaps.sort()
 406        return gaps
 407            
 408    def check_gap(self, tracks=None, verbose=None):
 409        """ Check if there is a track with a gap, flag it if yes """
 410        if tracks is None:
 411            tracks = self.get_track_list()
 412        gaped = []
 413        for track in tracks:
 414            if self.gaped_track( track ):
 415                gaped.append(track)
 416        if verbose is None:
 417            verbose = self.epicure.verbose
 418        if verbose > 0 and len(gaped)>0:
 419            ut.show_warning("Gap in track(s) "+str(gaped)+"\n"
 420            +"Consider doing sanity_check in Editing onglet to fix it")
 421        return gaped
 422
 423    def get_track_indexes(self, track_id):
 424        """ Get indexes of track_id tracks position in the arrays """
 425        if isinstance( track_id,  int ):
 426            return (np.flatnonzero( self.track_data[:,0] == track_id ) )
 427        return (np.flatnonzero( np.isin( self.track_data[:,0], track_id ) ) )
 428    
 429    def get_track_indexes_onframes( self, track_id, frames ):
 430        """ Get indexes of track_id tracks position in the arrays """
 431        if isinstance( frames, int ):
 432            frames = [frames]
 433        if isinstance( track_id,  int ):
 434            return (np.flatnonzero( (self.track_data[:,0] == track_id) * np.isin( self.track_data[:,1], frames) ) )
 435        return (np.flatnonzero( np.isin( self.track_data[:,0], track_id ) * np.isin( self.track_data[:,1], frames) ) )
 436
 437    def get_track_indexes_from_frame(self, track_id, frame):
 438        """ Get indexes of track_id tracks position in the arrays from the given frame """
 439        if type(track_id) == int:
 440            return (np.argwhere( (self.track_data[:,0] == track_id)*(self.track_data[:,1]>= frame) )).flatten()
 441        return (np.argwhere( np.isin( self.track_data[:,0], track_id )*(self.track_data[:,1]>= frame) )).flatten()
 442
 443    def get_index(self, track_id, frame ):
 444        """ Get index of track_id at given frame """
 445        if np.isscalar(track_id):
 446            track_id = [track_id]
 447        return np.argwhere( (np.isin(self.track_data[:,0], track_id))*(self.track_data[:,1] == frame) )
 448
 449    def get_small_tracks(self, max_length=1):
 450        """ Get tracks smaller than the given threshold """
 451        labels = []
 452        lengths = []
 453        positions = []
 454        for lab in self.get_track_list():
 455            indexes = self.get_track_indexes(lab)
 456            length = len(indexes)
 457            if length <= max_length:
 458                pos = self.mean_position( indexes, only_first=False )
 459                labels.append(lab)
 460                lengths.append(length)
 461                positions.append(pos)
 462        return labels, lengths, positions
 463
 464    def get_track_data(self, track_id):
 465        """ Return the data of track track_id """
 466        indexes = self.get_track_indexes( track_id )
 467        track = self.track_data[indexes,]
 468        return track
 469    
 470    def get_track_column( self, track_id, column ):
 471        """ Return the chosen column (frame, x, y, label) of track track_id """
 472        indexes = self.get_track_indexes( track_id )
 473        if column == "frame":
 474            return self.track_data[indexes, 1]
 475        if column == "label":
 476            return self.track_data[indexes, 0]
 477        if column == "pos":
 478            return self.track_data[indexes, 2:4]
 479        if column == "fullpos":
 480            return self.track_data[indexes, 1:4]
 481        track = self.track_data[indexes]
 482        return track
 483
 484    def get_frame_data( self, track_id, ind ):
 485        """ Get ind-th data of track track_id """
 486        track = self.get_track_data( track_id )
 487        return track[ind]
 488    
 489    def get_middle_position( self, track_id, framea, frameb ):
 490        """ Get track position in middle of frame a and frame b """
 491        inda = self.get_index( track_id, framea ) 
 492        indb = self.get_index( track_id, frameb )
 493        return self.mean_position( np.ravel( np.vstack((inda, indb)) ), only_first=False )
 494
 495    def get_position( self, track_id, frame ):
 496        """ Get position of the track at given frame """
 497        ind = self.get_index( track_id, frame )
 498        ind = ind.flatten()[0] ## ensure it's single element
 499        x,y = self.track_data[ind,2], self.track_data[ind,3]
 500        return [int(x), int(y)]
 501
 502    def get_full_position( self, track_id, frame ):
 503        """ Get position of the track at given frame, with the frame itself """
 504        ind = self.get_index( track_id, frame )
 505        ind = ind.flatten()[0] ## ensure it's single element
 506        x,y = self.track_data[ind,2], self.track_data[ind,3]
 507        return [frame,x,y]
 508
 509    def mean_position(self, indexes, only_first=False):
 510        """ Mean positions of tracks at indexes """
 511        if len(indexes) <= 0:
 512            return None
 513        track = self.track_data[indexes,]
 514        ## keep only the first frame of the tracks
 515        if only_first:
 516            min_frame = np.min(track[:,1])
 517            track = track[track[:,1]==min_frame,]
 518        return ( int(np.mean(track[:,1])), int(np.mean(track[:,2])), int(np.mean(track[:,3])) )
 519
 520    def get_first_frame(self, track_id):
 521        """ Returns first frame where track_id is present """
 522        track = self.get_track_data( track_id )
 523        if len(track) <= 0:
 524            return None
 525        return int( np.min(track[:,1]) )
 526
 527    def is_in_frame( self, track_id, frame ):
 528        """ Returns if track_id is present at given frame """
 529        track = self.get_track_data( track_id )
 530        if len(track) > 0:
 531            return frame in track[:,1]
 532        return False
 533    
 534    def get_last_frame(self, track_id):
 535        """ Returns last frame where track_id is present """
 536        track = self.get_track_data( track_id )
 537        if len(track) > 0:
 538            return int(np.max(track[:,1]))
 539        return None
 540    
 541    def get_extreme_frames(self, track_id):
 542        """ Returns the first and last frames where track_id is present """
 543        track = self.get_track_data( track_id )
 544        if track.shape[0] > 0:
 545            return (int(np.min(track[:,1])), int(np.max(track[:,1])) )
 546        return None, None
 547
 548    def get_mean_position(self, track_id, only_first=False):
 549        """ Get mean position of the track """
 550        indexes = self.get_track_indexes( track_id )
 551        return self.mean_position( indexes, only_first )
 552
 553    def update_centroid(self, track_id, frame, ind=None, cx=None, cy=None):
 554        """ Update track at given frame """
 555        if ind is None:
 556            ind = self.get_index( track_id, frame )
 557        if cx is None:
 558            prop = ut.getPropLabel( self.epicure.seg[frame], track_id )
 559            self.track_data[ind, 2:4] = prop.centroid[1]
 560        else:
 561            self.track_data[ind, 2] = cx
 562            self.track_data[ind, 3] = cy
 563
 564    def replace_on_frames( self, tids, new_tids, frames ):
 565        """ Replace the id tid by new_tid in all given frames """
 566        ind = self.get_track_indexes_onframes( tids, frames )
 567        cur_track = np.copy(self.track_data[ind])
 568        new_ids = np.repeat(-1, len(ind))
 569        for tid, new_tid in zip(tids, new_tids):
 570            self.update_graph_frames( tid, cur_track[cur_track[:,0]==tid,1] )
 571            new_ids[cur_track[:,0]==tid] = new_tid
 572        self.track_data[ind, 0] = new_ids
 573        
 574    def swap_frame_id(self, tid, otid, frame):
 575        """ Swap the ids of two tracks at frame """
 576        ind = int(self.get_index(tid, frame))
 577        oind = int(self.get_index(otid, frame))
 578        ## check if one of the label is an extreme of a track and potentially in the graph
 579        for track_index in [tid, otid]:
 580            min_frame, max_frame = self.get_extreme_frames( track_index )
 581            if (min_frame == frame) or (max_frame == frame):
 582                self.update_graph( track_index, frame )
 583        self.track_data[[ind, oind],0] = [otid, tid]
 584
 585    def update_track_on_frame(self, track_ids, frame):
 586        """ Update (add or modify) tracks at given frame """
 587        frame_table = ut.labels_table( labimg = np.where(np.isin(self.epicure.seg[frame], track_ids), self.epicure.seg[frame], 0), properties=self.properties )
 588        for x, y, tid in zip(frame_table["centroid-0"], frame_table["centroid-1"], frame_table["label"]):
 589            index = self.get_index(tid, frame)
 590            if len(index) > 0:
 591                self.update_centroid( tid, frame, index, int(x), int(y) )
 592            else:
 593                cur_cell = np.array( [[tid, frame, int(x), int(y)]] )
 594                self.track_data = np.append(self.track_data, cur_cell, axis=0)
 595
 596    def add_tracks_fromindices( self, indices, track_ids ):
 597        """ Add tracks of given track ids from the indices"""
 598        new_data = np.empty( (0,4), int )
 599        for tid in np.unique(track_ids):
 600            keep = track_ids == tid 
 601            for frame in np.unique( indices[0][keep] ):
 602                cent0 = np.mean( indices[1][keep] ) 
 603                cent1 = np.mean( indices[2][keep] ) 
 604                new_data = np.append( new_data, np.array([[tid, frame, int(cent0), int(cent1)]]), axis=0 )
 605        self.track_data = np.append( self.track_data, new_data, axis=0)
 606    
 607    def add_one_frame(self, track_ids, frame, refresh=True):
 608        """ Add one frame from track """
 609        for tid in track_ids:
 610            frame_table = ut.labels_table( np.uint8(self.epicure.seg[frame]==tid), properties=self.properties ) 
 611            cur_cell = np.array( [tid, frame, int(frame_table["centroid-0"]), int(frame_table["centroid-1"])], dtype=np.uint32 )
 612            cur_cell = np.expand_dims(cur_cell, axis=0)
 613            self.track_data = np.append(self.track_data, cur_cell, axis=0)
 614
 615    def remove_one_frame( self, track_id, frame, handle_gaps=False, refresh=True ):
 616        """ 
 617        Remove one frame from track(s) 
 618        """
 619        inds = self.get_index( track_id, frame )
 620        if np.isscalar(track_id):
 621            track_id = [track_id]
 622        check_for_gaps = False
 623        for tid in track_id:
 624            ## removed frame is in the extremity of a track, can be in the graph
 625            first_frame, last_frame = self.get_extreme_frames( tid )
 626            if first_frame is None:
 627                continue
 628            if (first_frame == frame) or (last_frame == frame):
 629                self.update_graph( tid, frame )
 630            else:
 631                check_for_gaps = True
 632        self.track_data = np.delete( self.track_data, inds, axis=0 )
 633        ## gaps might have been created in the tracks, for now doesn't allow it so split the tracks
 634        if handle_gaps and check_for_gaps:
 635            gaped = self.check_gap( track_id, verbose=0 )
 636            if len(gaped) > 0:
 637                self.epicure.fix_gaps( gaped )
 638        
 639    def get_current_value(self, track_id, frame):
 640        ind = self.get_index(track_id, frame)
 641        centx, centy = self.track_data[ind, 2:4].astype(int).flatten()
 642        return self.epicure.seg[frame, centx, centy]
 643
 644    def clear_graph( self ):
 645        """ Check the state of the graph and removes non existing keys or values """
 646        if self.graph is None:
 647            return
 648        keys = list(self.graph.keys())
 649        for key in keys:
 650            if key not in self.track_data[:,0]:
 651                del self.graph[key]
 652            else:
 653                vals = self.graph[key]
 654                if isinstance(vals, list):
 655                    for val in vals:
 656                        if val not in self.track_data[:,0]:
 657                            del self.graph[key]
 658                            break
 659                else:
 660                    if vals not in self.track_data[:,0]:
 661                        del self.graph[key]
 662
 663    def set_graph(self, graph):
 664        """ Set the current graph (eg imported from TrackMate XML file) """
 665        self.graph = graph
 666        ## set the divisions from the graph
 667        self.epicure.inspecting.get_divisions()
 668
 669    def update_graph_frames( self, track_id, frames ):
 670        """ Update graph when one label was deleted at given frames """
 671        fframe = np.min(frames)
 672        lframe = np.max(frames)
 673        self.update_graph( track_id, fframe )
 674        self.update_graph( track_id, lframe )
 675
 676    def update_graph(self, track_id, frame):
 677        """ Update graph if deleted label was linked at that frame, assume keys are unique """
 678        if self.graph is not None:
 679            ## handles current node is last of his branch
 680            parents = self.last_in_graph( track_id, frame )
 681            current_label = self.get_current_value( track_id, frame )
 682            for parent in parents:
 683                if current_label == 0:
 684                    del self.graph[parent]
 685                else:
 686                    self.update_child( parent, track_id, current_label )
 687            ## handles when current track is first frame of a division
 688            if self.first_in_graph( track_id, frame ):
 689                if current_label == 0:
 690                    del self.graph[track_id]
 691                else:
 692                    self.update_key( track_id, current_label ) 
 693
 694    def update_child(self, parent, prev_key, new_key):
 695        """ Change the value of a key in the graph """
 696        if isinstance(self.graph[parent], list):
 697            self.graph[parent] = [new_key if val == prev_key else val for val in self.graph[parent]]
 698        else:
 699            if self.graph[parent] == prev_key:
 700                self.graph[parent] = new_key
 701
 702    def update_key(self, prev_key, new_key):
 703        """ Change the value of a key in the graph """
 704        self.graph[new_key] = self.graph.pop(prev_key)
 705
 706    def is_parent( self, cur_id ):
 707        """ Return if the current id is in the graph (as a parent, so in values) """
 708        if self.graph is None:
 709            return False
 710        return any( cur_id in vals if isinstance(vals, list) else cur_id in [vals] for vals in self.graph.values() )
 711
 712    def add_division( self, childa, childb, parent ):
 713        """ Add info of a division to the graph of divisions/merges """
 714        if self.graph is None:
 715            self.graph = {}
 716        self.graph.update({childa: [parent], childb: [parent]})
 717
 718    def remove_division( self, parent ):
 719        """ Remove a division event from the graph """
 720        self.graph = {key: vals for key, vals in self.graph.items() if not ( self.graph_parent(key) == parent )  }
 721
 722    def last_in_graph(self, track_id, frame=None, check_last=True):
 723        """ Check if given label and frame is the last of a branch, in the graph """
 724        if check_last:
 725            return [key for key, vals in self.graph.items() if track_id in (vals if isinstance(vals, list) else [vals]) and self.get_last_frame(track_id) == frame]
 726        return [key for key, vals in self.graph.items() if track_id in (vals if isinstance(vals, list) else [vals])]
 727
 728    def first_in_graph(self, track_id, frame=None, check_first=True):
 729        """ Check if the given label and frame is the first in the branch so the node in the graph """
 730        if check_first:
 731            return track_id in self.graph and self.get_first_frame(track_id) == frame
 732        return track_id in self.graph
 733
 734    def remove_on_frames( self, track_ids, frames ):
 735        """ Remove tracks with given id on specified frames """
 736        track_ids = track_ids.tolist()
 737        if 0 in track_ids:
 738            track_ids.remove(0)
 739        inds = self.get_track_indexes_onframes( track_ids, frames )
 740        for tid in track_ids:
 741            self.update_graph_frames( tid, frames )
 742        self.track_data = np.delete( self.track_data, inds, axis=0 )
 743
 744    def remove_tracks(self, track_ids):
 745        """ Remove track with given id """
 746        inds = self.get_track_indexes(track_ids)
 747        self.track_data = np.delete(self.track_data, inds, axis=0)
 748        self.remove_ids_from_graph( track_ids )
 749    
 750    def remove_ids_from_graph( self, track_ids ):
 751        """ Remove all ids from the graph """
 752        track_ids_set = set( track_ids )
 753        if self.graph is not None:
 754            self.graph = {
 755                key: vals for key, vals in self.graph.items()
 756                if (key not in track_ids_set) and ( not any( val in track_ids_set for val in (vals if isinstance(vals, list) else [vals])) )
 757            }
 758    
 759    def is_single_parent( self, cur_id ):
 760        """ Return if the current id is in the graph (as a single parent, not a merge) """
 761        if self.graph is None:
 762            return False
 763        return any( cur_id in [vals] if not isinstance(vals, list) else (cur_id in vals and len(vals)==1) for vals in self.graph.values() )
 764
 765       
 766    def build_tracks(self, track_df):
 767        """ Create tracks from dataframe (after tracking) """
 768        track = track_df[["track_id", "frame", "centroid-0", "centroid-1"]]
 769        #frame_prop = frame_table[["tree_id", "label", "nframes", "group"]]
 770        return np.array(track, int), None #dict(frame_prop)
 771
 772    def create_tracks(self):
 773        """ Create tracks from labels (without tracking) """
 774        #track_table = np.empty( (0,4), int )   
 775        labels = self.epicure.seg
 776        total = self.epicure.nframes
 777        if self.epicure.process_parallel:
 778            track_tables = Parallel( n_jobs=self.epicure.nparallel ) (
 779                delayed(ut.labels_to_table)(frame, iframe ) for iframe, frame in enumerate(labels)
 780            )
 781        else:
 782            track_tables = [ ut.labels_to_table( frame, iframe) for iframe, frame in progress(enumerate(labels), total=total) ]
 783        track_table = np.concatenate( [ tab for tab in track_tables if tab.shape[0] != 0 ], axis=0 ) # handle empty frame
 784        return track_table, None # track_prop
 785
 786    def add_track_features(self, labels):
 787        """ Add features specific to tracks (eg nframes) """
 788        nframes = np.zeros(len(labels), int)
 789        if self.epicure.verbose > 2:
 790            print("REPLACE BY COUNT METHOD")
 791        for track_id in np.unique(labels):
 792            cur_track = np.argwhere(labels == track_id)
 793            nframes[ list(cur_track) ] = len(cur_track)
 794        return nframes
 795    
 796
 797    ##########################################
 798    #### Tracking functions
 799
 800    def changed_start(self, i):
 801        """ Ensures that end frame > start frame """
 802        if i > self.end_frame.value():
 803            self.end_frame.setValue(i+1)
 804
 805    def changed_end(self, i):
 806        if i < self.start_frame.value():
 807            self.start_frame.setValue(i-1)
 808
 809    def find_parents(self, labels, twoframes):
 810        """ Find in the first frame the parents of labels from second frame """
 811        
 812        if self.track_choice.currentText() == "Laptrack-Centroids":
 813            return self.laptrack_centroids_twoframes(labels, twoframes, loose=True)
 814        
 815        if self.track_choice.currentText() == "Laptrack-Overlaps":
 816            return self.laptrack_overlaps_twoframes(labels, twoframes, loose=True)
 817        
 818
 819    def do_tracking(self):
 820        """ Start the tracking with the selected options """
 821        if self.frame_range.isChecked():
 822            start = self.start_frame.value()
 823            end = self.end_frame.value()
 824        else:
 825            start = 0
 826            end = self.nframes-1
 827        start_time = ut.start_time()
 828        self.viewer.window._status_bar._toggle_activity_dock(True)
 829        self.epicure.inspecting.reset_all_events()
 830        
 831        if self.track_choice.currentText() == "Laptrack-Centroids":
 832            if self.epicure.verbose > 1:
 833                print("Starting track with Laptrack-Centroids")
 834            self.laptrack_centroids( start, end )
 835            self.epicure.tracked = 1
 836        if self.track_choice.currentText() == "Laptrack-Overlaps":
 837            if self.epicure.verbose > 1:
 838                print("Starting track with Laptrack-Centroids")
 839            self.laptrack_overlaps( start, end )
 840            self.epicure.tracked = 1
 841        
 842        self.epicure.finish_update(contour=2)
 843        #self.epicure.reset_free_label()
 844        self.viewer.window._status_bar._toggle_activity_dock(False)
 845        if self.epicure.verbose > 0:
 846            ut.show_duration( start_time, header="Tracking done in " )
 847
 848    def show_trackoptions(self):
 849        self.gLapCentroids.setVisible(self.track_choice.currentText() == "Laptrack-Centroids")
 850        if laptrack_over:
 851            self.gLapOverlap.setVisible(self.track_choice.currentText() == "Laptrack-Overlaps")
 852
 853    def relabel_nonunique_labels(self, track_df):
 854        """ After tracking, some track can be splitted and get same label, fix that """
 855        inittids = np.unique(track_df["track_id"])
 856        labtracks = []
 857        saved_data = np.copy(self.epicure.seglayer.data)
 858        mframes = []
 859        tids = []
 860        used = np.unique( saved_data )
 861        all_labels = np.unique(track_df["label"])
 862        for tid in inittids:
 863            cdf = track_df[track_df["track_id"]==tid]
 864            #print(cdf)
 865            min_frame = np.min( cdf["frame"] )
 866            #labtrack = int( cdf["label"][cdf["frame"]==min_frame] )
 867            for lab in np.unique(cdf["label"]):
 868                labtracks.append(lab)
 869                mframes.append( min_frame )
 870                tids.append(tid)
 871        if len(labtracks) != len(np.unique(labtracks)):
 872            ## some labels are present several times
 873            used = used.tolist()
 874            for lab in all_labels :
 875                indexes = list(np.where(np.array(labtracks)==lab)[0])
 876                if len(indexes)>1:
 877                    minframes = [mframes[indy] for indy in range(len(labtracks)) if labtracks[indy]==lab]
 878                    indmin = indexes[ np.argmin( minframes ) ]
 879                    ## for the other(s), change the label
 880                    newvals = ut.get_free_labels( used, len(indexes) )
 881                    used = used + newvals
 882                    for i, ind in enumerate(indexes):
 883                        if ind != indmin:
 884                            tid = tids[ind]
 885                            newval = newvals[i]
 886                            track_df.loc[ (track_df["track_id"]==tid)  & (track_df["label"]==lab) , "label"] = newval
 887                            for frame in track_df["frame"][(track_df["track_id"]==tid) & (track_df["label"]==newval)]:
 888                                mask = (saved_data[frame]==lab)
 889                                self.epicure.seglayer.data[frame][mask] = newval
 890        
 891
 892    def relabel_trackids(self, track_df, splitdf, mergedf):
 893        """ Change the trackids to take the first label of each track """
 894        start_time = ut.start_time()
 895        new_trackids = track_df['track_id'].copy()
 896        new_splitdf = splitdf.copy()
 897        new_mergedf = mergedf.copy()
 898        
 899        unique_track_ids = np.unique(track_df['track_id'])
 900        if ut.version_python_minor(10):
 901            ## from python3.10, get futurewarning on groupby without group_keys and include_groups keywords
 902            first_labels = track_df.groupby('track_id', group_keys=False).apply(lambda x: x.loc[x['frame'].idxmin(), 'label'], include_groups=False).to_dict()
 903        else:
 904            first_labels = track_df.groupby('track_id').apply(lambda x: x.loc[x['frame'].idxmin(), 'label']).to_dict()
 905        
 906        for tid in unique_track_ids:
 907            newval = first_labels[tid]
 908            if tid != newval:
 909                new_trackids[track_df['track_id'] == tid] = newval
 910                if not new_splitdf.empty:
 911                    new_splitdf.loc[splitdf["parent_track_id"] == tid, "parent_track_id"] = newval
 912                    new_splitdf.loc[splitdf["child_track_id"] == tid, "child_track_id"] = newval
 913                if not new_mergedf.empty:
 914                    new_mergedf.loc[mergedf["parent_track_id"] == tid, "parent_track_id"] = newval
 915                    new_mergedf.loc[mergedf["child_track_id"] == tid, "child_track_id"] = newval
 916        if self.epicure.verbose > 1:
 917            ut.show_duration( start_time, header="Relabeling done in " )            
 918        return new_trackids, new_splitdf, new_mergedf
 919
 920    def change_labels(self, track_df):
 921        """ Change the labels at each frame according to tracks """
 922        for frame, frame_df in track_df.groupby("frame"):
 923            self.change_frame_labels(frame, frame_df)
 924
 925    def change_frame_labels(self, frame, frame_df):
 926        """ Change the labels at given frame according to tracks """
 927        track_ids = frame_df['track_id'].astype(int).values
 928        old_labels = frame_df["label"].astype(int).values
 929        seglayer = np.copy(self.epicure.seglayer.data[frame])
 930        for old_lab, new_lab in zip(old_labels, track_ids):
 931            mask = (seglayer==old_lab)
 932            self.epicure.seglayer.data[frame][mask] = new_lab
 933
 934    def label_to_dataframe( self, labimg, frame ):
 935        """ from label, get dataframe of centroids with properties """
 936        df = pd.DataFrame( ut.labels_table(labimg, properties=self.region_properties) )
 937        if df.shape[0] == 0:
 938            ## no labels in this frame
 939            return None
 940        df["frame"] = frame
 941        return df
 942    
 943    def optical_flow( self, img0, img1, radius ):
 944        """ Compute the optical flow between two images """
 945        v, u = optical_flow_ilk( img0, img1, radius=radius)
 946        return v, u
 947    
 948    def apply_flow( self, flowv, flowu, labimg ):
 949        """ Apply the calculated optical flow on a label image """
 950        nr, nc = labimg.shape
 951        rowc, colc = np.meshgrid( np.arange(nr), np.arange(nc), indexing="ij" )
 952        lab_reg = warp( labimg, np.array( [rowc+flowv, colc+flowu] ), order=0, mode="edge" )
 953        return lab_reg
 954    
 955    def labels_to_centroids( self, start_frame, end_frame ):
 956        """ Get centroids of each cell in dataframe """
 957        regionprops = [
 958            result
 959            for frame in range(start_frame, end_frame + 1)
 960            if (result := self.label_to_dataframe(self.epicure.seg[frame], frame)) is not None
 961        ]
 962        return pd.concat(regionprops)
 963    
 964    def labels_to_centroids_flow(self, start_frame, end_frame):
 965        """ Get centroids of each cell in dataframe """
 966        regionprops = []    
 967        radius = float( self.drift_radius.text() )
 968        if self.epicure.verbose > 1:
 969            if self.drift_correction.isChecked():
 970                print( "Apply drift correction to tracking with optical flow of radius "+str(radius) )
 971        prev_movie = None
 972        flow_v = None
 973        for frame in range(start_frame, end_frame+1):
 974            if self.drift_correction.isChecked():
 975                cur_movie = self.epicure.img[frame]
 976                if frame > start_frame:
 977                    v, u = self.optical_flow( prev_movie, cur_movie, radius )
 978                    if flow_v is None:
 979                        flow_v = v
 980                        flow_u = u
 981                    else:
 982                        flow_v = flow_v + v
 983                        flow_u = flow_u + u
 984                prev_movie = cur_movie
 985            clabel = self.epicure.seg[frame]  
 986            df = self.label_to_dataframe( clabel, frame )
 987            if flow_v is not None:
 988                c0 = np.array( np.floor( df["centroid-0"] ), dtype="uint8" )
 989                c1 = np.array( np.floor( df["centroid-1"] ), dtype="uint8" )
 990                df["centroid-0"] = df["centroid-0"] - flow_v[c0,c1]
 991                df["centroid-1"] = df["centroid-1"] - flow_u[c0,c1]
 992            regionprops.append(df)
 993        regionprops_df = pd.concat(regionprops)
 994        return regionprops_df
 995    
 996    def labels_flow(self, start_frame, end_frame ):
 997        """ Get registered label image corrected for optical flow """
 998        radius = float( self.drift_radius.text() )
 999        flow_v = None
1000        prev_movie = None
1001        res_labels = []
1002        for frame in range(start_frame, end_frame+1):
1003            cur_movie = self.epicure.img[frame]
1004            if prev_movie is not None:
1005                v, u = self.optical_flow( prev_movie, cur_movie, radius )
1006                if flow_v is None:
1007                    flow_v = v
1008                    flow_u = u
1009                else:
1010                    flow_v = flow_v + v
1011                    flow_u = flow_u + u
1012            prev_movie = cur_movie
1013            clabel = np.copy( self.epicure.seg[frame] ) 
1014            if flow_v is not None:         
1015                clabel = self.apply_flow( flow_v, flow_u, clabel )
1016            res_labels.append( clabel )
1017        res_labels = np.array(res_labels)
1018        return res_labels
1019
1020    def labels_ready(self, start_frame, end_frame, locked=True):
1021        """ Get labels of unlocked cells to track """
1022        if self.drift_correction.isChecked():
1023            return self.labels_flow( start_frame, end_frame )
1024        res_labels = self.epicure.seg[start_frame:end_frame+1] 
1025        return res_labels
1026    
1027    def label_frame_todf( self, frame ):
1028        """ For current frame, get label frame image then dataframe of centroids """
1029        clabel = self.epicure.seg[frame] #self.current_label_frame(frame)
1030        return self.label_to_dataframe( clabel, frame )
1031    
1032    def current_label_frame( self, frame ):
1033        """ For current frame, get label frame image """
1034        clabel = None
1035        #if self.track_only_in_roi.isChecked():
1036        #    clabel = self.epicure.only_current_roi(frame)
1037        if clabel is None:
1038            clabel = self.epicure.seg[frame]
1039        return clabel
1040
1041    def after_tracking( self, track_df, split_df, merge_df, progress_bar, indprogress ):
1042        """ Steps after tracking: get/show the graph from the track_df """
1043        if "frame_y" in track_df.keys():
1044            track_df["frame"] = track_df["frame_y"]
1045        graph = None
1046        progress_bar.set_description( "Update labels and tracks" )
1047        ## shift all by 1 so that doesn't start at 0
1048        if len(split_df) > 0:
1049            split_df[["parent_track_id"]] += 1
1050            split_df[["child_track_id"]] += 1
1051        if len(merge_df) > 0:
1052            merge_df[["parent_track_id"]] += 1
1053            merge_df[["child_track_id"]] += 1
1054        track_df[["track_id"]] += 1
1055       
1056        ## relabel if some track have the same label
1057        self.relabel_nonunique_labels(track_df)
1058        ## relabel track ids so that they are equal to the first label of the track
1059        newtids, split_df, merge_df = self.relabel_trackids( track_df, split_df, merge_df )
1060        track_df["track_id"] = newtids
1061        self.change_labels( track_df )
1062
1063        # create graph of division/merging
1064        self.graph = to_napari_graph(split_df, merge_df)
1065
1066        progress_bar.update(indprogress+1)
1067        
1068        ## update display if active
1069        self.replace_tracks( track_df )
1070
1071        progress_bar.update(indprogress+2)
1072        ## update the list of events, or others 
1073        self.epicure.updates_after_tracking()
1074        progress_bar.update(indprogress+3)
1075        return graph
1076
1077############ Laptrack centroids option
1078    
1079    def create_laptrack_centroids(self):
1080        """ GUI of the laptrack option """
1081        self.gLapCentroids, glap_layout = wid.group_layout( "Laptrack-Centroids" )
1082        mdist, self.max_dist = wid.value_line( "Max distance", "15.0", "Maximal distance between two labels in consecutive frames to link them (in pixels)" )
1083        glap_layout.addLayout(mdist)
1084        ## splitting ~ cell division
1085        scost, self.splitting_cost = wid.value_line( "Splitting cutoff", "1", "Weight to split a track in two (increasing it favors division)" )
1086        glap_layout.addLayout(scost)
1087        ## merging ~ error ?
1088        mcost, self.merging_cost = wid.value_line( "Merging cutoff", "0", "Weight to merge to labels together" )
1089        glap_layout.addLayout(mcost)
1090
1091        add_feat, self.check_penalties, self.bpenalties = wid.checkgroup_help( "Add features cost", True, "Add cell features in the tracking calculation", None )
1092        self.create_penalties()
1093        glap_layout.addWidget(self.check_penalties)
1094        glap_layout.addWidget(self.bpenalties)
1095        self.gLapCentroids.setLayout(glap_layout)
1096
1097    def show_penalties(self):
1098        self.bpenalties.setVisible(not self.bpenalties.isVisible())
1099
1100    def create_penalties(self):
1101        pen_layout = QVBoxLayout()
1102        areaCost, self.area_cost = wid.value_line( "Area difference", "2", "Weight of the difference of area between two labels to link them (0 to ignore)" )
1103        pen_layout.addLayout(areaCost)
1104        solidCost, self.solidity_cost = wid.value_line( "Solidity difference", "0", "Weight of the difference of solidity between two labels to link them (0 to ignore)" )
1105        pen_layout.addLayout(solidCost)
1106        self.bpenalties.setLayout(pen_layout)
1107
1108    def laptrack_centroids_twoframes(self, labels, twoframes, loose=False):
1109        """ Perform tracking of two frames with strict parameters """
1110        laptrack = LaptrackCentroids(self, self.epicure)
1111        laptrack.max_distance = float(self.max_dist.text()) 
1112        if loose:
1113            laptrack.max_distance = min(50, laptrack.max_distance) ## more probable to find a parent
1114        self.region_properties = ["label", "centroid"]
1115        #if self.check_penalties.isChecked():
1116        #    self.region_properties.append("area")
1117        #    self.region_properties.append("solidity")
1118        #    laptrack.penal_area = float(self.area_cost.text())
1119        #    laptrack.penal_solidity = float(self.solidity_cost.text())
1120        #laptrack.set_region_properties(with_extra=self.check_penalties.isChecked())
1121        laptrack.set_region_properties(with_extra=False)
1122            
1123        df = self.twoframes_centroid(twoframes)
1124        if set(np.unique(df["label"])) == set(labels):
1125            ## no other labels
1126            return [None]*len(labels) 
1127        laptrack.splitting_cost = False ## disable splitting option
1128        laptrack.merging_cost = False ## disable merging option
1129        parent_labels = laptrack.twoframes_track(df, labels)
1130        return parent_labels
1131    
1132    def twoframes_centroid(self, img):
1133        """ Get centroids of first frame only """
1134        df0 = self.label_to_dataframe( img[0], 0 )
1135        df1 = self.label_to_dataframe( img[1], 1 )
1136        return pd.concat([df0, df1])
1137    
1138    def laptrack_centroids(self, start, end):
1139        """ Perform track with laptrack option and chosen parameters """
1140        ## Laptrack tracker
1141        laptrack = LaptrackCentroids(self, self.epicure)
1142        laptrack.max_distance = float(self.max_dist.text())
1143        laptrack.splitting_cost = float(self.splitting_cost.text())
1144        laptrack.merging_cost = float(self.merging_cost.text())
1145        self.region_properties = ["label", "centroid"]
1146        if self.check_penalties.isChecked():
1147            self.region_properties.append("area")
1148            self.region_properties.append("solidity")
1149            laptrack.penal_area = float(self.area_cost.text())
1150            laptrack.penal_solidity = float(self.solidity_cost.text())
1151        laptrack.set_region_properties(with_extra=self.check_penalties.isChecked())
1152
1153        progress_bar = progress(total=7)
1154        progress_bar.set_description( "Prepare tracking" )
1155        if self.epicure.verbose > 1:
1156            print("Convert labels to centroids: use track info ?")
1157        self.undrifted = False
1158        if self.drift_correction.isChecked():
1159            df = self.labels_to_centroids_flow( start, end )
1160        else:
1161            df = self.labels_to_centroids( start, end )
1162        progress_bar.update(1)
1163        if self.epicure.verbose > 1:
1164            print("GO tracking")
1165        progress_bar.set_description( "Do tracking with LapTrack Centroids" )
1166        track_df, split_df, merge_df = laptrack.track_centroids(df)
1167        progress_bar.update(2)
1168        if self.epicure.verbose > 1:
1169            print("After tracking, update everything")
1170        self.after_tracking(track_df, split_df, merge_df, progress_bar, 2)
1171        progress_bar.update(6)
1172        progress_bar.close()
1173    
1174############ Laptrack overlap option
1175
1176    def create_laptrack_overlap(self):
1177        """ GUI of the laptrack overlap option """
1178        self.gLapOverlap, glap_layout = wid.group_layout( "Laptrack-Overlaps" )
1179        miou, self.min_iou = wid.value_line( "Min IOU", "0.1", "Minimum Intersection Over Union score to link to labels together" )
1180        glap_layout.addLayout(miou)
1181        
1182        scost, self.split_cost = wid.value_line( "Splitting cost", "0.2", "Weight of linking a parent label with two labels (increasing it for more divisions)" )
1183        glap_layout.addLayout(scost)
1184        
1185        mcost, self.merg_cost = wid.value_line( "Merging cost", "0", "Weight of merging two parent labels into one" )
1186        glap_layout.addLayout(mcost)
1187
1188        self.gLapOverlap.setLayout(glap_layout)
1189
1190    def laptrack_overlaps(self, start, end):
1191        """ Perform track with laptrack overlap option and chosen parameters """
1192        ## Laptrack tracker
1193        laptrack = LaptrackOverlaps(self, self.epicure)
1194        miniou = float(self.min_iou.text())
1195        if miniou >= 1.0:
1196            miniou = 1.0
1197        laptrack.cost_cutoff = 1.0 - miniou
1198        laptrack.splitting_cost = float(self.split_cost.text())
1199        laptrack.merging_cost = float(self.merg_cost.text())
1200        self.region_properties = ["label", "centroid"]
1201
1202        progress_bar = progress(total=6)
1203        progress_bar.set_description( "Prepare tracking" )
1204        labels = self.labels_ready( start, end )
1205        self.undrifted = False
1206        progress_bar.update(1)
1207        progress_bar.set_description( "Do tracking with LapTrack Overlaps" )
1208        track_df, split_df, merge_df = laptrack.track_overlaps( labels )
1209        progress_bar.update(2)
1210        
1211        ## get dataframe of coordinates to create the graph 
1212        df = self.labels_to_centroids( start, end )
1213        self.undrifted = True
1214        progress_bar.update(3)
1215        coordinate_df = df.set_index(["frame", "label"])
1216        tdf = track_df.set_index(["frame", "label"])
1217        track_df2 = pd.merge( tdf, coordinate_df, right_index=True, left_index=True).reset_index()
1218        self.after_tracking( track_df2, split_df, merge_df, progress_bar, 3 )
1219        progress_bar.update(6)
1220        progress_bar.close()
1221    
1222    def laptrack_overlaps_twoframes(self, labels, twoframes, loose=False):
1223        """ Perform tracking of two frames with strict parameters """
1224        laptrack = LaptrackOverlaps(self, self.epicure)
1225        miniou = min( float(self.min_iou.text()), 0.9999 ) ## ensure that miniou is < 1
1226        laptrack.cost_cutoff = 1.0 - miniou
1227        if loose:
1228            laptrack.cost_cutoff = 0.95 ## more probable to find a parent/child
1229        self.region_properties = ["label", "centroid"]
1230
1231        laptrack.splitting_cost = False ## disable splitting option
1232        laptrack.merging_cost = False ## disable merging option
1233        parent_labels = laptrack.twoframes_track(twoframes, labels)
1234        return parent_labels