|
| 1 | +import re |
| 2 | +import numpy as np |
| 3 | +import pandas as pd |
| 4 | +from pathlib import Path |
| 5 | +import pickle |
| 6 | +import ruamel.yaml as yaml |
| 7 | + |
| 8 | + |
| 9 | +class PoseEstimation: |
| 10 | + |
| 11 | + def __init__(self, dlc_dir=None, pkl_path=None, h5_path=None, yml_path=None, |
| 12 | + filename_prefix=''): |
| 13 | + if dlc_dir is None: |
| 14 | + assert pkl_path and h5_path and yml_path, \ |
| 15 | + ('If "dlc_dir" is not provided, then pkl_path, h5_path, and yml_path ' |
| 16 | + + 'must be provided') |
| 17 | + else: |
| 18 | + self.dlc_dir = Path(dlc_dir) |
| 19 | + if self.dlc_dir.stem != 'videos': |
| 20 | + self.dlc_dir = dlc_dir / 'videos' |
| 21 | + assert self.dlc_dir.exists(), f'Unable to find {dlc_dir}' |
| 22 | + |
| 23 | + # meta file: pkl - info about this DLC run (input video, configuration, etc.) |
| 24 | + if pkl_path is None: |
| 25 | + pkl_paths = list(self.dlc_dir.rglob(f'{filename_prefix}*meta.pickle')) |
| 26 | + assert len(pkl_paths) == 1, ('Unable to find one unique .pickle file in: ' |
| 27 | + + f'{dlc_dir} - Found: {len(pkl_paths)}') |
| 28 | + self.pkl_path = pkl_paths[0] |
| 29 | + else: |
| 30 | + self.pkl_path = Path(pkl_path) |
| 31 | + assert self.pkl_path.exists() |
| 32 | + |
| 33 | + # data file: h5 - body part outputs from the DLC post estimation step |
| 34 | + if h5_path is None: |
| 35 | + h5_paths = list(self.dlc_dir.rglob(f'{filename_prefix}*.h5')) |
| 36 | + assert len(h5_paths) == 1, ('Unable to find one unique .h5 file in: ' |
| 37 | + + f'{dlc_dir} - Found: {len(h5_paths)}') |
| 38 | + self.h5_path = h5_paths[0] |
| 39 | + else: |
| 40 | + self.h5_path = Path(h5_path) |
| 41 | + assert self.h5_path.exists() |
| 42 | + |
| 43 | + assert self.pkl_path.stem == self.h5_path.stem + '_meta', \ |
| 44 | + (f'Mismatching h5 ({self.h5_path.stem}) and pickle {self.pkl_path.stem}') |
| 45 | + |
| 46 | + # config file: yaml - configuration for invoking the DLC post estimation step |
| 47 | + if yml_path is None: |
| 48 | + yml_paths = list(self.dlc_dir.parent.glob(f'{filename_prefix}*.yaml')) |
| 49 | + # remove the one we save |
| 50 | + yml_paths = [val for val in yml_paths if not val.stem == "dlc_config_file"] |
| 51 | + assert len(yml_paths) == 1, ('Unable to find one unique .yaml file in: ' |
| 52 | + + f'{dlc_dir} - Found: {len(yml_paths)}') |
| 53 | + self.yml_path = yml_paths[0] |
| 54 | + else: |
| 55 | + self.yml_path = Path(yml_path) |
| 56 | + assert self.yml_path.exists() |
| 57 | + |
| 58 | + self._pkl = None |
| 59 | + self._rawdata = None |
| 60 | + self._yml = None |
| 61 | + self._data = None |
| 62 | + |
| 63 | + train_idx = np.where((np.array(self.yml['TrainingFraction'])*100 |
| 64 | + ).astype(int) == int(self.pkl['training set fraction' |
| 65 | + ] * 100))[0][0] |
| 66 | + train_iter = int(self.pkl['Scorer'].split('_')[-1]) |
| 67 | + |
| 68 | + self.model = {'Scorer': self.pkl['Scorer'], |
| 69 | + 'Task': self.yml['Task'], |
| 70 | + 'date': self.yml['date'], |
| 71 | + 'iteration': self.pkl['iteration (active-learning)'], |
| 72 | + 'shuffle': int(re.search('shuffle(\d+)', |
| 73 | + self.pkl['Scorer']).groups()[0]), |
| 74 | + 'snapshotindex': self.yml['snapshotindex'], |
| 75 | + 'trainingsetindex': train_idx, |
| 76 | + 'training_iteration': train_iter} |
| 77 | + |
| 78 | + self.fps = self.pkl['fps'] |
| 79 | + self.nframes = self.pkl['nframes'] |
| 80 | + |
| 81 | + self.creation_time = self.h5_path.stat().st_mtime |
| 82 | + |
| 83 | + @property |
| 84 | + def pkl(self): |
| 85 | + if self._pkl is None: |
| 86 | + with open(self.pkl_path, 'rb') as f: |
| 87 | + self._pkl = pickle.load(f) |
| 88 | + return self._pkl['data'] |
| 89 | + |
| 90 | + @property |
| 91 | + def yml(self): |
| 92 | + if self._yml is None: |
| 93 | + with open(self.yml_path, 'rb') as f: |
| 94 | + self._yml = yaml.safe_load(f) |
| 95 | + return self._yml |
| 96 | + |
| 97 | + @property |
| 98 | + def rawdata(self): |
| 99 | + if self._rawdata is None: |
| 100 | + self._rawdata = pd.read_hdf(self.h5_path) |
| 101 | + return self._rawdata |
| 102 | + |
| 103 | + @property |
| 104 | + def data(self): |
| 105 | + if self._data is None: |
| 106 | + self._data = self.reformat_rawdata() |
| 107 | + return self._data |
| 108 | + |
| 109 | + @property |
| 110 | + def df(self): |
| 111 | + top_level = self.rawdata.columns.levels[0][0] |
| 112 | + return self.rawdata.get(top_level) |
| 113 | + |
| 114 | + @property |
| 115 | + def body_parts(self): |
| 116 | + return self.df.columns.levels[0] |
| 117 | + |
| 118 | + def reformat_rawdata(self): |
| 119 | + error_message = (f'Total frames from .h5 file ({len(self.rawdata)}) differs ' |
| 120 | + + f'from .pickle ({self.pkl["nframes"]})') |
| 121 | + assert len(self.rawdata) == self.pkl['nframes'], error_message |
| 122 | + |
| 123 | + top_level = self.rawdata.columns.levels[0][0] |
| 124 | + |
| 125 | + body_parts_position = {} |
| 126 | + for body_part in self.body_parts: |
| 127 | + body_parts_position[body_part] = {c: self.df.get(body_part).get(c).values |
| 128 | + for c in self.df.get(body_part).columns} |
| 129 | + |
| 130 | + return body_parts_position |
| 131 | + |
| 132 | + |
| 133 | +def do_pose_estimation(video_filepaths, dlc_model, project_path, output_dir, |
| 134 | + videotype=None, gputouse=None, save_as_csv=False, batchsize=None, |
| 135 | + cropping=None, TFGPUinference=True, dynamic=(False, 0.5, 10), |
| 136 | + robust_nframes=False, allow_growth=False, use_shelve=False, |
| 137 | + modelprefix="", # need from paramset |
| 138 | + ): |
| 139 | + """Launch DLC's analyze_videos within element-deeplabcut |
| 140 | + :param video_filepaths: list of videos to analyze |
| 141 | + :param dlc_model: element-deeplabcut dlc.Model dict |
| 142 | + :param project_path: path to project config.yml |
| 143 | + :param output_dir: where to save output |
| 144 | + Remaining parameters are DLC's defaults |
| 145 | + """ |
| 146 | + from deeplabcut.pose_estimation_tensorflow import analyze_videos |
| 147 | + |
| 148 | + # ---- Build and save DLC configuration (yaml) file ---- |
| 149 | + dlc_config = dlc_model['config_template'] |
| 150 | + dlc_project_path = Path(project_path) |
| 151 | + assert dlc_project_path.exists(), (f'DLC project path ({dlc_project_path}) not ' |
| 152 | + + 'found on this machine') |
| 153 | + dlc_config['project_path'] = dlc_project_path.as_posix() |
| 154 | + |
| 155 | + # ---- Write DLC and basefolder yaml (config) files ---- |
| 156 | + |
| 157 | + # Write dlc config file to base (data) folder |
| 158 | + # This is important for parsing the DLC in datajoint imaging |
| 159 | + output_dir.mkdir(exist_ok=True) |
| 160 | + dlc_cfg_filepath = output_dir / 'dlc_config_file.yaml' |
| 161 | + with open(dlc_cfg_filepath, 'w') as f: |
| 162 | + yaml.dump(dlc_config, f) |
| 163 | + |
| 164 | + # ---- Trigger DLC prediction job ---- |
| 165 | + analyze_videos(config=dlc_cfg_filepath, videos=video_filepaths, |
| 166 | + shuffle=dlc_model['shuffle'], |
| 167 | + trainingsetindex=dlc_model['training_fract_idx'], |
| 168 | + destfolder=output_dir, |
| 169 | + videotype=None, gputouse=None, save_as_csv=False, batchsize=None, |
| 170 | + cropping=None, TFGPUinference=True, dynamic=(False, 0.5, 10), |
| 171 | + robust_nframes=False, allow_growth=False, use_shelve=False) |
0 commit comments