| 20 | |
| 21 | |
| 22 | class TrackletManager: |
| 23 | def __init__(self, config, min_swap_len=2, min_tracklet_len=2, max_gap=0): |
| 24 | """ |
| 25 | |
| 26 | Parameters |
| 27 | ---------- |
| 28 | config : str |
| 29 | Path to a configuration file. |
| 30 | min_swap_len : float, optional (default=2) |
| 31 | Minimum swap length. |
| 32 | Swaps shorter than 2 frames are discarded by default. |
| 33 | min_tracklet_len : float, optional (default=2) |
| 34 | Minimum tracklet length. |
| 35 | Tracklets shorter than 2 frames are discarded by default. |
| 36 | max_gap : int, optional (default = 0). |
| 37 | Number of frames to consider when filling in missing data. |
| 38 | |
| 39 | Examples |
| 40 | -------- |
| 41 | |
| 42 | manager = TrackletManager(config_path, min_swap_frac=0, min_tracklet_frac=0) |
| 43 | |
| 44 | manager.load_tracklets_from_pickle(filename) |
| 45 | # Alternatively |
| 46 | manager.load_tracklets_from_h5(filename) |
| 47 | |
| 48 | manager.find_swapping_bodypart_pairs() |
| 49 | """ |
| 50 | self.config = config |
| 51 | self.cfg = auxiliaryfunctions.read_config(config) |
| 52 | self.min_swap_len = min_swap_len |
| 53 | self.min_tracklet_len = min_tracklet_len |
| 54 | self.max_gap = max_gap |
| 55 | |
| 56 | self.filename = "" |
| 57 | self.data = None |
| 58 | self.xy = None |
| 59 | self._xy = None |
| 60 | self.prob = None |
| 61 | self.nframes = 0 |
| 62 | self.times = [] |
| 63 | self.scorer = None |
| 64 | self.bodyparts = [] |
| 65 | self.nindividuals = len(self.cfg["individuals"]) |
| 66 | self.individuals = [] |
| 67 | self.tracklet2id = [] |
| 68 | self.tracklet2bp = [] |
| 69 | self.swapping_pairs = [] |
| 70 | self.swapping_bodyparts = [] |
| 71 | self._label_pairs = None |
| 72 | |
| 73 | def _load_tracklets(self, tracklets, auto_fill): |
| 74 | header = tracklets.pop("header") |
| 75 | self.scorer = header.get_level_values("scorer").unique().to_list() |
| 76 | bodyparts = header.get_level_values("bodyparts") |
| 77 | bodyparts_multi = [bp for bp in self.cfg["multianimalbodyparts"] if bp in bodyparts] |
| 78 | bodyparts_single = self.cfg["uniquebodyparts"] |
| 79 | mask_multi = bodyparts.isin(bodyparts_multi) |