dlc2action.data.input_store
Specific realisations of dlc2action.data.base_store.InputStore are defined here
1# 2# Copyright 2020-2022 by A. Mathis Group and contributors. All rights reserved. 3# 4# This project and all its files are licensed under GNU AGPLv3 or later version. A copy is included in dlc2action/LICENSE.AGPL. 5# 6""" 7Specific realisations of `dlc2action.data.base_store.InputStore` are defined here 8""" 9 10from dlc2action.data.base_store import PoseInputStore 11from typing import Dict, List, Tuple, Union, Set, Optional, Iterable 12from dlc2action.utils import TensorDict, strip_suffix, strip_prefix 13import numpy as np 14from collections import defaultdict 15import torch 16import os 17import pickle 18from abc import abstractmethod 19import pandas as pd 20from p_tqdm import p_map 21from dlc2action import options 22from PIL import Image 23from tqdm import tqdm 24import mimetypes 25from pims import PyAVReaderIndexed 26 27 28class GeneralInputStore(PoseInputStore): 29 """ 30 A generalized realization of a PoseInputStore 31 32 Assumes the following file structure: 33 ``` 34 data_path 35 ├── video1DLC1000.pickle 36 ├── video2DLC400.pickle 37 ├── video1_features.pt 38 └── video2_features.pt 39 ``` 40 Here `data_suffix` is `{'DLC1000.pickle', 'DLC400.pickle'}` and `feature_suffix` (optional) is `'_features.pt'`. 41 """ 42 43 data_suffix = None 44 45 def __init__( 46 self, 47 video_order: List = None, 48 data_path: Union[Set, str] = None, 49 file_paths: Set = None, 50 data_suffix: Union[Set, str] = None, 51 data_prefix: Union[Set, str] = None, 52 feature_suffix: str = None, 53 convert_int_indices: bool = True, 54 feature_save_path: str = None, 55 canvas_shape: List = None, 56 len_segment: int = 128, 57 overlap: int = 0, 58 feature_extraction: str = "kinematic", 59 ignored_clips: List = None, 60 ignored_bodyparts: List = None, 61 default_agent_name: str = "ind0", 62 key_objects: Tuple = None, 63 likelihood_threshold: float = 0, 64 num_cpus: int = None, 65 frame_limit: int = 1, 66 normalize: bool = False, 67 feature_extraction_pars: Dict = None, 68 centered: bool = False, 69 transpose_features: bool = False, 70 *args, 71 **kwargs, 72 ) -> None: 73 """ 74 Parameters 75 ---------- 76 video_order : list, optional 77 a list of video ids that should be processed in the same order (not passed if creating from key objects 78 data_path : str | set, optional 79 the path to the folder where the pose and feature files are stored or a set of such paths 80 (not passed if creating from key objects or from `file_paths`) 81 file_paths : set, optional 82 a set of string paths to the pose and feature files 83 (not passed if creating from key objects or from `data_path`) 84 data_suffix : str | set, optional 85 the suffix or the set of suffices such that the pose files are named {video_id}{data_suffix} 86 (not passed if creating from key objects or if irrelevant for the dataset) 87 data_prefix : str | set, optional 88 the prefix or the set of prefixes such that the pose files for different video views of the same 89 clip are named {prefix}{sep}{video_id}{data_suffix} (not passed if creating from key objects 90 or if irrelevant for the dataset) 91 feature_suffix : str | set, optional 92 the suffix or the set of suffices such that the additional feature files are named 93 {video_id}{feature_suffix} (and placed at the data_path folder) 94 convert_int_indices : bool, default True 95 if `True`, convert any integer key `i` in feature files to `'ind{i}'` 96 feature_save_path : str, optional 97 the path to the folder where pre-processed files are stored (not passed if creating from key objects) 98 canvas_shape : List, default [1, 1] 99 the canvas size where the pose is defined 100 len_segment : int, default 128 101 the length of the segments in which the data should be cut (in frames) 102 overlap : int, default 0 103 the length of the overlap between neighboring segments (in frames) 104 feature_extraction : str, default 'kinematic' 105 the feature extraction method (see options.feature_extractors for available options) 106 ignored_clips : list, optional 107 list of strings of clip ids to ignore 108 ignored_bodyparts : list, optional 109 list of strings of bodypart names to ignore 110 default_agent_name : str, default 'ind0' 111 the agent name used as default in the pose files for a single agent 112 key_objects : tuple, optional 113 a tuple of key objects 114 likelihood_threshold : float, default 0 115 coordinate values with likelihoods less than this value will be set to 'unknown' 116 num_cpus : int, optional 117 the number of cpus to use in data processing 118 frame_limit : int, default 1 119 clips shorter than this number of frames will be ignored 120 feature_extraction_pars : dict, optional 121 parameters of the feature extractor 122 """ 123 124 super().__init__() 125 self.loaded_max = 0 126 if feature_extraction_pars is None: 127 feature_extraction_pars = {} 128 if ignored_clips is None: 129 ignored_clips = [] 130 self.bodyparts = [] 131 self.visibility = None 132 self.normalize = normalize 133 134 if canvas_shape is None: 135 canvas_shape = [1, 1] 136 if isinstance(data_suffix, str): 137 data_suffix = [data_suffix] 138 if isinstance(data_prefix, str): 139 data_prefix = [data_prefix] 140 if isinstance(data_path, str): 141 data_path = [data_path] 142 if isinstance(feature_suffix, str): 143 feature_suffix = [feature_suffix] 144 145 self.video_order = video_order 146 self.centered = centered 147 self.feature_extraction = feature_extraction 148 self.len_segment = int(len_segment) 149 self.data_suffices = data_suffix 150 self.data_prefixes = data_prefix 151 self.feature_suffix = feature_suffix 152 self.convert_int_indices = convert_int_indices 153 if overlap < 1: 154 overlap = overlap * len_segment 155 self.overlap = int(overlap) 156 self.canvas_shape = canvas_shape 157 self.default_agent_name = default_agent_name 158 self.feature_save_path = feature_save_path 159 self.data_suffices = data_suffix 160 self.data_prefixes = data_prefix 161 self.likelihood_threshold = likelihood_threshold 162 self.num_cpus = num_cpus 163 self.frame_limit = frame_limit 164 self.transpose = transpose_features 165 166 self.ram = False 167 self.min_frames = {} 168 self.original_coordinates = np.array([]) 169 170 self.file_paths = self._get_file_paths(file_paths, data_path) 171 172 self.extractor = options.feature_extractors[self.feature_extraction]( 173 self, 174 **feature_extraction_pars, 175 ) 176 177 self.canvas_center = np.array(canvas_shape) // 2 178 179 if ignored_clips is not None: 180 self.ignored_clips = ignored_clips 181 else: 182 self.ignored_clips = [] 183 if ignored_bodyparts is not None: 184 self.ignored_bodyparts = ignored_bodyparts 185 else: 186 self.ignored_bodyparts = [] 187 188 self.step = self.len_segment - self.overlap 189 if self.step < 0: 190 raise ValueError( 191 f"The overlap value ({self.overlap}) cannot be larger than len_segment ({self.len_segment}" 192 ) 193 194 if self.feature_save_path is None and data_path is not None: 195 self.feature_save_path = os.path.join(data_path[0], "trimmed") 196 197 if key_objects is None and self.video_order is not None: 198 print("Computing input features...") 199 self.data = self._load_data() 200 elif key_objects is not None: 201 self.load_from_key_objects(key_objects) 202 203 def __getitem__(self, ind: int) -> Dict: 204 prompt = self.data[ind] 205 if not self.ram: 206 with open(prompt, "rb") as f: 207 prompt = pickle.load(f) 208 return prompt 209 210 def __len__(self) -> int: 211 if self.data is None: 212 raise RuntimeError("The input store data has not been initialized!") 213 return len(self.data) 214 215 @classmethod 216 def _get_file_paths(cls, file_paths: Set, data_path: Union[str, Set]) -> List: 217 """ 218 Get a set of relevant files 219 220 Parameters 221 ---------- 222 file_paths : set 223 a set of filepaths to include 224 data_path : str | set 225 the path to a folder that contains relevant files (a single path or a set) 226 227 Returns 228 ------- 229 file_paths : list 230 a list of relevant file paths (input and feature files that follow the dataset naming pattern) 231 """ 232 233 if file_paths is None: 234 file_paths = [] 235 file_paths = list(file_paths) 236 if data_path is not None: 237 if isinstance(data_path, str): 238 data_path = [data_path] 239 for folder in data_path: 240 file_paths += [os.path.join(folder, x) for x in os.listdir(folder)] 241 return file_paths 242 243 def get_folder(self, video_id: str) -> str: 244 """ 245 Get the input folder that the file with this video id was read from 246 247 Parameters 248 ---------- 249 video_id : str 250 the video id 251 252 Returns 253 ------- 254 folder : str 255 the path to the directory that contains the input file associated with the video id 256 """ 257 258 for file in self.file_paths: 259 if ( 260 strip_prefix( 261 strip_suffix(os.path.basename(file), self.data_suffices), 262 self.data_prefixes, 263 ) 264 == video_id 265 ): 266 return os.path.dirname(file) 267 268 def remove(self, indices: List) -> None: 269 """ 270 Remove the samples corresponding to indices 271 272 Parameters 273 ---------- 274 indices : int 275 a list of integer indices to remove 276 """ 277 278 if len(indices) > 0: 279 mask = np.ones(len(self.original_coordinates)) 280 mask[indices] = 0 281 mask = mask.astype(bool) 282 for file in self.data[~mask]: 283 os.remove(file) 284 self.original_coordinates = self.original_coordinates[mask] 285 self.data = self.data[mask] 286 if self.metadata is not None: 287 self.metadata = self.metadata[mask] 288 289 def key_objects(self) -> Tuple: 290 """ 291 Return a tuple of the key objects necessary to re-create the Store 292 293 Returns 294 ------- 295 key_objects : tuple 296 a tuple of key objects 297 """ 298 299 for k, v in self.min_frames.items(): 300 self.min_frames[k] = dict(v) 301 for k, v in self.max_frames.items(): 302 self.max_frames[k] = dict(v) 303 return ( 304 self.original_coordinates, 305 dict(self.min_frames), 306 dict(self.max_frames), 307 self.data, 308 self.visibility, 309 self.step, 310 self.file_paths, 311 self.len_segment, 312 self.metadata, 313 ) 314 315 def load_from_key_objects(self, key_objects: Tuple) -> None: 316 """ 317 Load the information from a tuple of key objects 318 319 Parameters 320 ---------- 321 key_objects : tuple 322 a tuple of key objects 323 """ 324 325 ( 326 self.original_coordinates, 327 self.min_frames, 328 self.max_frames, 329 self.data, 330 self.visibility, 331 self.step, 332 self.file_paths, 333 self.len_segment, 334 self.metadata, 335 ) = key_objects 336 337 def to_ram(self) -> None: 338 """ 339 Transfer the data samples to RAM if they were previously stored as file paths 340 """ 341 342 if self.ram: 343 return 344 345 data = p_map(lambda x: self[x], list(range(len(self))), num_cpus=self.num_cpus) 346 # data = [load(x) for x in self.data] 347 self.data = TensorDict(data) 348 self.ram = True 349 350 def get_original_coordinates(self) -> np.ndarray: 351 """ 352 Return the original coordinates array 353 354 Returns 355 ------- 356 np.ndarray 357 an array that contains the coordinates of the data samples in original input data (video id, clip id, 358 start frame) 359 """ 360 361 return self.original_coordinates 362 363 def create_subsample(self, indices: List, ssl_indices: List = None): 364 """ 365 Create a new store that contains a subsample of the data 366 367 Parameters 368 ---------- 369 indices : list 370 the indices to be included in the subsample 371 ssl_indices : list, optional 372 the indices to be included in the subsample without the annotation data 373 """ 374 375 if ssl_indices is None: 376 ssl_indices = [] 377 new = self.new() 378 new.original_coordinates = self.original_coordinates[indices + ssl_indices] 379 new.min_frames = self.min_frames 380 new.max_frames = self.max_frames 381 new.data = self.data[indices + ssl_indices] 382 new.visibility = self.visibility 383 new.step = self.step 384 new.file_paths = self.file_paths 385 new.len_segment = self.len_segment 386 if self.metadata is None: 387 new.metadata = None 388 else: 389 new.metadata = self.metadata[indices + ssl_indices] 390 return new 391 392 def get_video_id(self, coords: Tuple) -> str: 393 """ 394 Get the video id from an element of original coordinates 395 396 Parameters 397 ---------- 398 coords : tuple 399 an element of the original coordinates array 400 401 Returns 402 ------- 403 video_id: str 404 the id of the video that the coordinates point to 405 """ 406 407 video_name = coords[0].split("---")[0] 408 return video_name 409 410 def get_clip_id(self, coords: Tuple) -> str: 411 """ 412 Get the clip id from an element of original coordinates 413 414 Parameters 415 ---------- 416 coords : tuple 417 an element of the original coordinates array 418 419 Returns 420 ------- 421 clip_id : str 422 the id of the clip that the coordinates point to 423 """ 424 425 clip_id = coords[0].split("---")[1] 426 return clip_id 427 428 def get_clip_length(self, video_id: str, clip_id: str) -> int: 429 """ 430 Get the clip length from the id 431 432 Parameters 433 ---------- 434 video_id : str 435 the video id 436 clip_id : str 437 the clip id 438 439 Returns 440 ------- 441 clip_length : int 442 the length of the clip 443 """ 444 445 inds = clip_id.split("+") 446 max_frame = min([self.max_frames[video_id][x] for x in inds]) 447 min_frame = max([self.min_frames[video_id][x] for x in inds]) 448 return max_frame - min_frame + 1 449 450 def get_clip_start_end(self, coords: Tuple) -> Tuple[int, int]: 451 """ 452 Get the clip start and end frames from an element of original coordinates 453 454 Parameters 455 ---------- 456 coords : tuple 457 an element of original coordinates array 458 459 Returns 460 ------- 461 start : int 462 the start frame of the clip that the coordinates point to 463 end : int 464 the end frame of the clip that the coordinates point to 465 """ 466 467 l = self.get_clip_length_from_coords(coords) 468 i = coords[1] 469 start = int(i) * self.step 470 end = min(start + self.len_segment, l) 471 return start, end 472 473 def get_clip_start(self, video_name: str, clip_id: str) -> int: 474 """ 475 Get the clip start frame from the video id and the clip id 476 477 Parameters 478 ---------- 479 video_name : str 480 the video id 481 clip_id : str 482 the clip id 483 484 Returns 485 ------- 486 clip_start : int 487 the start frame of the clip 488 """ 489 490 return max( 491 [self.min_frames[video_name][clip_id_k] for clip_id_k in clip_id.split("+")] 492 ) 493 494 def get_visibility( 495 self, video_id: str, clip_id: str, start: int, end: int, score: int 496 ) -> float: 497 """ 498 Get the fraction of the frames in that have a visibility score better than a hard_threshold 499 500 For example, in the case of keypoint data the visibility score can be the number of identified keypoints. 501 502 Parameters 503 ---------- 504 video_id : str 505 the video id of the frames 506 clip_id : str 507 the clip id of the frames 508 start : int 509 the start frame 510 end : int 511 the end frame 512 score : float 513 the visibility score hard_threshold 514 515 Returns 516 ------- 517 frac_visible: float 518 the fraction of frames with visibility above the hard_threshold 519 """ 520 521 s = 0 522 for ind_k in clip_id.split("+"): 523 s += np.sum(self.visibility[video_id][ind_k][start:end] > score) / ( 524 end - start 525 ) 526 return s / len(clip_id.split("+")) 527 528 def get_annotation_objects(self) -> Dict: 529 """ 530 Get a dictionary of objects necessary to create an AnnotationStore 531 532 Returns 533 ------- 534 annotation_objects : dict 535 a dictionary of objects to be passed to the AnnotationStore constructor where the keys are the names of 536 the objects 537 """ 538 539 min_frames = self.min_frames 540 max_frames = self.max_frames 541 num_bp = self.visibility 542 return { 543 "min_frames": min_frames, 544 "max_frames": max_frames, 545 "visibility": num_bp, 546 } 547 548 @classmethod 549 def get_file_ids( 550 cls, 551 data_suffix: Union[Set, str] = None, 552 data_path: Union[Set, str] = None, 553 data_prefix: Union[Set, str] = None, 554 file_paths: Set = None, 555 feature_suffix: Set = None, 556 *args, 557 **kwargs, 558 ) -> List: 559 """ 560 Process data parameters and return a list of ids of the videos that should 561 be processed by the __init__ function 562 563 Parameters 564 ---------- 565 data_suffix : set | str, optional 566 the suffix (or a set of suffixes) of the input data files 567 data_path : set | str, optional 568 the path to the folder where the pose and feature files are stored or a set of such paths 569 (not passed if creating from key objects or from `file_paths`) 570 data_prefix : set | str, optional 571 the prefix or the set of prefixes such that the pose files for different video views of the same 572 clip are named {prefix}{sep}{video_id}{data_suffix} (not passed if creating from key objects 573 or if irrelevant for the dataset) 574 file_paths : set, optional 575 a set of string paths to the pose and feature files 576 feature_suffix : str | set, optional 577 the suffix or the set of suffices such that the additional feature files are named 578 {video_id}{feature_suffix} (and placed at the `data_path` folder or at `file_paths`) 579 580 Returns 581 ------- 582 video_ids : list 583 a list of video file ids 584 """ 585 586 if data_suffix is None: 587 if cls.data_suffix is not None: 588 data_suffix = cls.data_suffix 589 else: 590 raise ValueError("Cannot get video ids without the data suffix!") 591 if feature_suffix is None: 592 feature_suffix = [] 593 if data_prefix is None: 594 data_prefix = "" 595 if isinstance(data_suffix, str): 596 data_suffix = [data_suffix] 597 else: 598 data_suffix = [x for x in data_suffix] 599 data_suffix = tuple(data_suffix) 600 if isinstance(data_prefix, str): 601 data_prefix = data_prefix 602 else: 603 data_prefix = tuple([x for x in data_prefix]) 604 if isinstance(feature_suffix, str): 605 feature_suffix = [feature_suffix] 606 if file_paths is None: 607 file_paths = [] 608 if data_path is not None: 609 if isinstance(data_path, str): 610 data_path = [data_path] 611 file_paths = [] 612 for folder in data_path: 613 file_paths += [os.path.join(folder, x) for x in os.listdir(folder)] 614 basenames = [os.path.basename(f) for f in file_paths] 615 ids = set() 616 for f in file_paths: 617 if f.endswith(data_suffix) and os.path.basename(f).startswith(data_prefix): 618 bn = os.path.basename(f) 619 video_id = strip_prefix(strip_suffix(bn, data_suffix), data_prefix) 620 if all([video_id + s in basenames for s in feature_suffix]): 621 ids.add(video_id) 622 ids = sorted(ids) 623 return ids 624 625 def get_bodyparts(self) -> List: 626 """ 627 Get a list of bodypart names 628 629 Parameters 630 ---------- 631 data_dict : dict 632 the data dictionary (passed to feature extractor) 633 clip_id : str 634 the clip id 635 636 Returns 637 ------- 638 bodyparts : list 639 a list of string or integer body part names 640 """ 641 642 return [x for x in self.bodyparts if x not in self.ignored_bodyparts] 643 644 def get_coords(self, data_dict: Dict, clip_id: str, bodypart: str) -> np.ndarray: 645 """ 646 Get the coordinates array of a specific bodypart in a specific clip 647 648 Parameters 649 ---------- 650 data_dict : dict 651 the data dictionary (passed to feature extractor) 652 clip_id : str 653 the clip id 654 bodypart : str 655 the name of the body part 656 657 Returns 658 ------- 659 coords : np.ndarray 660 the coordinates array of shape (#timesteps, #coordinates) 661 """ 662 663 columns = [x for x in data_dict[clip_id].columns if x != "likelihood"] 664 xy_coord = ( 665 data_dict[clip_id] 666 .xs(bodypart, axis=0, level=1, drop_level=False)[columns] 667 .values 668 ) 669 return xy_coord 670 671 def get_n_frames(self, data_dict: Dict, clip_id: str) -> int: 672 """ 673 Get the length of the clip 674 675 Parameters 676 ---------- 677 data_dict : dict 678 the data dictionary (passed to feature extractor) 679 clip_id : str 680 the clip id 681 682 Returns 683 ------- 684 n_frames : int 685 the length of the clip 686 """ 687 688 if clip_id in data_dict: 689 return len(data_dict[clip_id].groupby(level=0)) 690 else: 691 return min( 692 [len(data_dict[ind_k].groupby(level=0)) for ind_k in clip_id.split("+")] 693 ) 694 695 def _filter(self, data_dict: Dict) -> Tuple[Dict, Dict, Dict]: 696 """ 697 Apply filters to a data dictionary + normalize the values and generate frame index dictionaries 698 699 The filters include filling nan values, applying length and likelihood thresholds and removing 700 ignored clip ids. 701 """ 702 703 new_data_dict = {} 704 keys = list(data_dict.keys()) 705 for key in keys: 706 if key == "loaded": 707 continue 708 coord = data_dict.pop(key) 709 if key in self.ignored_clips: 710 continue 711 num_frames = len(coord.index.unique(level=0)) 712 if num_frames < self.frame_limit: 713 continue 714 if "likelihood" in coord.columns: 715 columns = list(coord.columns) 716 columns.remove("likelihood") 717 coord.loc[ 718 coord["likelihood"] < self.likelihood_threshold, columns 719 ] = np.nan 720 if not isinstance(self.centered, Iterable): 721 self.centered = [ 722 bool(self.centered) 723 for dim in ["x", "y", "z"] 724 if dim in coord.columns 725 ] 726 for i, dim in enumerate(["x", "y", "z"]): 727 if dim in coord.columns: 728 if self.centered[i]: 729 coord[dim] = coord[dim] + self.canvas_shape[i] // 2 730 # coord.loc[coord[dim] < -self.canvas_shape[i] * 3 // 2, dim] = np.nan 731 # coord.loc[coord[dim] > self.canvas_shape[i] * 3 // 2, dim] = np.nan 732 coord = coord.sort_index(level=0) 733 for bp in coord.index.unique(level=1): 734 coord.loc[coord.index.isin([bp], level=1)] = coord[ 735 coord.index.isin([bp], level=1) 736 ].interpolate() 737 dims = [x for x in coord.columns if x != "likelihood"] 738 mask = ~coord[dims[0]].isna() 739 for dim in dims[1:]: 740 mask = mask & (~coord[dim].isna()) 741 mean = coord.loc[mask].groupby(level=0).mean() 742 for frame in set(coord.index.get_level_values(0)): 743 if frame not in mean.index: 744 mean.loc[frame] = [np.nan for _ in mean.columns] 745 mean = mean.interpolate() 746 mean[mean.isna()] = 0 747 for dim in coord.columns: 748 if dim == "likelihood": 749 continue 750 coord.loc[coord[dim].isna(), dim] = mean.loc[ 751 coord.loc[coord[dim].isna()].index.get_level_values(0) 752 ][dim].to_numpy() 753 if np.sum(self.canvas_shape) > 0: 754 for i, dim in enumerate(["x", "y", "z"]): 755 if dim in coord.columns: 756 coord[dim] = ( 757 coord[dim] - self.canvas_shape[i] // 2 758 ) / self.canvas_shape[0] 759 new_data_dict[key] = coord 760 max_frames = {} 761 min_frames = {} 762 for key, value in new_data_dict.items(): 763 max_frames[key] = max(value.index.unique(0)) 764 min_frames[key] = min(value.index.unique(0)) 765 if "loaded" in data_dict: 766 new_data_dict["loaded"] = data_dict["loaded"] 767 return new_data_dict, min_frames, max_frames 768 769 def _get_files_from_ids(self): 770 files = defaultdict(lambda: []) 771 used_prefixes = defaultdict(lambda: []) 772 for f in self.file_paths: 773 if f.endswith(tuple([x for x in self.data_suffices])): 774 bn = os.path.basename(f) 775 video_id = strip_prefix( 776 strip_suffix(bn, self.data_suffices), self.data_prefixes 777 ) 778 ok = True 779 if self.data_prefixes is not None: 780 for p in self.data_prefixes: 781 if bn.startswith(p): 782 if p not in used_prefixes[video_id]: 783 used_prefixes[video_id].append(p) 784 else: 785 ok = False 786 break 787 if not ok: 788 continue 789 files[video_id].append(f) 790 files = [files[x] for x in self.video_order] 791 return files 792 793 def _make_trimmed_data(self, keypoint_dict: Dict) -> Tuple[List, Dict, List]: 794 """ 795 Cut a keypoint dictionary into overlapping pieces of equal length 796 """ 797 798 X = [] 799 original_coordinates = [] 800 lengths = defaultdict(lambda: {}) 801 if not os.path.exists(self.feature_save_path): 802 try: 803 os.mkdir(self.feature_save_path) 804 except FileExistsError: 805 pass 806 order = sorted(list(keypoint_dict.keys())) 807 for v_id in order: 808 keypoints = keypoint_dict[v_id] 809 v_len = min([len(x) for x in keypoints.values()]) 810 sp = np.arange(0, v_len, self.step) 811 pad = sp[-1] + self.len_segment - v_len 812 video_id, clip_id = v_id.split("---") 813 for key in keypoints: 814 if len(keypoints[key]) > v_len: 815 keypoints[key] = keypoints[key][:v_len] 816 if len(keypoints[key].shape) == 2: 817 keypoints[key] = np.pad(keypoints[key], ((0, pad), (0, 0))) 818 else: 819 keypoints[key] = np.pad( 820 keypoints[key], ((0, pad), (0, 0), (0, 0), (0, 0)) 821 ) 822 for i, start in enumerate(sp): 823 sample_dict = {} 824 original_coordinates.append((v_id, i)) 825 for key in keypoints: 826 sample_dict[key] = keypoints[key][start : start + self.len_segment] 827 sample_dict[key] = torch.tensor(np.array(sample_dict[key])).float() 828 sample_dict[key] = sample_dict[key].permute( 829 (*range(1, len(sample_dict[key].shape)), 0) 830 ) 831 name = os.path.join(self.feature_save_path, f"{v_id}_{start}.pickle") 832 X.append(name) 833 lengths[video_id][clip_id] = v_len 834 with open(name, "wb") as f: 835 pickle.dump(sample_dict, f) 836 return X, dict(lengths), original_coordinates 837 838 def _load_saved_features(self, video_id: str): 839 """ 840 Load saved features file `(#frames, #features)` 841 """ 842 843 basenames = [os.path.basename(x) for x in self.file_paths] 844 loaded_features_cat = [] 845 self.feature_suffix = sorted(self.feature_suffix) 846 for feature_suffix in self.feature_suffix: 847 i = basenames.index(os.path.basename(video_id) + feature_suffix) 848 path = self.file_paths[i] 849 if not os.path.exists(path): 850 raise RuntimeError(f"Did not find a feature file for {video_id}!") 851 extension = feature_suffix.split(".")[-1] 852 if extension in ["pickle", "pkl"]: 853 with open(path, "rb") as f: 854 loaded_features = pickle.load(f) 855 elif extension in ["pt", "pth"]: 856 loaded_features = torch.load(path) 857 elif extension == "npy": 858 loaded_features = np.load(path, allow_pickle=True).item() 859 else: 860 raise ValueError( 861 f"Found feature file in an unrecognized format: .{extension}. \n " 862 "Please save with torch (as .pt or .pth), numpy (as .npy) or pickle (as .pickle or .pkl)." 863 ) 864 loaded_features_cat.append(loaded_features) 865 keys = list(loaded_features_cat[0].keys()) 866 loaded_features = {} 867 for k in keys: 868 if k in ["min_frames", "max_frames", "video_tag"]: 869 loaded_features[k] = loaded_features_cat[0][k] 870 continue 871 features = [] 872 for x in loaded_features_cat: 873 if not isinstance(x[k], torch.Tensor): 874 features.append(torch.from_numpy(x[k])) 875 else: 876 features.append(x[k]) 877 a = torch.cat(features) 878 if self.transpose: 879 a = a.T 880 loaded_features[k] = a 881 return loaded_features 882 883 def get_likelihood( 884 self, data_dict: Dict, clip_id: str, bodypart: str 885 ) -> Union[np.ndarray, None]: 886 """ 887 Get the likelihood values 888 889 Parameters 890 ---------- 891 data_dict : dict 892 the data dictionary 893 clip_id : str 894 the clip id 895 bodypart : str 896 the name of the body part 897 898 Returns 899 ------- 900 likelihoods: np.ndarrray | None 901 `None` if the dataset doesn't have likelihoods or an array of shape (#timestamps) 902 """ 903 904 if "likelihood" in data_dict[clip_id].columns: 905 likelihood = ( 906 data_dict[clip_id] 907 .xs(bodypart, axis=0, level=1, drop_level=False) 908 .values[:, -1] 909 ) 910 return likelihood 911 else: 912 return None 913 914 def _get_video_metadata(self, metadata_list: Optional[List]): 915 """ 916 Make a single metadata dictionary from a list of dictionaries recieved from different data prefixes 917 """ 918 919 if metadata_list is None: 920 return None 921 else: 922 return metadata_list[0] 923 924 def get_indices(self, tag: int) -> List: 925 """ 926 Get a list of indices of samples that have a specific meta tag 927 928 Parameters 929 ---------- 930 tag : int 931 the meta tag for the subsample (`None` for the whole dataset) 932 933 Returns 934 ------- 935 indices : list 936 a list of indices that meet the criteria 937 """ 938 939 if tag is None: 940 return list(range(len(self.data))) 941 else: 942 return list(np.where(self.metadata == tag)[0]) 943 944 def get_tags(self) -> List: 945 """ 946 Get a list of all meta tags 947 948 Returns 949 ------- 950 tags: List 951 a list of unique meta tag values 952 """ 953 954 if self.metadata is None: 955 return [None] 956 else: 957 return list(np.unique(self.metadata)) 958 959 def get_tag(self, idx: int) -> Union[int, None]: 960 """ 961 Return a tag object corresponding to an index 962 963 Tags can carry meta information (like annotator id) and are accepted by models that require 964 that information. When a tag is `None`, it is not passed to the model. 965 966 Parameters 967 ---------- 968 idx : int 969 the index 970 971 Returns 972 ------- 973 tag : int 974 the tag object 975 """ 976 977 if self.metadata is None or idx is None: 978 return None 979 else: 980 return self.metadata[idx] 981 982 @abstractmethod 983 def _load_data(self) -> None: 984 """ 985 Load input data and generate data prompts 986 """ 987 988 989class FileInputStore(GeneralInputStore): 990 """ 991 An implementation of `dlc2action.data.InputStore` for datasets where each input data file corresponds to one video 992 """ 993 994 def _count_bodyparts( 995 self, data: Dict, stripped_name: str, max_frames: Dict 996 ) -> Dict: 997 """ 998 Create a visibility score dictionary (with a score from 0 to 1 assigned to each frame of each clip) 999 """ 1000 1001 result = {stripped_name: {}} 1002 prefixes = list(data.keys()) 1003 for ind in data[prefixes[0]]: 1004 res = 0 1005 for _, data_dict in data.items(): 1006 num_bp = len(data_dict[ind].index.unique(level=1)) 1007 coords = ( 1008 data_dict[ind].values.reshape( 1009 -1, num_bp, len(data_dict[ind].columns) 1010 )[: max_frames[ind], :, 0] 1011 != 0 1012 ) 1013 res = np.sum(coords, axis=1) + res 1014 result[stripped_name][ind] = (res / len(prefixes)) / coords.shape[1] 1015 return result 1016 1017 def _generate_features(self, data: Dict, video_id: str) -> Dict: 1018 """ 1019 Generate features from the raw coordinates 1020 """ 1021 1022 features = defaultdict(lambda: {}) 1023 loaded_common = [] 1024 for prefix, data_dict in data.items(): 1025 if prefix == "": 1026 prefix = None 1027 if "loaded" in data_dict: 1028 loaded_common.append(torch.tensor(data_dict.pop("loaded"))) 1029 1030 key_features = self.extractor.extract_features( 1031 data_dict, video_id, prefix=prefix 1032 ) 1033 for f_key in key_features: 1034 features[f_key].update(key_features[f_key]) 1035 if len(loaded_common) > 0: 1036 loaded_common = torch.cat(loaded_common, dim=1) 1037 else: 1038 loaded_common = None 1039 if self.feature_suffix is not None: 1040 loaded_features = self._load_saved_features(video_id) 1041 for clip_id, feature_tensor in loaded_features.items(): 1042 if not isinstance(feature_tensor, torch.Tensor): 1043 feature_tensor = torch.tensor(feature_tensor) 1044 if self.convert_int_indices and ( 1045 isinstance(clip_id, int) or isinstance(clip_id, np.integer) 1046 ): 1047 clip_id = f"ind{clip_id}" 1048 key1 = f"{os.path.basename(video_id)}---{clip_id}" 1049 if key1 in features: 1050 try: 1051 key2 = list(features[key1].keys())[0] 1052 n_frames = features[key1][key2].shape[0] 1053 if feature_tensor.shape[0] != n_frames: 1054 n = feature_tensor.shape[0] - n_frames 1055 if ( 1056 abs(n) > 2 1057 and abs(feature_tensor.shape[1] - n_frames) <= 2 1058 ): 1059 feature_tensor = feature_tensor.T 1060 # If off by <=2 frames, just clip the end 1061 elif n > 0 and n <= 2: 1062 feature_tensor = feature_tensor[:n_frames, :] 1063 elif n < 0 and n >= -2: 1064 filler = feature_tensor[-2:-1, :] 1065 for i in range(n_frames - feature_tensor.shape[0]): 1066 feature_tensor = torch.cat( 1067 [feature_tensor, filler], 0 1068 ) 1069 else: 1070 raise RuntimeError( 1071 f"Number of frames in precomputed features with shape" 1072 f" {feature_tensor.shape} is inconsistent with generated features!" 1073 ) 1074 if loaded_common is not None: 1075 if feature_tensor.shape[0] == loaded_common.shape[0]: 1076 feature_tensor = torch.cat( 1077 [feature_tensor, loaded_common], dim=1 1078 ) 1079 elif feature_tensor.shape[0] == loaded_common.shape[1]: 1080 feature_tensor = torch.cat( 1081 [feature_tensor.T, loaded_common], dim=1 1082 ) 1083 else: 1084 raise ValueError( 1085 "The features from the data file and from the feature file have a different number of frames!" 1086 ) 1087 features[key1]["loaded"] = feature_tensor 1088 except ValueError: 1089 raise RuntimeError( 1090 "Individuals in precomputed features are inconsistent " 1091 "with generated features" 1092 ) 1093 elif loaded_common is not None: 1094 for key in features: 1095 features[key]["loaded"] = loaded_common 1096 return features 1097 1098 def _load_data(self) -> np.array: 1099 """ 1100 Load input data and generate data prompts 1101 """ 1102 1103 if self.video_order is None: 1104 return None 1105 1106 files = defaultdict(lambda: []) 1107 for f in self.file_paths: 1108 if f.endswith(tuple([x for x in self.data_suffices])): 1109 bn = os.path.basename(f) 1110 video_id = strip_prefix( 1111 strip_suffix(bn, self.data_suffices), self.data_prefixes 1112 ) 1113 files[video_id].append(f) 1114 files = [files[x] for x in self.video_order] 1115 1116 def make_data_dictionary(filenames): 1117 data = {} 1118 stored_maxes = defaultdict(lambda: []) 1119 min_frames, max_frames = {}, {} 1120 name = strip_suffix(filenames[0], self.data_suffices) 1121 name = os.path.basename(name) 1122 stripped_name = strip_prefix(name, self.data_prefixes) 1123 metadata_list = [] 1124 for filename in filenames: 1125 name = strip_suffix(filename, self.data_suffices) 1126 name = os.path.basename(name) 1127 prefix = strip_suffix(name, [stripped_name]) 1128 data_new, tag = self._open_data(filename, self.default_agent_name) 1129 data_new, min_frames, max_frames = self._filter(data_new) 1130 data[prefix] = data_new 1131 for key, val in max_frames.items(): 1132 stored_maxes[key].append(val) 1133 metadata_list.append(tag) 1134 video_tag = self._get_video_metadata(metadata_list) 1135 sample_df = list(list(data.values())[0].values())[0] 1136 self.bodyparts = sorted(list(sample_df.index.unique(1))) 1137 smallest_maxes = dict.fromkeys(stored_maxes) 1138 for key, val in stored_maxes.items(): 1139 smallest_maxes[key] = np.amin(val) 1140 data_dict = self._generate_features(data, stripped_name) 1141 bp_dict = self._count_bodyparts( 1142 data=data, stripped_name=stripped_name, max_frames=smallest_maxes 1143 ) 1144 min_frames = {stripped_name: min_frames} # name is e.g. 20190707T1126-1226 1145 max_frames = {stripped_name: max_frames} 1146 names, lengths, coords = self._make_trimmed_data(data_dict) 1147 return names, lengths, coords, bp_dict, min_frames, max_frames, video_tag 1148 1149 dict_list = p_map(make_data_dictionary, files, num_cpus=self.num_cpus) 1150 # dict_list = tqdm([make_data_dictionary(f) for f in files]) 1151 1152 self.visibility = {} 1153 self.min_frames = {} 1154 self.max_frames = {} 1155 self.original_coordinates = [] 1156 self.metadata = [] 1157 X = [] 1158 for ( 1159 names, 1160 lengths, 1161 coords, 1162 bp_dictionary, 1163 min_frames, 1164 max_frames, 1165 metadata, 1166 ) in dict_list: 1167 X += names 1168 self.original_coordinates += coords 1169 self.visibility.update(bp_dictionary) 1170 self.min_frames.update(min_frames) 1171 self.max_frames.update(max_frames) 1172 if metadata is not None: 1173 self.metadata += metadata 1174 del dict_list 1175 if len(self.metadata) != len(self.original_coordinates): 1176 self.metadata = None 1177 else: 1178 self.metadata = np.array(self.metadata) 1179 1180 self.min_frames = dict(self.min_frames) 1181 self.max_frames = dict(self.max_frames) 1182 self.original_coordinates = np.array(self.original_coordinates) 1183 return np.array(X) 1184 1185 @abstractmethod 1186 def _open_data( 1187 self, filename: str, default_clip_name: str 1188 ) -> Tuple[Dict, Optional[Dict]]: 1189 """ 1190 Load the keypoints from filename and organize them in a dictionary 1191 1192 In `data_dictionary`, the keys are clip ids and the values are `pandas` dataframes with two-level indices. 1193 The first level is the frame numbers and the second is the body part names. The dataframes should have from 1194 two to four columns labeled `"x"`, `"y"` and (optionally) `"z"` and `"likelihood"`. Each frame should have 1195 information on all the body parts. You don't have to filter the data in any way or fill the nans, it will 1196 be done automatically. 1197 1198 Parameters 1199 ---------- 1200 filename : str 1201 path to the pose file 1202 default_clip_name : str 1203 the name to assign to a clip if it does not have a name in the raw data 1204 1205 Returns 1206 ------- 1207 data dictionary : dict 1208 a dictionary where the keys are clip ids and the values are keypoint dataframes (see above for details) 1209 metadata_dictionary : dict 1210 a dictionary where the keys are clip ids and the values are metadata objects (can be any additional information, 1211 like the annotator tag; for no metadata pass `None`) 1212 """ 1213 1214 1215class SequenceInputStore(GeneralInputStore): 1216 """ 1217 An implementation of `dlc2action.data.InputStore` for datasets where input data files correspond to multiple videos 1218 """ 1219 1220 def _count_bodyparts( 1221 self, data: Dict, stripped_name: str, max_frames: Dict 1222 ) -> Dict: 1223 """ 1224 Create a visibility score dictionary (with a score from 0 to 1 assigned to each frame of each clip) 1225 """ 1226 1227 result = {stripped_name: {}} 1228 for ind in data.keys(): 1229 num_bp = len(data[ind].index.unique(level=1)) 1230 coords = ( 1231 data[ind].values.reshape(-1, num_bp, len(data[ind].columns))[ 1232 : max_frames[ind], :, 0 1233 ] 1234 != 0 1235 ) 1236 res = np.sum(coords, axis=1) 1237 result[stripped_name][ind] = res / coords.shape[1] 1238 return result 1239 1240 def _generate_features(self, data: Dict, name: str) -> Dict: 1241 """ 1242 Generate features for an individual 1243 """ 1244 1245 features = self.extractor.extract_features(data, name, prefix=None) 1246 if self.feature_suffix is not None: 1247 loaded_features = self._load_saved_features(name) 1248 for clip_id, feature_tensor in loaded_features.items(): 1249 if not isinstance(feature_tensor, torch.Tensor): 1250 feature_tensor = torch.tensor(feature_tensor) 1251 if self.convert_int_indices and ( 1252 isinstance(clip_id, int) or isinstance(clip_id, np.integer) 1253 ): 1254 clip_id = f"ind{clip_id}" 1255 key1 = f"{os.path.basename(name)}---{clip_id}" 1256 if key1 in features: 1257 try: 1258 key2 = list(features[key1].keys())[0] 1259 n_frames = features[key1][key2].shape[0] 1260 if feature_tensor.shape[0] != n_frames: 1261 n = feature_tensor.shape[0] - n_frames 1262 if ( 1263 abs(n) > 2 1264 and abs(feature_tensor.shape[1] - n_frames) <= 2 1265 ): 1266 feature_tensor = feature_tensor.T 1267 # If off by <=2 frames, just clip the end 1268 elif n > 0 and n <= 2: 1269 feature_tensor = feature_tensor[:n_frames, :] 1270 elif n < 0 and n >= -2: 1271 filler = feature_tensor[-2:-1, :] 1272 for i in range(n_frames - feature_tensor.shape[0]): 1273 feature_tensor = torch.cat( 1274 [feature_tensor, filler], 0 1275 ) 1276 else: 1277 raise RuntimeError( 1278 print( 1279 f"Number of frames in precomputed features with shape" 1280 f" {feature_tensor.shape} is inconsistent with generated features!" 1281 ) 1282 ) 1283 features[key1]["loaded"] = feature_tensor 1284 except ValueError: 1285 raise RuntimeError( 1286 print( 1287 "Individuals in precomputed features are inconsistent " 1288 "with generated features" 1289 ) 1290 ) 1291 return features 1292 1293 def _load_data(self) -> np.array: 1294 """ 1295 Load input data and generate data prompts 1296 """ 1297 1298 if self.video_order is None: 1299 return None 1300 1301 files = [] 1302 for f in self.file_paths: 1303 if os.path.basename(f) in self.video_order: 1304 files.append(f) 1305 1306 def make_data_dictionary(seq_tuple): 1307 seq_id, sequence = seq_tuple 1308 data, tag = self._get_data(seq_id, sequence, self.default_agent_name) 1309 if "loaded" in data.keys(): 1310 loaded_features = data.pop("loaded") 1311 data, min_frames, max_frames = self._filter(data) 1312 sample_df = list(data.values())[0] 1313 self.bodyparts = sorted(list(sample_df.index.unique(1))) 1314 data_dict = self._generate_features(data, seq_id) 1315 for key in data_dict.keys(): 1316 data_dict[key]["loaded"] = loaded_features 1317 bp_dict = self._count_bodyparts( 1318 data=data, stripped_name=seq_id, max_frames=max_frames 1319 ) 1320 min_frames = {seq_id: min_frames} # name is e.g. 20190707T1126-1226 1321 max_frames = {seq_id: max_frames} 1322 names, lengths, coords = self._make_trimmed_data(data_dict) 1323 return names, lengths, coords, bp_dict, min_frames, max_frames, tag 1324 1325 seq_tuples = [] 1326 for file in files: 1327 opened = self._open_file(file) 1328 seq_tuples += opened 1329 dict_list = p_map( 1330 make_data_dictionary, sorted(seq_tuples), num_cpus=self.num_cpus 1331 ) 1332 # dict_list = tqdm([make_data_dictionary(f) for f in files]) 1333 1334 self.visibility = {} 1335 self.min_frames = {} 1336 self.max_frames = {} 1337 self.original_coordinates = [] 1338 self.metadata = [] 1339 X = [] 1340 for ( 1341 names, 1342 lengths, 1343 coords, 1344 bp_dictionary, 1345 min_frames, 1346 max_frames, 1347 metadata, 1348 ) in dict_list: 1349 X += names 1350 self.original_coordinates += coords 1351 self.visibility.update(bp_dictionary) 1352 self.min_frames.update(min_frames) 1353 self.max_frames.update(max_frames) 1354 if metadata is not None: 1355 self.metadata += metadata 1356 del dict_list 1357 1358 if len(self.metadata) != len(self.original_coordinates): 1359 self.metadata = None 1360 else: 1361 self.metadata = np.array(self.metadata) 1362 self.min_frames = dict(self.min_frames) 1363 self.max_frames = dict(self.max_frames) 1364 self.original_coordinates = np.array(self.original_coordinates) 1365 return np.array(X) 1366 1367 @classmethod 1368 def get_file_ids( 1369 cls, 1370 filenames: Set = None, 1371 data_path: Union[str, Set] = None, 1372 file_paths: Set = None, 1373 *args, 1374 **kwargs, 1375 ) -> List: 1376 """ 1377 Process data parameters and return a list of ids of the videos that should 1378 be processed by the __init__ function 1379 1380 Parameters 1381 ---------- 1382 filenames : set, optional 1383 a set of string filenames to search for (only basenames, not the whole paths) 1384 data_path : str | set, optional 1385 the path to the folder where the pose and feature files are stored or a set of such paths 1386 (not passed if creating from key objects or from `file_paths`) 1387 file_paths : set, optional 1388 a set of string paths to the pose and feature files 1389 (not passed if creating from key objects or from `data_path`) 1390 1391 Returns 1392 ------- 1393 video_ids : list 1394 a list of video file ids 1395 """ 1396 1397 if file_paths is None: 1398 file_paths = [] 1399 if data_path is not None: 1400 if isinstance(data_path, str): 1401 data_path = [data_path] 1402 file_paths = [] 1403 for folder in data_path: 1404 file_paths += [os.path.join(folder, x) for x in os.listdir(folder)] 1405 ids = set() 1406 for f in file_paths: 1407 if os.path.basename(f) in filenames: 1408 ids.add(os.path.basename(f)) 1409 ids = sorted(ids) 1410 return ids 1411 1412 @abstractmethod 1413 def _open_file(self, filename: str) -> List: 1414 """ 1415 Open a file and make a list of sequences 1416 1417 The sequence objects should contain information about all clips in one video. The sequences and 1418 video ids will be processed in the `_get_data` function. 1419 1420 Parameters 1421 ---------- 1422 filename : str 1423 the name of the file 1424 1425 Returns 1426 ------- 1427 video_tuples : list 1428 a list of video tuples: `(video_id, sequence)` 1429 """ 1430 1431 @abstractmethod 1432 def _get_data( 1433 self, video_id: str, sequence, default_agent_name: str 1434 ) -> Tuple[Dict, Optional[Dict]]: 1435 """ 1436 Get the keypoint dataframes from a sequence 1437 1438 The sequences and video ids are generated in the `_open_file` function. 1439 In `data_dictionary`, the keys are clip ids and the values are `pandas` dataframes with two-level indices. 1440 The first level is the frame numbers and the second is the body part names. The dataframes should have from 1441 two to four columns labeled `"x"`, `"y"` and (optionally) `"z"` and `"likelihood"`. Each frame should have 1442 information on all the body parts. You don't have to filter the data in any way or fill the nans, it will 1443 be done automatically. 1444 1445 Parameters 1446 ---------- 1447 video_id : str 1448 the video id 1449 sequence 1450 an object containing information about all clips in one video 1451 default_agent_name 1452 1453 Returns 1454 ------- 1455 data dictionary : dict 1456 a dictionary where the keys are clip ids and the values are keypoint dataframes (see above for details) 1457 metadata_dictionary : dict 1458 a dictionary where the keys are clip ids and the values are metadata objects (can be any additional information, 1459 like the annotator tag; for no metadata pass `None`) 1460 """ 1461 1462 1463class DLCTrackStore(FileInputStore): 1464 """ 1465 DLC track data 1466 1467 Assumes the following file structure: 1468 ``` 1469 data_path 1470 ├── video1DLC1000.pickle 1471 ├── video2DLC400.pickle 1472 ├── video1_features.pt 1473 └── video2_features.pt 1474 ``` 1475 Here `data_suffix` is `{'DLC1000.pickle', 'DLC400.pickle'}` and `feature_suffix` (optional) is `'_features.pt'`. 1476 1477 The feature files should to be dictionaries where keys are clip IDs (e.g. animal names) and values are 1478 feature values (arrays of shape `(#frames, #features)`). If the arrays are shaped as `(#features, #frames)`, 1479 set `transpose_features` to `True`. 1480 1481 The files can be saved with `numpy.save()` (with `.npy` extension), `torch.save()` (with `.pt` extension) or 1482 with `pickle.dump()` (with `.pickle` or `.pkl` extension). 1483 """ 1484 1485 def _open_data( 1486 self, filename: str, default_agent_name: str 1487 ) -> Tuple[Dict, Optional[Dict]]: 1488 """ 1489 Load the keypoints from filename and organize them in a dictionary 1490 1491 In `data_dictionary`, the keys are clip ids and the values are `pandas` dataframes with two-level indices. 1492 The first level is the frame numbers and the second is the body part names. The dataframes should have from 1493 two to four columns labeled `"x"`, `"y"` and (optionally) `"z"` and `"likelihood"`. Each frame should have 1494 information on all the body parts. You don't have to filter the data in any way or fill the nans, it will 1495 be done automatically. 1496 1497 Parameters 1498 ---------- 1499 filename : str 1500 path to the pose file 1501 default_clip_name : str 1502 the name to assign to a clip if it does not have a name in the raw data 1503 1504 Returns 1505 ------- 1506 data dictionary : dict 1507 a dictionary where the keys are clip ids and the values are keypoint dataframes (see above for details) 1508 metadata_dictionary : dict 1509 a dictionary where the keys are clip ids and the values are metadata objects (can be any additional information, 1510 like the annotator tag; for no metadata pass `None`) 1511 """ 1512 1513 if filename.endswith("h5"): 1514 temp = pd.read_hdf(filename) 1515 temp = temp.droplevel("scorer", axis=1) 1516 else: 1517 temp = pd.read_csv(filename, header=[1, 2]) 1518 temp.columns.names = ["bodyparts", "coords"] 1519 if "individuals" not in temp.columns.names: 1520 old_idx = temp.columns.to_frame() 1521 old_idx.insert(0, "individuals", self.default_agent_name) 1522 temp.columns = pd.MultiIndex.from_frame(old_idx) 1523 df = temp.stack(["individuals", "bodyparts"]) 1524 idx = pd.MultiIndex.from_product( 1525 [df.index.levels[0], df.index.levels[1], df.index.levels[2]], 1526 names=df.index.names, 1527 ) 1528 df = df.reindex(idx).fillna(value=0) 1529 animals = sorted(list(df.index.levels[1])) 1530 dic = {} 1531 for ind in animals: 1532 coord = df.iloc[df.index.get_level_values(1) == ind].droplevel(1) 1533 coord = coord[["x", "y", "likelihood"]] 1534 dic[ind] = coord 1535 1536 return dic, None 1537 1538 1539class DLCTrackletStore(FileInputStore): 1540 """ 1541 DLC tracklet data 1542 1543 Assumes the following file structure: 1544 ``` 1545 data_path 1546 ├── video1DLC1000.pickle 1547 ├── video2DLC400.pickle 1548 ├── video1_features.pt 1549 └── video2_features.pt 1550 ``` 1551 Here `data_suffix` is `{'DLC1000.pickle', 'DLC400.pickle'}` and `feature_suffix` (optional) is `'_features.pt'`. 1552 1553 The feature files should to be dictionaries where keys are clip IDs (e.g. animal names) and values are 1554 feature values (arrays of shape `(#frames, #features)`). If the arrays are shaped as `(#features, #frames)`, 1555 set `transpose_features` to `True`. 1556 1557 The files can be saved with `numpy.save()` (with `.npy` extension), `torch.save()` (with `.pt` extension) or 1558 with `pickle.dump()` (with `.pickle` or `.pkl` extension). 1559 """ 1560 1561 def _open_data( 1562 self, filename: str, default_agent_name: str 1563 ) -> Tuple[Dict, Optional[Dict]]: 1564 """ 1565 Load the keypoints from filename and organize them in a dictionary 1566 1567 In `data_dictionary`, the keys are clip ids and the values are `pandas` dataframes with two-level indices. 1568 The first level is the frame numbers and the second is the body part names. The dataframes should have from 1569 two to four columns labeled `"x"`, `"y"` and (optionally) `"z"` and `"likelihood"`. Each frame should have 1570 information on all the body parts. You don't have to filter the data in any way or fill the nans, it will 1571 be done automatically. 1572 1573 Parameters 1574 ---------- 1575 filename : str 1576 path to the pose file 1577 default_clip_name : str 1578 the name to assign to a clip if it does not have a name in the raw data 1579 1580 Returns 1581 ------- 1582 data dictionary : dict 1583 a dictionary where the keys are clip ids and the values are keypoint dataframes (see above for details) 1584 metadata_dictionary : dict 1585 a dictionary where the keys are clip ids and the values are metadata objects (can be any additional information, 1586 like the annotator tag; for no metadata pass `None`) 1587 """ 1588 1589 output = {} 1590 with open(filename, "rb") as f: 1591 data_p = pickle.load(f) 1592 header = data_p["header"] 1593 bodyparts = header.unique("bodyparts") 1594 1595 keys = sorted([key for key in data_p.keys() if key != "header"]) 1596 min_frames = defaultdict(lambda: 10**5) 1597 max_frames = defaultdict(lambda: 0) 1598 for tr_id in keys: 1599 coords = {} 1600 fr_i = int(list(data_p[tr_id].keys())[0][5:]) - 1 1601 for frame in data_p[tr_id]: 1602 count = 0 1603 while int(frame[5:]) > fr_i + 1: 1604 count += 1 1605 fr_i = fr_i + 1 1606 if count <= 3: 1607 for bp, name in enumerate(bodyparts): 1608 coords[(fr_i, name)] = coords[(fr_i - 1, name)] 1609 else: 1610 for bp, name in enumerate(bodyparts): 1611 coords[(fr_i, name)] = np.zeros( 1612 coords[(fr_i - 1, name)].shape 1613 ) 1614 fr_i = int(frame[5:]) 1615 if fr_i > max_frames[f"ind{tr_id}"]: 1616 max_frames[f"ind{tr_id}"] = fr_i 1617 if fr_i < min_frames[f"ind{tr_id}"]: 1618 min_frames[f"ind{tr_id}"] = fr_i 1619 for bp, name in enumerate(bodyparts): 1620 coords[(fr_i, name)] = data_p[tr_id][frame][bp][:3] 1621 1622 output[f"ind{tr_id}"] = pd.DataFrame( 1623 data=coords, index=["x", "y", "likelihood"] 1624 ).T 1625 return output, None 1626 1627 1628class PKUMMDInputStore(FileInputStore): 1629 """ 1630 PKU-MMD data 1631 1632 Assumes the following file structure: 1633 ``` 1634 data_path 1635 ├── 0073-R.txt 1636 ... 1637 └── 0274-L.txt 1638 ``` 1639 """ 1640 1641 data_suffix = ".txt" 1642 1643 def __init__( 1644 self, 1645 video_order: str = None, 1646 data_path: Union[str, Set] = None, 1647 file_paths: Set = None, 1648 feature_save_path: str = None, 1649 feature_extraction: str = "kinematic", 1650 len_segment: int = 128, 1651 overlap: int = 0, 1652 key_objects: Tuple = None, 1653 num_cpus: int = None, 1654 interactive: bool = False, 1655 feature_extraction_pars: Dict = None, 1656 *args, 1657 **kwargs, 1658 ) -> None: 1659 """ 1660 Parameters 1661 ---------- 1662 video_order : list, optional 1663 a list of video ids that should be processed in the same order (not passed if creating from key objects) 1664 data_path : str | set, optional 1665 the path to the folder where the pose and feature files are stored or a set of such paths 1666 (not passed if creating from key objects or from `file_paths`) 1667 file_paths : set, optional 1668 a set of string paths to the pose and feature files 1669 (not passed if creating from key objects or from `data_path`) 1670 feature_save_path : str, optional 1671 the path to the folder where pre-processed files are stored (not passed if creating from key objects) 1672 feature_extraction : str, default 'kinematic' 1673 the feature extraction method (run options.feature_extractors to see available options) 1674 len_segment : int, default 128 1675 the length of the segments in which the data should be cut (in frames) 1676 overlap : int, default 0 1677 the length of the overlap between neighboring segments (in frames) 1678 interactive : bool, default False 1679 if True, distances between two agents are included; if False, only the first agent features are computed 1680 key_objects : tuple, optional 1681 a tuple of key objects 1682 num_cpus : int, optional 1683 the number of cpus to use in data processing 1684 feature_extraction_pars : dict, optional 1685 parameters of the feature extractor 1686 """ 1687 1688 if feature_extraction_pars is None: 1689 feature_extraction_pars = {} 1690 feature_extraction_pars["interactive"] = interactive 1691 self.interactive = interactive 1692 super().__init__( 1693 video_order, 1694 data_path, 1695 file_paths, 1696 data_suffix=".txt", 1697 data_prefix=None, 1698 feature_suffix=None, 1699 convert_int_indices=False, 1700 feature_save_path=feature_save_path, 1701 canvas_shape=[1, 1, 1], 1702 len_segment=len_segment, 1703 overlap=overlap, 1704 feature_extraction=feature_extraction, 1705 ignored_clips=None, 1706 ignored_bodyparts=None, 1707 default_agent_name="ind0", 1708 key_objects=key_objects, 1709 likelihood_threshold=0, 1710 num_cpus=num_cpus, 1711 frame_limit=1, 1712 interactive=interactive, 1713 feature_extraction_pars=feature_extraction_pars, 1714 ) 1715 1716 def _open_data( 1717 self, filename: str, default_clip_name: str 1718 ) -> Tuple[Dict, Optional[Dict]]: 1719 """ 1720 Load the keypoints from filename and organize them in a dictionary 1721 1722 In `data_dictionary`, the keys are clip ids and the values are `pandas` dataframes with two-level indices. 1723 The first level is the frame numbers and the second is the body part names. The dataframes should have from 1724 two to four columns labeled `"x"`, `"y"` and (optionally) `"z"` and `"likelihood"`. Each frame should have 1725 information on all the body parts. You don't have to filter the data in any way or fill the nans, it will 1726 be done automatically. 1727 1728 Parameters 1729 ---------- 1730 filename : str 1731 path to the pose file 1732 default_clip_name : str 1733 the name to assign to a clip if it does not have a name in the raw data 1734 1735 Returns 1736 ------- 1737 data dictionary : dict 1738 a dictionary where the keys are clip ids and the values are keypoint dataframes (see above for details) 1739 metadata_dictionary : dict 1740 a dictionary where the keys are clip ids and the values are metadata objects (can be any additional information, 1741 like the annotator tag; for no metadata pass `None`) 1742 """ 1743 1744 keypoint_dict = {"0": [], "1": []} 1745 with open(filename) as f: 1746 for line in f.readlines(): 1747 line_data = list(map(float, line.split())) 1748 line_data = np.array(line_data) 1749 line_data = line_data.reshape((2, 25, 3))[:, :, [0, 2, 1]] 1750 for ind in keypoint_dict: 1751 keypoint_dict[ind].append(line_data[int(ind)]) 1752 for ind in keypoint_dict: 1753 data = np.stack(keypoint_dict[ind]) 1754 mi = pd.MultiIndex.from_product( 1755 [list(range(data.shape[0])), list(range(data.shape[1]))] 1756 ) 1757 data = data.reshape((-1, 3)) 1758 keypoint_dict[ind] = pd.DataFrame( 1759 data=data, index=mi, columns=["x", "y", "z"] 1760 ) 1761 if not self.interactive: 1762 keypoint_dict.pop("1") 1763 return keypoint_dict, None 1764 1765 1766class CalMS21InputStore(SequenceInputStore): 1767 """ 1768 CalMS21 data 1769 1770 Use the `'random:test_from_name:{name}'` and `'val-from-name:{val_name}:test-from-name:{test_name}'` 1771 partitioning methods with `'train'`, `'test'` and `'unlabeled'` names to separate into train, test and validation 1772 subsets according to the original files. For example, with `'val-from-name:test:test-from-name:unlabeled'` 1773 the data from the test file will go into validation and the unlabeled files will be the test. 1774 1775 Assumes the following file structure: 1776 ``` 1777 data_path 1778 ├── calms21_task1_train.npy 1779 ├── calms21_task1_test.npy 1780 ├── calms21_task1_test_features.npy 1781 ├── calms21_task1_test_features.npy 1782 ├── calms21_unlabeled_videos_part1.npy 1783 ├── calms21_unlabeled_videos_part1.npy 1784 ├── calms21_unlabeled_videos_part2.npy 1785 └── calms21_unlabeled_videos_part3.npy 1786 ``` 1787 """ 1788 1789 def __init__( 1790 self, 1791 video_order: List = None, 1792 data_path: Union[Set, str] = None, 1793 file_paths: Set = None, 1794 task_n: int = 1, 1795 include_task1: bool = True, 1796 feature_save_path: str = None, 1797 len_segment: int = 128, 1798 overlap: int = 0, 1799 feature_extraction: str = "kinematic", 1800 key_objects: Dict = None, 1801 treba_files: bool = False, 1802 num_cpus: int = None, 1803 feature_extraction_pars: Dict = None, 1804 *args, 1805 **kwargs, 1806 ) -> None: 1807 """ 1808 Parameters 1809 ---------- 1810 video_order : list, optional 1811 a list of video ids that should be processed in the same order (not passed if creating from key objects) 1812 data_path : str | set, optional 1813 the path to the folder where the pose and feature files are stored or a set of such paths 1814 (not passed if creating from key objects or from `file_paths`) 1815 file_paths : set, optional 1816 a set of string paths to the pose and feature files 1817 (not passed if creating from key objects or from `data_path`) 1818 task_n : [1, 2] 1819 the number of the task 1820 include_task1 : bool, default True 1821 include task 1 data to training set 1822 feature_save_path : str, optional 1823 the path to the folder where pre-processed files are stored (not passed if creating from key objects) 1824 len_segment : int, default 128 1825 the length of the segments in which the data should be cut (in frames) 1826 overlap : int, default 0 1827 the length of the overlap between neighboring segments (in frames) 1828 feature_extraction : str, default 'kinematic' 1829 the feature extraction method (see options.feature_extractors for available options) 1830 ignored_bodyparts : list, optional 1831 list of strings of bodypart names to ignore 1832 key_objects : tuple, optional 1833 a tuple of key objects 1834 treba_files : bool, default False 1835 if `True`, TREBA feature files will be loaded 1836 num_cpus : int, optional 1837 the number of cpus to use in data processing 1838 feature_extraction_pars : dict, optional 1839 parameters of the feature extractor 1840 """ 1841 1842 self.task_n = int(task_n) 1843 self.include_task1 = include_task1 1844 self.treba_files = treba_files 1845 if feature_extraction_pars is not None: 1846 feature_extraction_pars["interactive"] = True 1847 1848 super().__init__( 1849 video_order, 1850 data_path, 1851 file_paths, 1852 data_prefix=None, 1853 feature_suffix=None, 1854 convert_int_indices=False, 1855 feature_save_path=feature_save_path, 1856 canvas_shape=[1024, 570], 1857 len_segment=len_segment, 1858 overlap=overlap, 1859 feature_extraction=feature_extraction, 1860 ignored_clips=None, 1861 ignored_bodyparts=None, 1862 default_agent_name="ind0", 1863 key_objects=key_objects, 1864 likelihood_threshold=0, 1865 num_cpus=num_cpus, 1866 frame_limit=1, 1867 feature_extraction_pars=feature_extraction_pars, 1868 ) 1869 1870 @classmethod 1871 def get_file_ids( 1872 cls, 1873 task_n: int = 1, 1874 include_task1: bool = False, 1875 treba_files: bool = False, 1876 data_path: Union[str, Set] = None, 1877 file_paths=None, 1878 *args, 1879 **kwargs, 1880 ) -> Iterable: 1881 """ 1882 Process data parameters and return a list of ids of the videos that should 1883 be processed by the __init__ function 1884 1885 Parameters 1886 ---------- 1887 task_n : {1, 2, 3} 1888 the index of the CalMS21 challenge task 1889 include_task1 : bool, default False 1890 if `True`, the training file of the task 1 will be loaded 1891 treba_files : bool, default False 1892 if `True`, the TREBA feature files will be loaded 1893 filenames : set, optional 1894 a set of string filenames to search for (only basenames, not the whole paths) 1895 data_path : str | set, optional 1896 the path to the folder where the pose and feature files are stored or a set of such paths 1897 (not passed if creating from key objects or from `file_paths`) 1898 file_paths : set, optional 1899 a set of string paths to the pose and feature files 1900 (not passed if creating from key objects or from `data_path`) 1901 1902 Returns 1903 ------- 1904 video_ids : list 1905 a list of video file ids 1906 """ 1907 1908 task_n = int(task_n) 1909 if task_n == 1: 1910 include_task1 = False 1911 files = [] 1912 if treba_files: 1913 postfix = "_features" 1914 else: 1915 postfix = "" 1916 files.append(f"calms21_task{task_n}_train{postfix}.npy") 1917 files.append(f"calms21_task{task_n}_test{postfix}.npy") 1918 if include_task1: 1919 files.append(f"calms21_task1_train{postfix}.npy") 1920 for i in range(1, 5): 1921 files.append(f"calms21_unlabeled_videos_part{i}{postfix}.npy") 1922 filenames = set(files) 1923 return SequenceInputStore.get_file_ids(filenames, data_path, file_paths) 1924 1925 def _open_file(self, filename: str) -> List: 1926 """ 1927 Open a file and make a list of sequences 1928 1929 The sequence objects should contain information about all clips in one video. The sequences and 1930 video ids will be processed in the `_get_data` function. 1931 1932 Parameters 1933 ---------- 1934 filename : str 1935 the name of the file 1936 1937 Returns 1938 ------- 1939 video_tuples : list 1940 a list of video tuples: `(video_id, sequence)` 1941 """ 1942 1943 if os.path.basename(filename).startswith("calms21_unlabeled_videos"): 1944 mode = "unlabeled" 1945 elif os.path.basename(filename).startswith(f"calms21_task{self.task_n}_test"): 1946 mode = "test" 1947 else: 1948 mode = "train" 1949 data_dict = np.load(filename, allow_pickle=True).item() 1950 data = {} 1951 keys = list(data_dict.keys()) 1952 for key in keys: 1953 data.update(data_dict[key]) 1954 data_dict.pop(key) 1955 dict_list = [(f'{mode}--{k.split("/")[-1]}', v) for k, v in data.items()] 1956 return dict_list 1957 1958 def _get_data( 1959 self, video_id: str, sequence, default_agent_name: str 1960 ) -> Tuple[Dict, Optional[Dict]]: 1961 """ 1962 Get the keypoint dataframes from a sequence 1963 1964 The sequences and video ids are generated in the `_open_file` function. 1965 In `data_dictionary`, the keys are clip ids and the values are `pandas` dataframes with two-level indices. 1966 The first level is the frame numbers and the second is the body part names. The dataframes should have from 1967 two to four columns labeled `"x"`, `"y"` and (optionally) `"z"` and `"likelihood"`. Each frame should have 1968 information on all the body parts. You don't have to filter the data in any way or fill the nans, it will 1969 be done automatically. 1970 1971 Parameters 1972 ---------- 1973 video_id : str 1974 the video id 1975 sequence 1976 an object containing information about all clips in one video 1977 default_agent_name 1978 1979 Returns 1980 ------- 1981 data dictionary : dict 1982 a dictionary where the keys are clip ids and the values are keypoint dataframes (see above for details) 1983 metadata_dictionary : dict 1984 a dictionary where the keys are clip ids and the values are metadata objects (can be any additional information, 1985 like the annotator tag; for no metadata pass `None`) 1986 """ 1987 1988 if "metadata" in sequence: 1989 annotator = sequence["metadata"]["annotator-id"] 1990 else: 1991 annotator = 0 1992 bodyparts = [ 1993 "nose", 1994 "left ear", 1995 "right ear", 1996 "neck", 1997 "left hip", 1998 "right hip", 1999 "tail", 2000 ] 2001 columns = ["x", "y"] 2002 if "keypoints" in sequence: 2003 sequence = sequence["keypoints"] 2004 index = pd.MultiIndex.from_product([range(sequence.shape[0]), bodyparts]) 2005 data = { 2006 "mouse1": pd.DataFrame( 2007 data=(sequence[:, 0, :, :]).transpose((0, 2, 1)).reshape(-1, 2), 2008 columns=columns, 2009 index=index, 2010 ), 2011 "mouse2": pd.DataFrame( 2012 data=(sequence[:, 1, :, :]).transpose((0, 2, 1)).reshape(-1, 2), 2013 columns=columns, 2014 index=index, 2015 ), 2016 } 2017 else: 2018 sequence = sequence["features"] 2019 mice = sequence[:, :-32].reshape((-1, 2, 2, 7)) 2020 index = pd.MultiIndex.from_product([range(mice.shape[0]), bodyparts]) 2021 data = { 2022 "mouse1": pd.DataFrame( 2023 data=(mice[:, 0, :, :]).transpose((0, 2, 1)).reshape(-1, 2), 2024 columns=columns, 2025 index=index, 2026 ), 2027 "mouse2": pd.DataFrame( 2028 data=(mice[:, 1, :, :]).transpose((0, 2, 1)).reshape(-1, 2), 2029 columns=columns, 2030 index=index, 2031 ), 2032 "loaded": sequence[:, -32:], 2033 } 2034 metadata = {k: annotator for k in data.keys()} 2035 return data, metadata 2036 2037 2038class Numpy3DInputStore(FileInputStore): 2039 """ 2040 3D data 2041 2042 Assumes the data files to be `numpy` arrays saved in `.npy` format with shape `(#frames, #keypoints, 3)`. 2043 2044 Assumes the following file structure: 2045 ``` 2046 data_path 2047 ├── video1_suffix1.npy 2048 ├── video2_suffix2.npy 2049 ├── video1_features.pt 2050 └── video2_features.pt 2051 ``` 2052 Here `data_suffix` is `{'_suffix1.npy', '_suffix1.npy'}` and `feature_suffix` (optional) is `'_features.pt'`. 2053 2054 The feature files should to be dictionaries where keys are clip IDs (e.g. animal names) and values are 2055 feature values (arrays of shape `(#frames, #features)`). If the arrays are shaped as `(#features, #frames)`, 2056 set `transpose_features` to `True`. 2057 2058 The files can be saved with `numpy.save()` (with `.npy` extension), `torch.save()` (with `.pt` extension) or 2059 with `pickle.dump()` (with `.pickle` or `.pkl` extension). 2060 """ 2061 2062 def __init__( 2063 self, 2064 video_order: List = None, 2065 data_path: Union[Set, str] = None, 2066 file_paths: Set = None, 2067 data_suffix: Union[Set, str] = None, 2068 data_prefix: Union[Set, str] = None, 2069 feature_suffix: Union[Set, str] = None, 2070 convert_int_indices: bool = True, 2071 feature_save_path: str = None, 2072 canvas_shape: List = None, 2073 len_segment: int = 128, 2074 overlap: int = 0, 2075 feature_extraction: str = "kinematic", 2076 ignored_clips: List = None, 2077 ignored_bodyparts: List = None, 2078 default_agent_name: str = "ind0", 2079 key_objects: Dict = None, 2080 likelihood_threshold: float = 0, 2081 num_cpus: int = None, 2082 frame_limit: int = 1, 2083 feature_extraction_pars: Dict = None, 2084 centered: bool = False, 2085 **kwargs, 2086 ) -> None: 2087 """ 2088 Parameters 2089 ---------- 2090 video_order : list, optional 2091 a list of video ids that should be processed in the same order (not passed if creating from key objects 2092 data_path : str | set, optional 2093 the path to the folder where the pose and feature files are stored or a set of such paths 2094 (not passed if creating from key objects or from `file_paths`) 2095 file_paths : set, optional 2096 a set of string paths to the pose and feature files 2097 (not passed if creating from key objects or from `data_path`) 2098 data_suffix : str | set, optional 2099 the suffix or the set of suffices such that the pose files are named {video_id}{data_suffix} 2100 (not passed if creating from key objects or if irrelevant for the dataset) 2101 data_prefix : str | set, optional 2102 the prefix or the set of prefixes such that the pose files for different video views of the same 2103 clip are named {prefix}{sep}{video_id}{data_suffix} (not passed if creating from key objects 2104 or if irrelevant for the dataset) 2105 feature_suffix : str | set, optional 2106 the suffix or the set of suffices such that the additional feature files are named 2107 {video_id}{feature_suffix} (and placed at the data_path folder) 2108 convert_int_indices : bool, default True 2109 if `True`, convert any integer key `i` in feature files to `'ind{i}'` 2110 feature_save_path : str, optional 2111 the path to the folder where pre-processed files are stored (not passed if creating from key objects) 2112 canvas_shape : List, default [1, 1] 2113 the canvas size where the pose is defined 2114 len_segment : int, default 128 2115 the length of the segments in which the data should be cut (in frames) 2116 overlap : int, default 0 2117 the length of the overlap between neighboring segments (in frames) 2118 feature_extraction : str, default 'kinematic' 2119 the feature extraction method (see options.feature_extractors for available options) 2120 ignored_clips : list, optional 2121 list of strings of clip ids to ignore 2122 ignored_bodyparts : list, optional 2123 list of strings of bodypart names to ignore 2124 default_agent_name : str, default 'ind0' 2125 the agent name used as default in the pose files for a single agent 2126 key_objects : tuple, optional 2127 a tuple of key objects 2128 likelihood_threshold : float, default 0 2129 coordinate values with likelihoods less than this value will be set to 'unknown' 2130 num_cpus : int, optional 2131 the number of cpus to use in data processing 2132 frame_limit : int, default 1 2133 clips shorter than this number of frames will be ignored 2134 feature_extraction_pars : dict, optional 2135 parameters of the feature extractor 2136 """ 2137 2138 super().__init__( 2139 video_order, 2140 data_path, 2141 file_paths, 2142 data_suffix=data_suffix, 2143 data_prefix=data_prefix, 2144 feature_suffix=feature_suffix, 2145 convert_int_indices=convert_int_indices, 2146 feature_save_path=feature_save_path, 2147 canvas_shape=canvas_shape, 2148 len_segment=len_segment, 2149 overlap=overlap, 2150 feature_extraction=feature_extraction, 2151 ignored_clips=ignored_clips, 2152 ignored_bodyparts=ignored_bodyparts, 2153 default_agent_name=default_agent_name, 2154 key_objects=key_objects, 2155 likelihood_threshold=likelihood_threshold, 2156 num_cpus=num_cpus, 2157 frame_limit=frame_limit, 2158 feature_extraction_pars=feature_extraction_pars, 2159 centered=centered, 2160 ) 2161 2162 def _open_data( 2163 self, filename: str, default_clip_name: str 2164 ) -> Tuple[Dict, Optional[Dict]]: 2165 """ 2166 Load the keypoints from filename and organize them in a dictionary 2167 2168 In `data_dictionary`, the keys are clip ids and the values are `pandas` dataframes with two-level indices. 2169 The first level is the frame numbers and the second is the body part names. The dataframes should have from 2170 two to four columns labeled `"x"`, `"y"` and (optionally) `"z"` and `"likelihood"`. Each frame should have 2171 information on all the body parts. You don't have to filter the data in any way or fill the nans, it will 2172 be done automatically. 2173 2174 Parameters 2175 ---------- 2176 filename : str 2177 path to the pose file 2178 default_clip_name : str 2179 the name to assign to a clip if it does not have a name in the raw data 2180 2181 Returns 2182 ------- 2183 data dictionary : dict 2184 a dictionary where the keys are clip ids and the values are keypoint dataframes (see above for details) 2185 metadata_dictionary : dict 2186 a dictionary where the keys are clip ids and the values are metadata objects (can be any additional information, 2187 like the annotator tag; for no metadata pass `None`) 2188 """ 2189 2190 data = np.load(filename) 2191 bodyparts = [str(i) for i in range(data.shape[1])] 2192 clip_id = self.default_agent_name 2193 columns = ["x", "y", "z"] 2194 index = pd.MultiIndex.from_product([range(data.shape[0]), bodyparts]) 2195 data_dict = { 2196 clip_id: pd.DataFrame( 2197 data=data.reshape(-1, 3), columns=columns, index=index 2198 ) 2199 } 2200 return data_dict, None 2201 2202 2203class LoadedFeaturesInputStore(GeneralInputStore): 2204 """ 2205 Non-pose feature files 2206 2207 The feature files should to be dictionaries where keys are clip IDs (e.g. animal names) and values are 2208 feature values (arrays of shape `(#frames, #features)`). If the arrays are shaped as `(#features, #frames)`, 2209 set `transpose_features` to `True`. 2210 2211 The files can be saved with `numpy.save()` (with `.npy` extension), `torch.save()` (with `.pt` extension) or 2212 with `pickle.dump()` (with `.pickle` or `.pkl` extension). 2213 2214 Assumes the following file structure: 2215 ``` 2216 data_path 2217 ├── video1_features.pt 2218 └── video2_features.pt 2219 ``` 2220 Here `feature_suffix` (optional) is `'_features.pt'`. 2221 """ 2222 2223 def __init__( 2224 self, 2225 video_order: List = None, 2226 data_path: Union[Set, str] = None, 2227 file_paths: Set = None, 2228 feature_suffix: Union[Set, str] = None, 2229 convert_int_indices: bool = True, 2230 feature_save_path: str = None, 2231 len_segment: int = 128, 2232 overlap: int = 0, 2233 ignored_clips: List = None, 2234 key_objects: Dict = None, 2235 num_cpus: int = None, 2236 frame_limit: int = 1, 2237 transpose_features: bool = False, 2238 **kwargs, 2239 ) -> None: 2240 """ 2241 Parameters 2242 ---------- 2243 video_order : list, optional 2244 a list of video ids that should be processed in the same order (not passed if creating from key objects 2245 data_path : str | set, optional 2246 the path to the folder where the pose and feature files are stored or a set of such paths 2247 (not passed if creating from key objects or from `file_paths`) 2248 file_paths : set, optional 2249 a set of string paths to the pose and feature files 2250 (not passed if creating from key objects or from `data_path`) 2251 feature_suffix : str | set, optional 2252 the suffix or the set of suffices such that the additional feature files are named 2253 {video_id}{feature_suffix} (and placed at the data_path folder) 2254 feature_save_path : str, optional 2255 the path to the folder where pre-processed files are stored (not passed if creating from key objects) 2256 len_segment : int, default 128 2257 the length of the segments in which the data should be cut (in frames) 2258 overlap : int, default 0 2259 the length of the overlap between neighboring segments (in frames) 2260 ignored_clips : list, optional 2261 list of strings of clip ids to ignore 2262 default_agent_name : str, default 'ind0' 2263 the agent name used as default in the pose files for a single agent 2264 key_objects : tuple, optional 2265 a tuple of key objects 2266 num_cpus : int, optional 2267 the number of cpus to use in data processing 2268 frame_limit : int, default 1 2269 clips shorter than this number of frames will be ignored 2270 feature_extraction_pars : dict, optional 2271 parameters of the feature extractor 2272 """ 2273 2274 super().__init__( 2275 video_order, 2276 data_path, 2277 file_paths, 2278 feature_suffix=feature_suffix, 2279 convert_int_indices=convert_int_indices, 2280 feature_save_path=feature_save_path, 2281 len_segment=len_segment, 2282 overlap=overlap, 2283 ignored_clips=ignored_clips, 2284 key_objects=key_objects, 2285 num_cpus=num_cpus, 2286 frame_limit=frame_limit, 2287 transpose_features=transpose_features, 2288 ) 2289 2290 def get_visibility( 2291 self, video_id: str, clip_id: str, start: int, end: int, score: int 2292 ) -> float: 2293 """ 2294 Get the fraction of the frames in that have a visibility score better than a hard_threshold 2295 2296 For example, in the case of keypoint data the visibility score can be the number of identified keypoints. 2297 2298 Parameters 2299 ---------- 2300 video_id : str 2301 the video id of the frames 2302 clip_id : str 2303 the clip id of the frames 2304 start : int 2305 the start frame 2306 end : int 2307 the end frame 2308 score : float 2309 the visibility score hard_threshold 2310 2311 Returns 2312 ------- 2313 frac_visible: float 2314 the fraction of frames with visibility above the hard_threshold 2315 """ 2316 2317 return 1 2318 2319 def _generate_features( 2320 self, video_id: str 2321 ) -> Tuple[Dict, Dict, Dict, Union[str, int]]: 2322 """ 2323 Generate features from the raw coordinates 2324 """ 2325 2326 features = defaultdict(lambda: {}) 2327 loaded_features = self._load_saved_features(video_id) 2328 min_frames = None 2329 max_frames = None 2330 video_tag = None 2331 for clip_id, feature_tensor in loaded_features.items(): 2332 if clip_id == "max_frames": 2333 max_frames = feature_tensor 2334 elif clip_id == "min_frames": 2335 min_frames = feature_tensor 2336 elif clip_id == "video_tag": 2337 video_tag = feature_tensor 2338 else: 2339 if not isinstance(feature_tensor, torch.Tensor): 2340 feature_tensor = torch.tensor(feature_tensor) 2341 if self.convert_int_indices and ( 2342 isinstance(clip_id, int) or isinstance(clip_id, np.integer) 2343 ): 2344 clip_id = f"ind{clip_id}" 2345 key = f"{os.path.basename(video_id)}---{clip_id}" 2346 features[key]["loaded"] = feature_tensor 2347 if min_frames is None: 2348 min_frames = {} 2349 for key, value in features.items(): 2350 video_id, clip_id = key.split("---") 2351 min_frames[clip_id] = 0 2352 if max_frames is None: 2353 max_frames = {} 2354 for key, value in features.items(): 2355 video_id, clip_id = key.split("---") 2356 max_frames[clip_id] = value["loaded"].shape[0] - 1 + min_frames[clip_id] 2357 return features, min_frames, max_frames, video_tag 2358 2359 def _load_data(self) -> np.array: 2360 """ 2361 Load input data and generate data prompts 2362 """ 2363 2364 if self.video_order is None: 2365 return None 2366 2367 files = [] 2368 for video_id in self.video_order: 2369 for f in self.file_paths: 2370 if f.endswith(tuple(self.feature_suffix)): 2371 bn = os.path.basename(f) 2372 if video_id == strip_suffix(bn, self.feature_suffix): 2373 files.append(f) 2374 2375 def make_data_dictionary(filename): 2376 name = strip_suffix(filename, self.feature_suffix) 2377 name = os.path.basename(name) 2378 data_dict, min_frames, max_frames, video_tag = self._generate_features(name) 2379 bp_dict = defaultdict(lambda: {}) 2380 for key, value in data_dict.items(): 2381 video_id, clip_id = key.split("---") 2382 bp_dict[video_id][clip_id] = 1 2383 min_frames = {name: min_frames} # name is e.g. 20190707T1126-1226 2384 max_frames = {name: max_frames} 2385 names, lengths, coords = self._make_trimmed_data(data_dict) 2386 return names, lengths, coords, bp_dict, min_frames, max_frames, video_tag 2387 2388 dict_list = p_map(make_data_dictionary, files, num_cpus=self.num_cpus) 2389 # dict_list = tqdm([make_data_dictionary(f) for f in files]) 2390 2391 self.visibility = {} 2392 self.min_frames = {} 2393 self.max_frames = {} 2394 self.original_coordinates = [] 2395 self.metadata = [] 2396 X = [] 2397 for ( 2398 names, 2399 lengths, 2400 coords, 2401 bp_dictionary, 2402 min_frames, 2403 max_frames, 2404 metadata, 2405 ) in dict_list: 2406 X += names 2407 self.original_coordinates += coords 2408 self.visibility.update(bp_dictionary) 2409 self.min_frames.update(min_frames) 2410 self.max_frames.update(max_frames) 2411 if metadata is not None: 2412 self.metadata += metadata 2413 del dict_list 2414 if len(self.metadata) != len(self.original_coordinates): 2415 self.metadata = None 2416 else: 2417 self.metadata = np.array(self.metadata) 2418 2419 self.min_frames = dict(self.min_frames) 2420 self.max_frames = dict(self.max_frames) 2421 self.original_coordinates = np.array(self.original_coordinates) 2422 return np.array(X) 2423 2424 @classmethod 2425 def get_file_ids( 2426 cls, 2427 data_path: Union[Set, str] = None, 2428 file_paths: Set = None, 2429 feature_suffix: Set = None, 2430 *args, 2431 **kwargs, 2432 ) -> List: 2433 """ 2434 Process data parameters and return a list of ids of the videos that should 2435 be processed by the __init__ function 2436 2437 Parameters 2438 ---------- 2439 data_suffix : set | str, optional 2440 the suffix (or a set of suffixes) of the input data files 2441 data_path : set | str, optional 2442 the path to the folder where the pose and feature files are stored or a set of such paths 2443 (not passed if creating from key objects or from `file_paths`) 2444 data_prefix : set | str, optional 2445 the prefix or the set of prefixes such that the pose files for different video views of the same 2446 clip are named {prefix}{sep}{video_id}{data_suffix} (not passed if creating from key objects 2447 or if irrelevant for the dataset) 2448 file_paths : set, optional 2449 a set of string paths to the pose and feature files 2450 feature_suffix : str | set, optional 2451 the suffix or the set of suffices such that the additional feature files are named 2452 {video_id}{feature_suffix} (and placed at the `data_path` folder or at `file_paths`) 2453 2454 Returns 2455 ------- 2456 video_ids : list 2457 a list of video file ids 2458 """ 2459 2460 if feature_suffix is None: 2461 feature_suffix = [] 2462 if isinstance(feature_suffix, str): 2463 feature_suffix = [feature_suffix] 2464 feature_suffix = tuple(feature_suffix) 2465 if file_paths is None: 2466 file_paths = [] 2467 if data_path is not None: 2468 if isinstance(data_path, str): 2469 data_path = [data_path] 2470 file_paths = [] 2471 for folder in data_path: 2472 file_paths += [os.path.join(folder, x) for x in os.listdir(folder)] 2473 ids = set() 2474 for f in file_paths: 2475 if f.endswith(feature_suffix): 2476 bn = os.path.basename(f) 2477 video_id = strip_suffix(bn, feature_suffix) 2478 ids.add(video_id) 2479 ids = sorted(ids) 2480 return ids 2481 2482 2483class SIMBAInputStore(FileInputStore): 2484 """ 2485 SIMBA paper format data 2486 2487 Assumes the following file structure 2488 2489 ``` 2490 data_path 2491 ├── Video1.csv 2492 ... 2493 └── Video9.csv 2494 ``` 2495 Here `data_suffix` is `.csv`. 2496 """ 2497 2498 def __init__( 2499 self, 2500 video_order: List = None, 2501 data_path: Union[Set, str] = None, 2502 file_paths: Set = None, 2503 data_prefix: Union[Set, str] = None, 2504 feature_suffix: str = None, 2505 feature_save_path: str = None, 2506 canvas_shape: List = None, 2507 len_segment: int = 128, 2508 overlap: int = 0, 2509 feature_extraction: str = "kinematic", 2510 ignored_clips: List = None, 2511 ignored_bodyparts: List = None, 2512 key_objects: Tuple = None, 2513 likelihood_threshold: float = 0, 2514 num_cpus: int = None, 2515 normalize: bool = False, 2516 feature_extraction_pars: Dict = None, 2517 centered: bool = False, 2518 data_suffix: str = None, 2519 use_features: bool = False, 2520 *args, 2521 **kwargs, 2522 ) -> None: 2523 """ 2524 Parameters 2525 ---------- 2526 video_order : list, optional 2527 a list of video ids that should be processed in the same order (not passed if creating from key objects 2528 data_path : str | set, optional 2529 the path to the folder where the pose and feature files are stored or a set of such paths 2530 (not passed if creating from key objects or from `file_paths`) 2531 file_paths : set, optional 2532 a set of string paths to the pose and feature files 2533 (not passed if creating from key objects or from `data_path`) 2534 data_suffix : str | set, optional 2535 the suffix or the set of suffices such that the pose files are named {video_id}{data_suffix} 2536 (not passed if creating from key objects or if irrelevant for the dataset) 2537 data_prefix : str | set, optional 2538 the prefix or the set of prefixes such that the pose files for different video views of the same 2539 clip are named {prefix}{sep}{video_id}{data_suffix} (not passed if creating from key objects 2540 or if irrelevant for the dataset) 2541 feature_suffix : str | set, optional 2542 the suffix or the set of suffices such that the additional feature files are named 2543 {video_id}{feature_suffix} (and placed at the data_path folder) 2544 feature_save_path : str, optional 2545 the path to the folder where pre-processed files are stored (not passed if creating from key objects) 2546 canvas_shape : List, default [1, 1] 2547 the canvas size where the pose is defined 2548 len_segment : int, default 128 2549 the length of the segments in which the data should be cut (in frames) 2550 overlap : int, default 0 2551 the length of the overlap between neighboring segments (in frames) 2552 feature_extraction : str, default 'kinematic' 2553 the feature extraction method (see options.feature_extractors for available options) 2554 ignored_clips : list, optional 2555 list of strings of clip ids to ignore 2556 ignored_bodyparts : list, optional 2557 list of strings of bodypart names to ignore 2558 key_objects : tuple, optional 2559 a tuple of key objects 2560 likelihood_threshold : float, default 0 2561 coordinate values with likelihoods less than this value will be set to 'unknown' 2562 num_cpus : int, optional 2563 the number of cpus to use in data processing 2564 feature_extraction_pars : dict, optional 2565 parameters of the feature extractor 2566 """ 2567 2568 self.use_features = use_features 2569 if feature_extraction_pars is not None: 2570 feature_extraction_pars["interactive"] = True 2571 super().__init__( 2572 video_order=video_order, 2573 data_path=data_path, 2574 file_paths=file_paths, 2575 data_suffix=data_suffix, 2576 data_prefix=data_prefix, 2577 feature_suffix=feature_suffix, 2578 convert_int_indices=False, 2579 feature_save_path=feature_save_path, 2580 canvas_shape=canvas_shape, 2581 len_segment=len_segment, 2582 overlap=overlap, 2583 feature_extraction=feature_extraction, 2584 ignored_clips=ignored_clips, 2585 ignored_bodyparts=ignored_bodyparts, 2586 default_agent_name="", 2587 key_objects=key_objects, 2588 likelihood_threshold=likelihood_threshold, 2589 num_cpus=num_cpus, 2590 min_frames=0, 2591 normalize=normalize, 2592 feature_extraction_pars=feature_extraction_pars, 2593 centered=centered, 2594 ) 2595 2596 def _open_data( 2597 self, filename: str, default_clip_name: str 2598 ) -> Tuple[Dict, Optional[Dict]]: 2599 data = pd.read_csv(filename) 2600 output = {} 2601 column_dict = {"x": "x", "y": "y", "z": "z", "p": "likelihood"} 2602 columns = [x for x in data.columns if x.split("_")[-1] in column_dict] 2603 animals = sorted(set([x.split("_")[-2] for x in columns])) 2604 coords = sorted(set([x.split("_")[-1] for x in columns])) 2605 names = sorted(set(["_".join(x.split("_")[:-2]) for x in columns])) 2606 for animal in animals: 2607 data_dict = {} 2608 for i, row in data.iterrows(): 2609 for col_name in names: 2610 data_dict[(i, col_name)] = [ 2611 row[f"{col_name}_{animal}_{coord}"] for coord in coords 2612 ] 2613 output[animal] = pd.DataFrame(data_dict).T 2614 output[animal].columns = [column_dict[x] for x in coords] 2615 if self.use_features: 2616 columns_to_avoid = [ 2617 x 2618 for x in data.columns 2619 if x.split("_")[-1] in column_dict 2620 or x.split("_")[-1].startswith("prediction") 2621 ] 2622 columns_to_avoid += ["scorer", "frames", "video_no"] 2623 output["loaded"] = ( 2624 data[[x for x in data.columns if x not in columns_to_avoid]] 2625 .interpolate() 2626 .values 2627 ) 2628 return output, None
29class GeneralInputStore(PoseInputStore): 30 """ 31 A generalized realization of a PoseInputStore 32 33 Assumes the following file structure: 34 ``` 35 data_path 36 ├── video1DLC1000.pickle 37 ├── video2DLC400.pickle 38 ├── video1_features.pt 39 └── video2_features.pt 40 ``` 41 Here `data_suffix` is `{'DLC1000.pickle', 'DLC400.pickle'}` and `feature_suffix` (optional) is `'_features.pt'`. 42 """ 43 44 data_suffix = None 45 46 def __init__( 47 self, 48 video_order: List = None, 49 data_path: Union[Set, str] = None, 50 file_paths: Set = None, 51 data_suffix: Union[Set, str] = None, 52 data_prefix: Union[Set, str] = None, 53 feature_suffix: str = None, 54 convert_int_indices: bool = True, 55 feature_save_path: str = None, 56 canvas_shape: List = None, 57 len_segment: int = 128, 58 overlap: int = 0, 59 feature_extraction: str = "kinematic", 60 ignored_clips: List = None, 61 ignored_bodyparts: List = None, 62 default_agent_name: str = "ind0", 63 key_objects: Tuple = None, 64 likelihood_threshold: float = 0, 65 num_cpus: int = None, 66 frame_limit: int = 1, 67 normalize: bool = False, 68 feature_extraction_pars: Dict = None, 69 centered: bool = False, 70 transpose_features: bool = False, 71 *args, 72 **kwargs, 73 ) -> None: 74 """ 75 Parameters 76 ---------- 77 video_order : list, optional 78 a list of video ids that should be processed in the same order (not passed if creating from key objects 79 data_path : str | set, optional 80 the path to the folder where the pose and feature files are stored or a set of such paths 81 (not passed if creating from key objects or from `file_paths`) 82 file_paths : set, optional 83 a set of string paths to the pose and feature files 84 (not passed if creating from key objects or from `data_path`) 85 data_suffix : str | set, optional 86 the suffix or the set of suffices such that the pose files are named {video_id}{data_suffix} 87 (not passed if creating from key objects or if irrelevant for the dataset) 88 data_prefix : str | set, optional 89 the prefix or the set of prefixes such that the pose files for different video views of the same 90 clip are named {prefix}{sep}{video_id}{data_suffix} (not passed if creating from key objects 91 or if irrelevant for the dataset) 92 feature_suffix : str | set, optional 93 the suffix or the set of suffices such that the additional feature files are named 94 {video_id}{feature_suffix} (and placed at the data_path folder) 95 convert_int_indices : bool, default True 96 if `True`, convert any integer key `i` in feature files to `'ind{i}'` 97 feature_save_path : str, optional 98 the path to the folder where pre-processed files are stored (not passed if creating from key objects) 99 canvas_shape : List, default [1, 1] 100 the canvas size where the pose is defined 101 len_segment : int, default 128 102 the length of the segments in which the data should be cut (in frames) 103 overlap : int, default 0 104 the length of the overlap between neighboring segments (in frames) 105 feature_extraction : str, default 'kinematic' 106 the feature extraction method (see options.feature_extractors for available options) 107 ignored_clips : list, optional 108 list of strings of clip ids to ignore 109 ignored_bodyparts : list, optional 110 list of strings of bodypart names to ignore 111 default_agent_name : str, default 'ind0' 112 the agent name used as default in the pose files for a single agent 113 key_objects : tuple, optional 114 a tuple of key objects 115 likelihood_threshold : float, default 0 116 coordinate values with likelihoods less than this value will be set to 'unknown' 117 num_cpus : int, optional 118 the number of cpus to use in data processing 119 frame_limit : int, default 1 120 clips shorter than this number of frames will be ignored 121 feature_extraction_pars : dict, optional 122 parameters of the feature extractor 123 """ 124 125 super().__init__() 126 self.loaded_max = 0 127 if feature_extraction_pars is None: 128 feature_extraction_pars = {} 129 if ignored_clips is None: 130 ignored_clips = [] 131 self.bodyparts = [] 132 self.visibility = None 133 self.normalize = normalize 134 135 if canvas_shape is None: 136 canvas_shape = [1, 1] 137 if isinstance(data_suffix, str): 138 data_suffix = [data_suffix] 139 if isinstance(data_prefix, str): 140 data_prefix = [data_prefix] 141 if isinstance(data_path, str): 142 data_path = [data_path] 143 if isinstance(feature_suffix, str): 144 feature_suffix = [feature_suffix] 145 146 self.video_order = video_order 147 self.centered = centered 148 self.feature_extraction = feature_extraction 149 self.len_segment = int(len_segment) 150 self.data_suffices = data_suffix 151 self.data_prefixes = data_prefix 152 self.feature_suffix = feature_suffix 153 self.convert_int_indices = convert_int_indices 154 if overlap < 1: 155 overlap = overlap * len_segment 156 self.overlap = int(overlap) 157 self.canvas_shape = canvas_shape 158 self.default_agent_name = default_agent_name 159 self.feature_save_path = feature_save_path 160 self.data_suffices = data_suffix 161 self.data_prefixes = data_prefix 162 self.likelihood_threshold = likelihood_threshold 163 self.num_cpus = num_cpus 164 self.frame_limit = frame_limit 165 self.transpose = transpose_features 166 167 self.ram = False 168 self.min_frames = {} 169 self.original_coordinates = np.array([]) 170 171 self.file_paths = self._get_file_paths(file_paths, data_path) 172 173 self.extractor = options.feature_extractors[self.feature_extraction]( 174 self, 175 **feature_extraction_pars, 176 ) 177 178 self.canvas_center = np.array(canvas_shape) // 2 179 180 if ignored_clips is not None: 181 self.ignored_clips = ignored_clips 182 else: 183 self.ignored_clips = [] 184 if ignored_bodyparts is not None: 185 self.ignored_bodyparts = ignored_bodyparts 186 else: 187 self.ignored_bodyparts = [] 188 189 self.step = self.len_segment - self.overlap 190 if self.step < 0: 191 raise ValueError( 192 f"The overlap value ({self.overlap}) cannot be larger than len_segment ({self.len_segment}" 193 ) 194 195 if self.feature_save_path is None and data_path is not None: 196 self.feature_save_path = os.path.join(data_path[0], "trimmed") 197 198 if key_objects is None and self.video_order is not None: 199 print("Computing input features...") 200 self.data = self._load_data() 201 elif key_objects is not None: 202 self.load_from_key_objects(key_objects) 203 204 def __getitem__(self, ind: int) -> Dict: 205 prompt = self.data[ind] 206 if not self.ram: 207 with open(prompt, "rb") as f: 208 prompt = pickle.load(f) 209 return prompt 210 211 def __len__(self) -> int: 212 if self.data is None: 213 raise RuntimeError("The input store data has not been initialized!") 214 return len(self.data) 215 216 @classmethod 217 def _get_file_paths(cls, file_paths: Set, data_path: Union[str, Set]) -> List: 218 """ 219 Get a set of relevant files 220 221 Parameters 222 ---------- 223 file_paths : set 224 a set of filepaths to include 225 data_path : str | set 226 the path to a folder that contains relevant files (a single path or a set) 227 228 Returns 229 ------- 230 file_paths : list 231 a list of relevant file paths (input and feature files that follow the dataset naming pattern) 232 """ 233 234 if file_paths is None: 235 file_paths = [] 236 file_paths = list(file_paths) 237 if data_path is not None: 238 if isinstance(data_path, str): 239 data_path = [data_path] 240 for folder in data_path: 241 file_paths += [os.path.join(folder, x) for x in os.listdir(folder)] 242 return file_paths 243 244 def get_folder(self, video_id: str) -> str: 245 """ 246 Get the input folder that the file with this video id was read from 247 248 Parameters 249 ---------- 250 video_id : str 251 the video id 252 253 Returns 254 ------- 255 folder : str 256 the path to the directory that contains the input file associated with the video id 257 """ 258 259 for file in self.file_paths: 260 if ( 261 strip_prefix( 262 strip_suffix(os.path.basename(file), self.data_suffices), 263 self.data_prefixes, 264 ) 265 == video_id 266 ): 267 return os.path.dirname(file) 268 269 def remove(self, indices: List) -> None: 270 """ 271 Remove the samples corresponding to indices 272 273 Parameters 274 ---------- 275 indices : int 276 a list of integer indices to remove 277 """ 278 279 if len(indices) > 0: 280 mask = np.ones(len(self.original_coordinates)) 281 mask[indices] = 0 282 mask = mask.astype(bool) 283 for file in self.data[~mask]: 284 os.remove(file) 285 self.original_coordinates = self.original_coordinates[mask] 286 self.data = self.data[mask] 287 if self.metadata is not None: 288 self.metadata = self.metadata[mask] 289 290 def key_objects(self) -> Tuple: 291 """ 292 Return a tuple of the key objects necessary to re-create the Store 293 294 Returns 295 ------- 296 key_objects : tuple 297 a tuple of key objects 298 """ 299 300 for k, v in self.min_frames.items(): 301 self.min_frames[k] = dict(v) 302 for k, v in self.max_frames.items(): 303 self.max_frames[k] = dict(v) 304 return ( 305 self.original_coordinates, 306 dict(self.min_frames), 307 dict(self.max_frames), 308 self.data, 309 self.visibility, 310 self.step, 311 self.file_paths, 312 self.len_segment, 313 self.metadata, 314 ) 315 316 def load_from_key_objects(self, key_objects: Tuple) -> None: 317 """ 318 Load the information from a tuple of key objects 319 320 Parameters 321 ---------- 322 key_objects : tuple 323 a tuple of key objects 324 """ 325 326 ( 327 self.original_coordinates, 328 self.min_frames, 329 self.max_frames, 330 self.data, 331 self.visibility, 332 self.step, 333 self.file_paths, 334 self.len_segment, 335 self.metadata, 336 ) = key_objects 337 338 def to_ram(self) -> None: 339 """ 340 Transfer the data samples to RAM if they were previously stored as file paths 341 """ 342 343 if self.ram: 344 return 345 346 data = p_map(lambda x: self[x], list(range(len(self))), num_cpus=self.num_cpus) 347 # data = [load(x) for x in self.data] 348 self.data = TensorDict(data) 349 self.ram = True 350 351 def get_original_coordinates(self) -> np.ndarray: 352 """ 353 Return the original coordinates array 354 355 Returns 356 ------- 357 np.ndarray 358 an array that contains the coordinates of the data samples in original input data (video id, clip id, 359 start frame) 360 """ 361 362 return self.original_coordinates 363 364 def create_subsample(self, indices: List, ssl_indices: List = None): 365 """ 366 Create a new store that contains a subsample of the data 367 368 Parameters 369 ---------- 370 indices : list 371 the indices to be included in the subsample 372 ssl_indices : list, optional 373 the indices to be included in the subsample without the annotation data 374 """ 375 376 if ssl_indices is None: 377 ssl_indices = [] 378 new = self.new() 379 new.original_coordinates = self.original_coordinates[indices + ssl_indices] 380 new.min_frames = self.min_frames 381 new.max_frames = self.max_frames 382 new.data = self.data[indices + ssl_indices] 383 new.visibility = self.visibility 384 new.step = self.step 385 new.file_paths = self.file_paths 386 new.len_segment = self.len_segment 387 if self.metadata is None: 388 new.metadata = None 389 else: 390 new.metadata = self.metadata[indices + ssl_indices] 391 return new 392 393 def get_video_id(self, coords: Tuple) -> str: 394 """ 395 Get the video id from an element of original coordinates 396 397 Parameters 398 ---------- 399 coords : tuple 400 an element of the original coordinates array 401 402 Returns 403 ------- 404 video_id: str 405 the id of the video that the coordinates point to 406 """ 407 408 video_name = coords[0].split("---")[0] 409 return video_name 410 411 def get_clip_id(self, coords: Tuple) -> str: 412 """ 413 Get the clip id from an element of original coordinates 414 415 Parameters 416 ---------- 417 coords : tuple 418 an element of the original coordinates array 419 420 Returns 421 ------- 422 clip_id : str 423 the id of the clip that the coordinates point to 424 """ 425 426 clip_id = coords[0].split("---")[1] 427 return clip_id 428 429 def get_clip_length(self, video_id: str, clip_id: str) -> int: 430 """ 431 Get the clip length from the id 432 433 Parameters 434 ---------- 435 video_id : str 436 the video id 437 clip_id : str 438 the clip id 439 440 Returns 441 ------- 442 clip_length : int 443 the length of the clip 444 """ 445 446 inds = clip_id.split("+") 447 max_frame = min([self.max_frames[video_id][x] for x in inds]) 448 min_frame = max([self.min_frames[video_id][x] for x in inds]) 449 return max_frame - min_frame + 1 450 451 def get_clip_start_end(self, coords: Tuple) -> Tuple[int, int]: 452 """ 453 Get the clip start and end frames from an element of original coordinates 454 455 Parameters 456 ---------- 457 coords : tuple 458 an element of original coordinates array 459 460 Returns 461 ------- 462 start : int 463 the start frame of the clip that the coordinates point to 464 end : int 465 the end frame of the clip that the coordinates point to 466 """ 467 468 l = self.get_clip_length_from_coords(coords) 469 i = coords[1] 470 start = int(i) * self.step 471 end = min(start + self.len_segment, l) 472 return start, end 473 474 def get_clip_start(self, video_name: str, clip_id: str) -> int: 475 """ 476 Get the clip start frame from the video id and the clip id 477 478 Parameters 479 ---------- 480 video_name : str 481 the video id 482 clip_id : str 483 the clip id 484 485 Returns 486 ------- 487 clip_start : int 488 the start frame of the clip 489 """ 490 491 return max( 492 [self.min_frames[video_name][clip_id_k] for clip_id_k in clip_id.split("+")] 493 ) 494 495 def get_visibility( 496 self, video_id: str, clip_id: str, start: int, end: int, score: int 497 ) -> float: 498 """ 499 Get the fraction of the frames in that have a visibility score better than a hard_threshold 500 501 For example, in the case of keypoint data the visibility score can be the number of identified keypoints. 502 503 Parameters 504 ---------- 505 video_id : str 506 the video id of the frames 507 clip_id : str 508 the clip id of the frames 509 start : int 510 the start frame 511 end : int 512 the end frame 513 score : float 514 the visibility score hard_threshold 515 516 Returns 517 ------- 518 frac_visible: float 519 the fraction of frames with visibility above the hard_threshold 520 """ 521 522 s = 0 523 for ind_k in clip_id.split("+"): 524 s += np.sum(self.visibility[video_id][ind_k][start:end] > score) / ( 525 end - start 526 ) 527 return s / len(clip_id.split("+")) 528 529 def get_annotation_objects(self) -> Dict: 530 """ 531 Get a dictionary of objects necessary to create an AnnotationStore 532 533 Returns 534 ------- 535 annotation_objects : dict 536 a dictionary of objects to be passed to the AnnotationStore constructor where the keys are the names of 537 the objects 538 """ 539 540 min_frames = self.min_frames 541 max_frames = self.max_frames 542 num_bp = self.visibility 543 return { 544 "min_frames": min_frames, 545 "max_frames": max_frames, 546 "visibility": num_bp, 547 } 548 549 @classmethod 550 def get_file_ids( 551 cls, 552 data_suffix: Union[Set, str] = None, 553 data_path: Union[Set, str] = None, 554 data_prefix: Union[Set, str] = None, 555 file_paths: Set = None, 556 feature_suffix: Set = None, 557 *args, 558 **kwargs, 559 ) -> List: 560 """ 561 Process data parameters and return a list of ids of the videos that should 562 be processed by the __init__ function 563 564 Parameters 565 ---------- 566 data_suffix : set | str, optional 567 the suffix (or a set of suffixes) of the input data files 568 data_path : set | str, optional 569 the path to the folder where the pose and feature files are stored or a set of such paths 570 (not passed if creating from key objects or from `file_paths`) 571 data_prefix : set | str, optional 572 the prefix or the set of prefixes such that the pose files for different video views of the same 573 clip are named {prefix}{sep}{video_id}{data_suffix} (not passed if creating from key objects 574 or if irrelevant for the dataset) 575 file_paths : set, optional 576 a set of string paths to the pose and feature files 577 feature_suffix : str | set, optional 578 the suffix or the set of suffices such that the additional feature files are named 579 {video_id}{feature_suffix} (and placed at the `data_path` folder or at `file_paths`) 580 581 Returns 582 ------- 583 video_ids : list 584 a list of video file ids 585 """ 586 587 if data_suffix is None: 588 if cls.data_suffix is not None: 589 data_suffix = cls.data_suffix 590 else: 591 raise ValueError("Cannot get video ids without the data suffix!") 592 if feature_suffix is None: 593 feature_suffix = [] 594 if data_prefix is None: 595 data_prefix = "" 596 if isinstance(data_suffix, str): 597 data_suffix = [data_suffix] 598 else: 599 data_suffix = [x for x in data_suffix] 600 data_suffix = tuple(data_suffix) 601 if isinstance(data_prefix, str): 602 data_prefix = data_prefix 603 else: 604 data_prefix = tuple([x for x in data_prefix]) 605 if isinstance(feature_suffix, str): 606 feature_suffix = [feature_suffix] 607 if file_paths is None: 608 file_paths = [] 609 if data_path is not None: 610 if isinstance(data_path, str): 611 data_path = [data_path] 612 file_paths = [] 613 for folder in data_path: 614 file_paths += [os.path.join(folder, x) for x in os.listdir(folder)] 615 basenames = [os.path.basename(f) for f in file_paths] 616 ids = set() 617 for f in file_paths: 618 if f.endswith(data_suffix) and os.path.basename(f).startswith(data_prefix): 619 bn = os.path.basename(f) 620 video_id = strip_prefix(strip_suffix(bn, data_suffix), data_prefix) 621 if all([video_id + s in basenames for s in feature_suffix]): 622 ids.add(video_id) 623 ids = sorted(ids) 624 return ids 625 626 def get_bodyparts(self) -> List: 627 """ 628 Get a list of bodypart names 629 630 Parameters 631 ---------- 632 data_dict : dict 633 the data dictionary (passed to feature extractor) 634 clip_id : str 635 the clip id 636 637 Returns 638 ------- 639 bodyparts : list 640 a list of string or integer body part names 641 """ 642 643 return [x for x in self.bodyparts if x not in self.ignored_bodyparts] 644 645 def get_coords(self, data_dict: Dict, clip_id: str, bodypart: str) -> np.ndarray: 646 """ 647 Get the coordinates array of a specific bodypart in a specific clip 648 649 Parameters 650 ---------- 651 data_dict : dict 652 the data dictionary (passed to feature extractor) 653 clip_id : str 654 the clip id 655 bodypart : str 656 the name of the body part 657 658 Returns 659 ------- 660 coords : np.ndarray 661 the coordinates array of shape (#timesteps, #coordinates) 662 """ 663 664 columns = [x for x in data_dict[clip_id].columns if x != "likelihood"] 665 xy_coord = ( 666 data_dict[clip_id] 667 .xs(bodypart, axis=0, level=1, drop_level=False)[columns] 668 .values 669 ) 670 return xy_coord 671 672 def get_n_frames(self, data_dict: Dict, clip_id: str) -> int: 673 """ 674 Get the length of the clip 675 676 Parameters 677 ---------- 678 data_dict : dict 679 the data dictionary (passed to feature extractor) 680 clip_id : str 681 the clip id 682 683 Returns 684 ------- 685 n_frames : int 686 the length of the clip 687 """ 688 689 if clip_id in data_dict: 690 return len(data_dict[clip_id].groupby(level=0)) 691 else: 692 return min( 693 [len(data_dict[ind_k].groupby(level=0)) for ind_k in clip_id.split("+")] 694 ) 695 696 def _filter(self, data_dict: Dict) -> Tuple[Dict, Dict, Dict]: 697 """ 698 Apply filters to a data dictionary + normalize the values and generate frame index dictionaries 699 700 The filters include filling nan values, applying length and likelihood thresholds and removing 701 ignored clip ids. 702 """ 703 704 new_data_dict = {} 705 keys = list(data_dict.keys()) 706 for key in keys: 707 if key == "loaded": 708 continue 709 coord = data_dict.pop(key) 710 if key in self.ignored_clips: 711 continue 712 num_frames = len(coord.index.unique(level=0)) 713 if num_frames < self.frame_limit: 714 continue 715 if "likelihood" in coord.columns: 716 columns = list(coord.columns) 717 columns.remove("likelihood") 718 coord.loc[ 719 coord["likelihood"] < self.likelihood_threshold, columns 720 ] = np.nan 721 if not isinstance(self.centered, Iterable): 722 self.centered = [ 723 bool(self.centered) 724 for dim in ["x", "y", "z"] 725 if dim in coord.columns 726 ] 727 for i, dim in enumerate(["x", "y", "z"]): 728 if dim in coord.columns: 729 if self.centered[i]: 730 coord[dim] = coord[dim] + self.canvas_shape[i] // 2 731 # coord.loc[coord[dim] < -self.canvas_shape[i] * 3 // 2, dim] = np.nan 732 # coord.loc[coord[dim] > self.canvas_shape[i] * 3 // 2, dim] = np.nan 733 coord = coord.sort_index(level=0) 734 for bp in coord.index.unique(level=1): 735 coord.loc[coord.index.isin([bp], level=1)] = coord[ 736 coord.index.isin([bp], level=1) 737 ].interpolate() 738 dims = [x for x in coord.columns if x != "likelihood"] 739 mask = ~coord[dims[0]].isna() 740 for dim in dims[1:]: 741 mask = mask & (~coord[dim].isna()) 742 mean = coord.loc[mask].groupby(level=0).mean() 743 for frame in set(coord.index.get_level_values(0)): 744 if frame not in mean.index: 745 mean.loc[frame] = [np.nan for _ in mean.columns] 746 mean = mean.interpolate() 747 mean[mean.isna()] = 0 748 for dim in coord.columns: 749 if dim == "likelihood": 750 continue 751 coord.loc[coord[dim].isna(), dim] = mean.loc[ 752 coord.loc[coord[dim].isna()].index.get_level_values(0) 753 ][dim].to_numpy() 754 if np.sum(self.canvas_shape) > 0: 755 for i, dim in enumerate(["x", "y", "z"]): 756 if dim in coord.columns: 757 coord[dim] = ( 758 coord[dim] - self.canvas_shape[i] // 2 759 ) / self.canvas_shape[0] 760 new_data_dict[key] = coord 761 max_frames = {} 762 min_frames = {} 763 for key, value in new_data_dict.items(): 764 max_frames[key] = max(value.index.unique(0)) 765 min_frames[key] = min(value.index.unique(0)) 766 if "loaded" in data_dict: 767 new_data_dict["loaded"] = data_dict["loaded"] 768 return new_data_dict, min_frames, max_frames 769 770 def _get_files_from_ids(self): 771 files = defaultdict(lambda: []) 772 used_prefixes = defaultdict(lambda: []) 773 for f in self.file_paths: 774 if f.endswith(tuple([x for x in self.data_suffices])): 775 bn = os.path.basename(f) 776 video_id = strip_prefix( 777 strip_suffix(bn, self.data_suffices), self.data_prefixes 778 ) 779 ok = True 780 if self.data_prefixes is not None: 781 for p in self.data_prefixes: 782 if bn.startswith(p): 783 if p not in used_prefixes[video_id]: 784 used_prefixes[video_id].append(p) 785 else: 786 ok = False 787 break 788 if not ok: 789 continue 790 files[video_id].append(f) 791 files = [files[x] for x in self.video_order] 792 return files 793 794 def _make_trimmed_data(self, keypoint_dict: Dict) -> Tuple[List, Dict, List]: 795 """ 796 Cut a keypoint dictionary into overlapping pieces of equal length 797 """ 798 799 X = [] 800 original_coordinates = [] 801 lengths = defaultdict(lambda: {}) 802 if not os.path.exists(self.feature_save_path): 803 try: 804 os.mkdir(self.feature_save_path) 805 except FileExistsError: 806 pass 807 order = sorted(list(keypoint_dict.keys())) 808 for v_id in order: 809 keypoints = keypoint_dict[v_id] 810 v_len = min([len(x) for x in keypoints.values()]) 811 sp = np.arange(0, v_len, self.step) 812 pad = sp[-1] + self.len_segment - v_len 813 video_id, clip_id = v_id.split("---") 814 for key in keypoints: 815 if len(keypoints[key]) > v_len: 816 keypoints[key] = keypoints[key][:v_len] 817 if len(keypoints[key].shape) == 2: 818 keypoints[key] = np.pad(keypoints[key], ((0, pad), (0, 0))) 819 else: 820 keypoints[key] = np.pad( 821 keypoints[key], ((0, pad), (0, 0), (0, 0), (0, 0)) 822 ) 823 for i, start in enumerate(sp): 824 sample_dict = {} 825 original_coordinates.append((v_id, i)) 826 for key in keypoints: 827 sample_dict[key] = keypoints[key][start : start + self.len_segment] 828 sample_dict[key] = torch.tensor(np.array(sample_dict[key])).float() 829 sample_dict[key] = sample_dict[key].permute( 830 (*range(1, len(sample_dict[key].shape)), 0) 831 ) 832 name = os.path.join(self.feature_save_path, f"{v_id}_{start}.pickle") 833 X.append(name) 834 lengths[video_id][clip_id] = v_len 835 with open(name, "wb") as f: 836 pickle.dump(sample_dict, f) 837 return X, dict(lengths), original_coordinates 838 839 def _load_saved_features(self, video_id: str): 840 """ 841 Load saved features file `(#frames, #features)` 842 """ 843 844 basenames = [os.path.basename(x) for x in self.file_paths] 845 loaded_features_cat = [] 846 self.feature_suffix = sorted(self.feature_suffix) 847 for feature_suffix in self.feature_suffix: 848 i = basenames.index(os.path.basename(video_id) + feature_suffix) 849 path = self.file_paths[i] 850 if not os.path.exists(path): 851 raise RuntimeError(f"Did not find a feature file for {video_id}!") 852 extension = feature_suffix.split(".")[-1] 853 if extension in ["pickle", "pkl"]: 854 with open(path, "rb") as f: 855 loaded_features = pickle.load(f) 856 elif extension in ["pt", "pth"]: 857 loaded_features = torch.load(path) 858 elif extension == "npy": 859 loaded_features = np.load(path, allow_pickle=True).item() 860 else: 861 raise ValueError( 862 f"Found feature file in an unrecognized format: .{extension}. \n " 863 "Please save with torch (as .pt or .pth), numpy (as .npy) or pickle (as .pickle or .pkl)." 864 ) 865 loaded_features_cat.append(loaded_features) 866 keys = list(loaded_features_cat[0].keys()) 867 loaded_features = {} 868 for k in keys: 869 if k in ["min_frames", "max_frames", "video_tag"]: 870 loaded_features[k] = loaded_features_cat[0][k] 871 continue 872 features = [] 873 for x in loaded_features_cat: 874 if not isinstance(x[k], torch.Tensor): 875 features.append(torch.from_numpy(x[k])) 876 else: 877 features.append(x[k]) 878 a = torch.cat(features) 879 if self.transpose: 880 a = a.T 881 loaded_features[k] = a 882 return loaded_features 883 884 def get_likelihood( 885 self, data_dict: Dict, clip_id: str, bodypart: str 886 ) -> Union[np.ndarray, None]: 887 """ 888 Get the likelihood values 889 890 Parameters 891 ---------- 892 data_dict : dict 893 the data dictionary 894 clip_id : str 895 the clip id 896 bodypart : str 897 the name of the body part 898 899 Returns 900 ------- 901 likelihoods: np.ndarrray | None 902 `None` if the dataset doesn't have likelihoods or an array of shape (#timestamps) 903 """ 904 905 if "likelihood" in data_dict[clip_id].columns: 906 likelihood = ( 907 data_dict[clip_id] 908 .xs(bodypart, axis=0, level=1, drop_level=False) 909 .values[:, -1] 910 ) 911 return likelihood 912 else: 913 return None 914 915 def _get_video_metadata(self, metadata_list: Optional[List]): 916 """ 917 Make a single metadata dictionary from a list of dictionaries recieved from different data prefixes 918 """ 919 920 if metadata_list is None: 921 return None 922 else: 923 return metadata_list[0] 924 925 def get_indices(self, tag: int) -> List: 926 """ 927 Get a list of indices of samples that have a specific meta tag 928 929 Parameters 930 ---------- 931 tag : int 932 the meta tag for the subsample (`None` for the whole dataset) 933 934 Returns 935 ------- 936 indices : list 937 a list of indices that meet the criteria 938 """ 939 940 if tag is None: 941 return list(range(len(self.data))) 942 else: 943 return list(np.where(self.metadata == tag)[0]) 944 945 def get_tags(self) -> List: 946 """ 947 Get a list of all meta tags 948 949 Returns 950 ------- 951 tags: List 952 a list of unique meta tag values 953 """ 954 955 if self.metadata is None: 956 return [None] 957 else: 958 return list(np.unique(self.metadata)) 959 960 def get_tag(self, idx: int) -> Union[int, None]: 961 """ 962 Return a tag object corresponding to an index 963 964 Tags can carry meta information (like annotator id) and are accepted by models that require 965 that information. When a tag is `None`, it is not passed to the model. 966 967 Parameters 968 ---------- 969 idx : int 970 the index 971 972 Returns 973 ------- 974 tag : int 975 the tag object 976 """ 977 978 if self.metadata is None or idx is None: 979 return None 980 else: 981 return self.metadata[idx] 982 983 @abstractmethod 984 def _load_data(self) -> None: 985 """ 986 Load input data and generate data prompts 987 """
A generalized realization of a PoseInputStore
Assumes the following file structure:
data_path
├── video1DLC1000.pickle
├── video2DLC400.pickle
├── video1_features.pt
└── video2_features.pt
Here data_suffix is {'DLC1000.pickle', 'DLC400.pickle'} and feature_suffix (optional) is '_features.pt'.
GeneralInputStore( video_order: List = None, data_path: Union[Set, str] = None, file_paths: Set = None, data_suffix: Union[Set, str] = None, data_prefix: Union[Set, str] = None, feature_suffix: str = None, convert_int_indices: bool = True, feature_save_path: str = None, canvas_shape: List = None, len_segment: int = 128, overlap: int = 0, feature_extraction: str = 'kinematic', ignored_clips: List = None, ignored_bodyparts: List = None, default_agent_name: str = 'ind0', key_objects: Tuple = None, likelihood_threshold: float = 0, num_cpus: int = None, frame_limit: int = 1, normalize: bool = False, feature_extraction_pars: Dict = None, centered: bool = False, transpose_features: bool = False, *args, **kwargs)
46 def __init__( 47 self, 48 video_order: List = None, 49 data_path: Union[Set, str] = None, 50 file_paths: Set = None, 51 data_suffix: Union[Set, str] = None, 52 data_prefix: Union[Set, str] = None, 53 feature_suffix: str = None, 54 convert_int_indices: bool = True, 55 feature_save_path: str = None, 56 canvas_shape: List = None, 57 len_segment: int = 128, 58 overlap: int = 0, 59 feature_extraction: str = "kinematic", 60 ignored_clips: List = None, 61 ignored_bodyparts: List = None, 62 default_agent_name: str = "ind0", 63 key_objects: Tuple = None, 64 likelihood_threshold: float = 0, 65 num_cpus: int = None, 66 frame_limit: int = 1, 67 normalize: bool = False, 68 feature_extraction_pars: Dict = None, 69 centered: bool = False, 70 transpose_features: bool = False, 71 *args, 72 **kwargs, 73 ) -> None: 74 """ 75 Parameters 76 ---------- 77 video_order : list, optional 78 a list of video ids that should be processed in the same order (not passed if creating from key objects 79 data_path : str | set, optional 80 the path to the folder where the pose and feature files are stored or a set of such paths 81 (not passed if creating from key objects or from `file_paths`) 82 file_paths : set, optional 83 a set of string paths to the pose and feature files 84 (not passed if creating from key objects or from `data_path`) 85 data_suffix : str | set, optional 86 the suffix or the set of suffices such that the pose files are named {video_id}{data_suffix} 87 (not passed if creating from key objects or if irrelevant for the dataset) 88 data_prefix : str | set, optional 89 the prefix or the set of prefixes such that the pose files for different video views of the same 90 clip are named {prefix}{sep}{video_id}{data_suffix} (not passed if creating from key objects 91 or if irrelevant for the dataset) 92 feature_suffix : str | set, optional 93 the suffix or the set of suffices such that the additional feature files are named 94 {video_id}{feature_suffix} (and placed at the data_path folder) 95 convert_int_indices : bool, default True 96 if `True`, convert any integer key `i` in feature files to `'ind{i}'` 97 feature_save_path : str, optional 98 the path to the folder where pre-processed files are stored (not passed if creating from key objects) 99 canvas_shape : List, default [1, 1] 100 the canvas size where the pose is defined 101 len_segment : int, default 128 102 the length of the segments in which the data should be cut (in frames) 103 overlap : int, default 0 104 the length of the overlap between neighboring segments (in frames) 105 feature_extraction : str, default 'kinematic' 106 the feature extraction method (see options.feature_extractors for available options) 107 ignored_clips : list, optional 108 list of strings of clip ids to ignore 109 ignored_bodyparts : list, optional 110 list of strings of bodypart names to ignore 111 default_agent_name : str, default 'ind0' 112 the agent name used as default in the pose files for a single agent 113 key_objects : tuple, optional 114 a tuple of key objects 115 likelihood_threshold : float, default 0 116 coordinate values with likelihoods less than this value will be set to 'unknown' 117 num_cpus : int, optional 118 the number of cpus to use in data processing 119 frame_limit : int, default 1 120 clips shorter than this number of frames will be ignored 121 feature_extraction_pars : dict, optional 122 parameters of the feature extractor 123 """ 124 125 super().__init__() 126 self.loaded_max = 0 127 if feature_extraction_pars is None: 128 feature_extraction_pars = {} 129 if ignored_clips is None: 130 ignored_clips = [] 131 self.bodyparts = [] 132 self.visibility = None 133 self.normalize = normalize 134 135 if canvas_shape is None: 136 canvas_shape = [1, 1] 137 if isinstance(data_suffix, str): 138 data_suffix = [data_suffix] 139 if isinstance(data_prefix, str): 140 data_prefix = [data_prefix] 141 if isinstance(data_path, str): 142 data_path = [data_path] 143 if isinstance(feature_suffix, str): 144 feature_suffix = [feature_suffix] 145 146 self.video_order = video_order 147 self.centered = centered 148 self.feature_extraction = feature_extraction 149 self.len_segment = int(len_segment) 150 self.data_suffices = data_suffix 151 self.data_prefixes = data_prefix 152 self.feature_suffix = feature_suffix 153 self.convert_int_indices = convert_int_indices 154 if overlap < 1: 155 overlap = overlap * len_segment 156 self.overlap = int(overlap) 157 self.canvas_shape = canvas_shape 158 self.default_agent_name = default_agent_name 159 self.feature_save_path = feature_save_path 160 self.data_suffices = data_suffix 161 self.data_prefixes = data_prefix 162 self.likelihood_threshold = likelihood_threshold 163 self.num_cpus = num_cpus 164 self.frame_limit = frame_limit 165 self.transpose = transpose_features 166 167 self.ram = False 168 self.min_frames = {} 169 self.original_coordinates = np.array([]) 170 171 self.file_paths = self._get_file_paths(file_paths, data_path) 172 173 self.extractor = options.feature_extractors[self.feature_extraction]( 174 self, 175 **feature_extraction_pars, 176 ) 177 178 self.canvas_center = np.array(canvas_shape) // 2 179 180 if ignored_clips is not None: 181 self.ignored_clips = ignored_clips 182 else: 183 self.ignored_clips = [] 184 if ignored_bodyparts is not None: 185 self.ignored_bodyparts = ignored_bodyparts 186 else: 187 self.ignored_bodyparts = [] 188 189 self.step = self.len_segment - self.overlap 190 if self.step < 0: 191 raise ValueError( 192 f"The overlap value ({self.overlap}) cannot be larger than len_segment ({self.len_segment}" 193 ) 194 195 if self.feature_save_path is None and data_path is not None: 196 self.feature_save_path = os.path.join(data_path[0], "trimmed") 197 198 if key_objects is None and self.video_order is not None: 199 print("Computing input features...") 200 self.data = self._load_data() 201 elif key_objects is not None: 202 self.load_from_key_objects(key_objects)
Parameters
video_order : list, optional
a list of video ids that should be processed in the same order (not passed if creating from key objects
data_path : str | set, optional
the path to the folder where the pose and feature files are stored or a set of such paths
(not passed if creating from key objects or from file_paths)
file_paths : set, optional
a set of string paths to the pose and feature files
(not passed if creating from key objects or from data_path)
data_suffix : str | set, optional
the suffix or the set of suffices such that the pose files are named {video_id}{data_suffix}
(not passed if creating from key objects or if irrelevant for the dataset)
data_prefix : str | set, optional
the prefix or the set of prefixes such that the pose files for different video views of the same
clip are named {prefix}{sep}{video_id}{data_suffix} (not passed if creating from key objects
or if irrelevant for the dataset)
feature_suffix : str | set, optional
the suffix or the set of suffices such that the additional feature files are named
{video_id}{feature_suffix} (and placed at the data_path folder)
convert_int_indices : bool, default True
if True, convert any integer key i in feature files to 'ind{i}'
feature_save_path : str, optional
the path to the folder where pre-processed files are stored (not passed if creating from key objects)
canvas_shape : List, default [1, 1]
the canvas size where the pose is defined
len_segment : int, default 128
the length of the segments in which the data should be cut (in frames)
overlap : int, default 0
the length of the overlap between neighboring segments (in frames)
feature_extraction : str, default 'kinematic'
the feature extraction method (see options.feature_extractors for available options)
ignored_clips : list, optional
list of strings of clip ids to ignore
ignored_bodyparts : list, optional
list of strings of bodypart names to ignore
default_agent_name : str, default 'ind0'
the agent name used as default in the pose files for a single agent
key_objects : tuple, optional
a tuple of key objects
likelihood_threshold : float, default 0
coordinate values with likelihoods less than this value will be set to 'unknown'
num_cpus : int, optional
the number of cpus to use in data processing
frame_limit : int, default 1
clips shorter than this number of frames will be ignored
feature_extraction_pars : dict, optional
parameters of the feature extractor
def
get_folder(self, video_id: str) -> str:
244 def get_folder(self, video_id: str) -> str: 245 """ 246 Get the input folder that the file with this video id was read from 247 248 Parameters 249 ---------- 250 video_id : str 251 the video id 252 253 Returns 254 ------- 255 folder : str 256 the path to the directory that contains the input file associated with the video id 257 """ 258 259 for file in self.file_paths: 260 if ( 261 strip_prefix( 262 strip_suffix(os.path.basename(file), self.data_suffices), 263 self.data_prefixes, 264 ) 265 == video_id 266 ): 267 return os.path.dirname(file)
Get the input folder that the file with this video id was read from
Parameters
video_id : str the video id
Returns
folder : str the path to the directory that contains the input file associated with the video id
def
remove(self, indices: List) -> None:
269 def remove(self, indices: List) -> None: 270 """ 271 Remove the samples corresponding to indices 272 273 Parameters 274 ---------- 275 indices : int 276 a list of integer indices to remove 277 """ 278 279 if len(indices) > 0: 280 mask = np.ones(len(self.original_coordinates)) 281 mask[indices] = 0 282 mask = mask.astype(bool) 283 for file in self.data[~mask]: 284 os.remove(file) 285 self.original_coordinates = self.original_coordinates[mask] 286 self.data = self.data[mask] 287 if self.metadata is not None: 288 self.metadata = self.metadata[mask]
Remove the samples corresponding to indices
Parameters
indices : int a list of integer indices to remove
def
key_objects(self) -> Tuple:
290 def key_objects(self) -> Tuple: 291 """ 292 Return a tuple of the key objects necessary to re-create the Store 293 294 Returns 295 ------- 296 key_objects : tuple 297 a tuple of key objects 298 """ 299 300 for k, v in self.min_frames.items(): 301 self.min_frames[k] = dict(v) 302 for k, v in self.max_frames.items(): 303 self.max_frames[k] = dict(v) 304 return ( 305 self.original_coordinates, 306 dict(self.min_frames), 307 dict(self.max_frames), 308 self.data, 309 self.visibility, 310 self.step, 311 self.file_paths, 312 self.len_segment, 313 self.metadata, 314 )
Return a tuple of the key objects necessary to re-create the Store
Returns
key_objects : tuple a tuple of key objects
def
load_from_key_objects(self, key_objects: Tuple) -> None:
316 def load_from_key_objects(self, key_objects: Tuple) -> None: 317 """ 318 Load the information from a tuple of key objects 319 320 Parameters 321 ---------- 322 key_objects : tuple 323 a tuple of key objects 324 """ 325 326 ( 327 self.original_coordinates, 328 self.min_frames, 329 self.max_frames, 330 self.data, 331 self.visibility, 332 self.step, 333 self.file_paths, 334 self.len_segment, 335 self.metadata, 336 ) = key_objects
Load the information from a tuple of key objects
Parameters
key_objects : tuple a tuple of key objects
def
to_ram(self) -> None:
338 def to_ram(self) -> None: 339 """ 340 Transfer the data samples to RAM if they were previously stored as file paths 341 """ 342 343 if self.ram: 344 return 345 346 data = p_map(lambda x: self[x], list(range(len(self))), num_cpus=self.num_cpus) 347 # data = [load(x) for x in self.data] 348 self.data = TensorDict(data) 349 self.ram = True
Transfer the data samples to RAM if they were previously stored as file paths
def
get_original_coordinates(self) -> numpy.ndarray:
351 def get_original_coordinates(self) -> np.ndarray: 352 """ 353 Return the original coordinates array 354 355 Returns 356 ------- 357 np.ndarray 358 an array that contains the coordinates of the data samples in original input data (video id, clip id, 359 start frame) 360 """ 361 362 return self.original_coordinates
Return the original coordinates array
Returns
np.ndarray an array that contains the coordinates of the data samples in original input data (video id, clip id, start frame)
def
create_subsample(self, indices: List, ssl_indices: List = None)
364 def create_subsample(self, indices: List, ssl_indices: List = None): 365 """ 366 Create a new store that contains a subsample of the data 367 368 Parameters 369 ---------- 370 indices : list 371 the indices to be included in the subsample 372 ssl_indices : list, optional 373 the indices to be included in the subsample without the annotation data 374 """ 375 376 if ssl_indices is None: 377 ssl_indices = [] 378 new = self.new() 379 new.original_coordinates = self.original_coordinates[indices + ssl_indices] 380 new.min_frames = self.min_frames 381 new.max_frames = self.max_frames 382 new.data = self.data[indices + ssl_indices] 383 new.visibility = self.visibility 384 new.step = self.step 385 new.file_paths = self.file_paths 386 new.len_segment = self.len_segment 387 if self.metadata is None: 388 new.metadata = None 389 else: 390 new.metadata = self.metadata[indices + ssl_indices] 391 return new
Create a new store that contains a subsample of the data
Parameters
indices : list the indices to be included in the subsample ssl_indices : list, optional the indices to be included in the subsample without the annotation data
def
get_video_id(self, coords: Tuple) -> str:
393 def get_video_id(self, coords: Tuple) -> str: 394 """ 395 Get the video id from an element of original coordinates 396 397 Parameters 398 ---------- 399 coords : tuple 400 an element of the original coordinates array 401 402 Returns 403 ------- 404 video_id: str 405 the id of the video that the coordinates point to 406 """ 407 408 video_name = coords[0].split("---")[0] 409 return video_name
Get the video id from an element of original coordinates
Parameters
coords : tuple an element of the original coordinates array
Returns
video_id: str the id of the video that the coordinates point to
def
get_clip_id(self, coords: Tuple) -> str:
411 def get_clip_id(self, coords: Tuple) -> str: 412 """ 413 Get the clip id from an element of original coordinates 414 415 Parameters 416 ---------- 417 coords : tuple 418 an element of the original coordinates array 419 420 Returns 421 ------- 422 clip_id : str 423 the id of the clip that the coordinates point to 424 """ 425 426 clip_id = coords[0].split("---")[1] 427 return clip_id
Get the clip id from an element of original coordinates
Parameters
coords : tuple an element of the original coordinates array
Returns
clip_id : str the id of the clip that the coordinates point to
def
get_clip_length(self, video_id: str, clip_id: str) -> int:
429 def get_clip_length(self, video_id: str, clip_id: str) -> int: 430 """ 431 Get the clip length from the id 432 433 Parameters 434 ---------- 435 video_id : str 436 the video id 437 clip_id : str 438 the clip id 439 440 Returns 441 ------- 442 clip_length : int 443 the length of the clip 444 """ 445 446 inds = clip_id.split("+") 447 max_frame = min([self.max_frames[video_id][x] for x in inds]) 448 min_frame = max([self.min_frames[video_id][x] for x in inds]) 449 return max_frame - min_frame + 1
Get the clip length from the id
Parameters
video_id : str the video id clip_id : str the clip id
Returns
clip_length : int the length of the clip
def
get_clip_start_end(self, coords: Tuple) -> Tuple[int, int]:
451 def get_clip_start_end(self, coords: Tuple) -> Tuple[int, int]: 452 """ 453 Get the clip start and end frames from an element of original coordinates 454 455 Parameters 456 ---------- 457 coords : tuple 458 an element of original coordinates array 459 460 Returns 461 ------- 462 start : int 463 the start frame of the clip that the coordinates point to 464 end : int 465 the end frame of the clip that the coordinates point to 466 """ 467 468 l = self.get_clip_length_from_coords(coords) 469 i = coords[1] 470 start = int(i) * self.step 471 end = min(start + self.len_segment, l) 472 return start, end
Get the clip start and end frames from an element of original coordinates
Parameters
coords : tuple an element of original coordinates array
Returns
start : int the start frame of the clip that the coordinates point to end : int the end frame of the clip that the coordinates point to
def
get_clip_start(self, video_name: str, clip_id: str) -> int:
474 def get_clip_start(self, video_name: str, clip_id: str) -> int: 475 """ 476 Get the clip start frame from the video id and the clip id 477 478 Parameters 479 ---------- 480 video_name : str 481 the video id 482 clip_id : str 483 the clip id 484 485 Returns 486 ------- 487 clip_start : int 488 the start frame of the clip 489 """ 490 491 return max( 492 [self.min_frames[video_name][clip_id_k] for clip_id_k in clip_id.split("+")] 493 )
Get the clip start frame from the video id and the clip id
Parameters
video_name : str the video id clip_id : str the clip id
Returns
clip_start : int the start frame of the clip
def
get_visibility( self, video_id: str, clip_id: str, start: int, end: int, score: int) -> float:
495 def get_visibility( 496 self, video_id: str, clip_id: str, start: int, end: int, score: int 497 ) -> float: 498 """ 499 Get the fraction of the frames in that have a visibility score better than a hard_threshold 500 501 For example, in the case of keypoint data the visibility score can be the number of identified keypoints. 502 503 Parameters 504 ---------- 505 video_id : str 506 the video id of the frames 507 clip_id : str 508 the clip id of the frames 509 start : int 510 the start frame 511 end : int 512 the end frame 513 score : float 514 the visibility score hard_threshold 515 516 Returns 517 ------- 518 frac_visible: float 519 the fraction of frames with visibility above the hard_threshold 520 """ 521 522 s = 0 523 for ind_k in clip_id.split("+"): 524 s += np.sum(self.visibility[video_id][ind_k][start:end] > score) / ( 525 end - start 526 ) 527 return s / len(clip_id.split("+"))
Get the fraction of the frames in that have a visibility score better than a hard_threshold
For example, in the case of keypoint data the visibility score can be the number of identified keypoints.
Parameters
video_id : str the video id of the frames clip_id : str the clip id of the frames start : int the start frame end : int the end frame score : float the visibility score hard_threshold
Returns
frac_visible: float the fraction of frames with visibility above the hard_threshold
def
get_annotation_objects(self) -> Dict:
529 def get_annotation_objects(self) -> Dict: 530 """ 531 Get a dictionary of objects necessary to create an AnnotationStore 532 533 Returns 534 ------- 535 annotation_objects : dict 536 a dictionary of objects to be passed to the AnnotationStore constructor where the keys are the names of 537 the objects 538 """ 539 540 min_frames = self.min_frames 541 max_frames = self.max_frames 542 num_bp = self.visibility 543 return { 544 "min_frames": min_frames, 545 "max_frames": max_frames, 546 "visibility": num_bp, 547 }
Get a dictionary of objects necessary to create an AnnotationStore
Returns
annotation_objects : dict a dictionary of objects to be passed to the AnnotationStore constructor where the keys are the names of the objects
@classmethod
def
get_file_ids( cls, data_suffix: Union[Set, str] = None, data_path: Union[Set, str] = None, data_prefix: Union[Set, str] = None, file_paths: Set = None, feature_suffix: Set = None, *args, **kwargs) -> List:
549 @classmethod 550 def get_file_ids( 551 cls, 552 data_suffix: Union[Set, str] = None, 553 data_path: Union[Set, str] = None, 554 data_prefix: Union[Set, str] = None, 555 file_paths: Set = None, 556 feature_suffix: Set = None, 557 *args, 558 **kwargs, 559 ) -> List: 560 """ 561 Process data parameters and return a list of ids of the videos that should 562 be processed by the __init__ function 563 564 Parameters 565 ---------- 566 data_suffix : set | str, optional 567 the suffix (or a set of suffixes) of the input data files 568 data_path : set | str, optional 569 the path to the folder where the pose and feature files are stored or a set of such paths 570 (not passed if creating from key objects or from `file_paths`) 571 data_prefix : set | str, optional 572 the prefix or the set of prefixes such that the pose files for different video views of the same 573 clip are named {prefix}{sep}{video_id}{data_suffix} (not passed if creating from key objects 574 or if irrelevant for the dataset) 575 file_paths : set, optional 576 a set of string paths to the pose and feature files 577 feature_suffix : str | set, optional 578 the suffix or the set of suffices such that the additional feature files are named 579 {video_id}{feature_suffix} (and placed at the `data_path` folder or at `file_paths`) 580 581 Returns 582 ------- 583 video_ids : list 584 a list of video file ids 585 """ 586 587 if data_suffix is None: 588 if cls.data_suffix is not None: 589 data_suffix = cls.data_suffix 590 else: 591 raise ValueError("Cannot get video ids without the data suffix!") 592 if feature_suffix is None: 593 feature_suffix = [] 594 if data_prefix is None: 595 data_prefix = "" 596 if isinstance(data_suffix, str): 597 data_suffix = [data_suffix] 598 else: 599 data_suffix = [x for x in data_suffix] 600 data_suffix = tuple(data_suffix) 601 if isinstance(data_prefix, str): 602 data_prefix = data_prefix 603 else: 604 data_prefix = tuple([x for x in data_prefix]) 605 if isinstance(feature_suffix, str): 606 feature_suffix = [feature_suffix] 607 if file_paths is None: 608 file_paths = [] 609 if data_path is not None: 610 if isinstance(data_path, str): 611 data_path = [data_path] 612 file_paths = [] 613 for folder in data_path: 614 file_paths += [os.path.join(folder, x) for x in os.listdir(folder)] 615 basenames = [os.path.basename(f) for f in file_paths] 616 ids = set() 617 for f in file_paths: 618 if f.endswith(data_suffix) and os.path.basename(f).startswith(data_prefix): 619 bn = os.path.basename(f) 620 video_id = strip_prefix(strip_suffix(bn, data_suffix), data_prefix) 621 if all([video_id + s in basenames for s in feature_suffix]): 622 ids.add(video_id) 623 ids = sorted(ids) 624 return ids
Process data parameters and return a list of ids of the videos that should be processed by the __init__ function
Parameters
data_suffix : set | str, optional
the suffix (or a set of suffixes) of the input data files
data_path : set | str, optional
the path to the folder where the pose and feature files are stored or a set of such paths
(not passed if creating from key objects or from file_paths)
data_prefix : set | str, optional
the prefix or the set of prefixes such that the pose files for different video views of the same
clip are named {prefix}{sep}{video_id}{data_suffix} (not passed if creating from key objects
or if irrelevant for the dataset)
file_paths : set, optional
a set of string paths to the pose and feature files
feature_suffix : str | set, optional
the suffix or the set of suffices such that the additional feature files are named
{video_id}{feature_suffix} (and placed at the data_path folder or at file_paths)
Returns
video_ids : list a list of video file ids
def
get_bodyparts(self) -> List:
626 def get_bodyparts(self) -> List: 627 """ 628 Get a list of bodypart names 629 630 Parameters 631 ---------- 632 data_dict : dict 633 the data dictionary (passed to feature extractor) 634 clip_id : str 635 the clip id 636 637 Returns 638 ------- 639 bodyparts : list 640 a list of string or integer body part names 641 """ 642 643 return [x for x in self.bodyparts if x not in self.ignored_bodyparts]
Get a list of bodypart names
Parameters
data_dict : dict the data dictionary (passed to feature extractor) clip_id : str the clip id
Returns
bodyparts : list a list of string or integer body part names
def
get_coords(self, data_dict: Dict, clip_id: str, bodypart: str) -> numpy.ndarray:
645 def get_coords(self, data_dict: Dict, clip_id: str, bodypart: str) -> np.ndarray: 646 """ 647 Get the coordinates array of a specific bodypart in a specific clip 648 649 Parameters 650 ---------- 651 data_dict : dict 652 the data dictionary (passed to feature extractor) 653 clip_id : str 654 the clip id 655 bodypart : str 656 the name of the body part 657 658 Returns 659 ------- 660 coords : np.ndarray 661 the coordinates array of shape (#timesteps, #coordinates) 662 """ 663 664 columns = [x for x in data_dict[clip_id].columns if x != "likelihood"] 665 xy_coord = ( 666 data_dict[clip_id] 667 .xs(bodypart, axis=0, level=1, drop_level=False)[columns] 668 .values 669 ) 670 return xy_coord
Get the coordinates array of a specific bodypart in a specific clip
Parameters
data_dict : dict the data dictionary (passed to feature extractor) clip_id : str the clip id bodypart : str the name of the body part
Returns
coords : np.ndarray the coordinates array of shape (#timesteps, #coordinates)
def
get_n_frames(self, data_dict: Dict, clip_id: str) -> int:
672 def get_n_frames(self, data_dict: Dict, clip_id: str) -> int: 673 """ 674 Get the length of the clip 675 676 Parameters 677 ---------- 678 data_dict : dict 679 the data dictionary (passed to feature extractor) 680 clip_id : str 681 the clip id 682 683 Returns 684 ------- 685 n_frames : int 686 the length of the clip 687 """ 688 689 if clip_id in data_dict: 690 return len(data_dict[clip_id].groupby(level=0)) 691 else: 692 return min( 693 [len(data_dict[ind_k].groupby(level=0)) for ind_k in clip_id.split("+")] 694 )
Get the length of the clip
Parameters
data_dict : dict the data dictionary (passed to feature extractor) clip_id : str the clip id
Returns
n_frames : int the length of the clip
def
get_likelihood( self, data_dict: Dict, clip_id: str, bodypart: str) -> Optional[numpy.ndarray]:
884 def get_likelihood( 885 self, data_dict: Dict, clip_id: str, bodypart: str 886 ) -> Union[np.ndarray, None]: 887 """ 888 Get the likelihood values 889 890 Parameters 891 ---------- 892 data_dict : dict 893 the data dictionary 894 clip_id : str 895 the clip id 896 bodypart : str 897 the name of the body part 898 899 Returns 900 ------- 901 likelihoods: np.ndarrray | None 902 `None` if the dataset doesn't have likelihoods or an array of shape (#timestamps) 903 """ 904 905 if "likelihood" in data_dict[clip_id].columns: 906 likelihood = ( 907 data_dict[clip_id] 908 .xs(bodypart, axis=0, level=1, drop_level=False) 909 .values[:, -1] 910 ) 911 return likelihood 912 else: 913 return None
Get the likelihood values
Parameters
data_dict : dict the data dictionary clip_id : str the clip id bodypart : str the name of the body part
Returns
likelihoods: np.ndarrray | None
None if the dataset doesn't have likelihoods or an array of shape (#timestamps)
def
get_indices(self, tag: int) -> List:
925 def get_indices(self, tag: int) -> List: 926 """ 927 Get a list of indices of samples that have a specific meta tag 928 929 Parameters 930 ---------- 931 tag : int 932 the meta tag for the subsample (`None` for the whole dataset) 933 934 Returns 935 ------- 936 indices : list 937 a list of indices that meet the criteria 938 """ 939 940 if tag is None: 941 return list(range(len(self.data))) 942 else: 943 return list(np.where(self.metadata == tag)[0])
Get a list of indices of samples that have a specific meta tag
Parameters
tag : int
the meta tag for the subsample (None for the whole dataset)
Returns
indices : list a list of indices that meet the criteria
def
get_tag(self, idx: int) -> Optional[int]:
960 def get_tag(self, idx: int) -> Union[int, None]: 961 """ 962 Return a tag object corresponding to an index 963 964 Tags can carry meta information (like annotator id) and are accepted by models that require 965 that information. When a tag is `None`, it is not passed to the model. 966 967 Parameters 968 ---------- 969 idx : int 970 the index 971 972 Returns 973 ------- 974 tag : int 975 the tag object 976 """ 977 978 if self.metadata is None or idx is None: 979 return None 980 else: 981 return self.metadata[idx]
Return a tag object corresponding to an index
Tags can carry meta information (like annotator id) and are accepted by models that require
that information. When a tag is None, it is not passed to the model.
Parameters
idx : int the index
Returns
tag : int the tag object
990class FileInputStore(GeneralInputStore): 991 """ 992 An implementation of `dlc2action.data.InputStore` for datasets where each input data file corresponds to one video 993 """ 994 995 def _count_bodyparts( 996 self, data: Dict, stripped_name: str, max_frames: Dict 997 ) -> Dict: 998 """ 999 Create a visibility score dictionary (with a score from 0 to 1 assigned to each frame of each clip) 1000 """ 1001 1002 result = {stripped_name: {}} 1003 prefixes = list(data.keys()) 1004 for ind in data[prefixes[0]]: 1005 res = 0 1006 for _, data_dict in data.items(): 1007 num_bp = len(data_dict[ind].index.unique(level=1)) 1008 coords = ( 1009 data_dict[ind].values.reshape( 1010 -1, num_bp, len(data_dict[ind].columns) 1011 )[: max_frames[ind], :, 0] 1012 != 0 1013 ) 1014 res = np.sum(coords, axis=1) + res 1015 result[stripped_name][ind] = (res / len(prefixes)) / coords.shape[1] 1016 return result 1017 1018 def _generate_features(self, data: Dict, video_id: str) -> Dict: 1019 """ 1020 Generate features from the raw coordinates 1021 """ 1022 1023 features = defaultdict(lambda: {}) 1024 loaded_common = [] 1025 for prefix, data_dict in data.items(): 1026 if prefix == "": 1027 prefix = None 1028 if "loaded" in data_dict: 1029 loaded_common.append(torch.tensor(data_dict.pop("loaded"))) 1030 1031 key_features = self.extractor.extract_features( 1032 data_dict, video_id, prefix=prefix 1033 ) 1034 for f_key in key_features: 1035 features[f_key].update(key_features[f_key]) 1036 if len(loaded_common) > 0: 1037 loaded_common = torch.cat(loaded_common, dim=1) 1038 else: 1039 loaded_common = None 1040 if self.feature_suffix is not None: 1041 loaded_features = self._load_saved_features(video_id) 1042 for clip_id, feature_tensor in loaded_features.items(): 1043 if not isinstance(feature_tensor, torch.Tensor): 1044 feature_tensor = torch.tensor(feature_tensor) 1045 if self.convert_int_indices and ( 1046 isinstance(clip_id, int) or isinstance(clip_id, np.integer) 1047 ): 1048 clip_id = f"ind{clip_id}" 1049 key1 = f"{os.path.basename(video_id)}---{clip_id}" 1050 if key1 in features: 1051 try: 1052 key2 = list(features[key1].keys())[0] 1053 n_frames = features[key1][key2].shape[0] 1054 if feature_tensor.shape[0] != n_frames: 1055 n = feature_tensor.shape[0] - n_frames 1056 if ( 1057 abs(n) > 2 1058 and abs(feature_tensor.shape[1] - n_frames) <= 2 1059 ): 1060 feature_tensor = feature_tensor.T 1061 # If off by <=2 frames, just clip the end 1062 elif n > 0 and n <= 2: 1063 feature_tensor = feature_tensor[:n_frames, :] 1064 elif n < 0 and n >= -2: 1065 filler = feature_tensor[-2:-1, :] 1066 for i in range(n_frames - feature_tensor.shape[0]): 1067 feature_tensor = torch.cat( 1068 [feature_tensor, filler], 0 1069 ) 1070 else: 1071 raise RuntimeError( 1072 f"Number of frames in precomputed features with shape" 1073 f" {feature_tensor.shape} is inconsistent with generated features!" 1074 ) 1075 if loaded_common is not None: 1076 if feature_tensor.shape[0] == loaded_common.shape[0]: 1077 feature_tensor = torch.cat( 1078 [feature_tensor, loaded_common], dim=1 1079 ) 1080 elif feature_tensor.shape[0] == loaded_common.shape[1]: 1081 feature_tensor = torch.cat( 1082 [feature_tensor.T, loaded_common], dim=1 1083 ) 1084 else: 1085 raise ValueError( 1086 "The features from the data file and from the feature file have a different number of frames!" 1087 ) 1088 features[key1]["loaded"] = feature_tensor 1089 except ValueError: 1090 raise RuntimeError( 1091 "Individuals in precomputed features are inconsistent " 1092 "with generated features" 1093 ) 1094 elif loaded_common is not None: 1095 for key in features: 1096 features[key]["loaded"] = loaded_common 1097 return features 1098 1099 def _load_data(self) -> np.array: 1100 """ 1101 Load input data and generate data prompts 1102 """ 1103 1104 if self.video_order is None: 1105 return None 1106 1107 files = defaultdict(lambda: []) 1108 for f in self.file_paths: 1109 if f.endswith(tuple([x for x in self.data_suffices])): 1110 bn = os.path.basename(f) 1111 video_id = strip_prefix( 1112 strip_suffix(bn, self.data_suffices), self.data_prefixes 1113 ) 1114 files[video_id].append(f) 1115 files = [files[x] for x in self.video_order] 1116 1117 def make_data_dictionary(filenames): 1118 data = {} 1119 stored_maxes = defaultdict(lambda: []) 1120 min_frames, max_frames = {}, {} 1121 name = strip_suffix(filenames[0], self.data_suffices) 1122 name = os.path.basename(name) 1123 stripped_name = strip_prefix(name, self.data_prefixes) 1124 metadata_list = [] 1125 for filename in filenames: 1126 name = strip_suffix(filename, self.data_suffices) 1127 name = os.path.basename(name) 1128 prefix = strip_suffix(name, [stripped_name]) 1129 data_new, tag = self._open_data(filename, self.default_agent_name) 1130 data_new, min_frames, max_frames = self._filter(data_new) 1131 data[prefix] = data_new 1132 for key, val in max_frames.items(): 1133 stored_maxes[key].append(val) 1134 metadata_list.append(tag) 1135 video_tag = self._get_video_metadata(metadata_list) 1136 sample_df = list(list(data.values())[0].values())[0] 1137 self.bodyparts = sorted(list(sample_df.index.unique(1))) 1138 smallest_maxes = dict.fromkeys(stored_maxes) 1139 for key, val in stored_maxes.items(): 1140 smallest_maxes[key] = np.amin(val) 1141 data_dict = self._generate_features(data, stripped_name) 1142 bp_dict = self._count_bodyparts( 1143 data=data, stripped_name=stripped_name, max_frames=smallest_maxes 1144 ) 1145 min_frames = {stripped_name: min_frames} # name is e.g. 20190707T1126-1226 1146 max_frames = {stripped_name: max_frames} 1147 names, lengths, coords = self._make_trimmed_data(data_dict) 1148 return names, lengths, coords, bp_dict, min_frames, max_frames, video_tag 1149 1150 dict_list = p_map(make_data_dictionary, files, num_cpus=self.num_cpus) 1151 # dict_list = tqdm([make_data_dictionary(f) for f in files]) 1152 1153 self.visibility = {} 1154 self.min_frames = {} 1155 self.max_frames = {} 1156 self.original_coordinates = [] 1157 self.metadata = [] 1158 X = [] 1159 for ( 1160 names, 1161 lengths, 1162 coords, 1163 bp_dictionary, 1164 min_frames, 1165 max_frames, 1166 metadata, 1167 ) in dict_list: 1168 X += names 1169 self.original_coordinates += coords 1170 self.visibility.update(bp_dictionary) 1171 self.min_frames.update(min_frames) 1172 self.max_frames.update(max_frames) 1173 if metadata is not None: 1174 self.metadata += metadata 1175 del dict_list 1176 if len(self.metadata) != len(self.original_coordinates): 1177 self.metadata = None 1178 else: 1179 self.metadata = np.array(self.metadata) 1180 1181 self.min_frames = dict(self.min_frames) 1182 self.max_frames = dict(self.max_frames) 1183 self.original_coordinates = np.array(self.original_coordinates) 1184 return np.array(X) 1185 1186 @abstractmethod 1187 def _open_data( 1188 self, filename: str, default_clip_name: str 1189 ) -> Tuple[Dict, Optional[Dict]]: 1190 """ 1191 Load the keypoints from filename and organize them in a dictionary 1192 1193 In `data_dictionary`, the keys are clip ids and the values are `pandas` dataframes with two-level indices. 1194 The first level is the frame numbers and the second is the body part names. The dataframes should have from 1195 two to four columns labeled `"x"`, `"y"` and (optionally) `"z"` and `"likelihood"`. Each frame should have 1196 information on all the body parts. You don't have to filter the data in any way or fill the nans, it will 1197 be done automatically. 1198 1199 Parameters 1200 ---------- 1201 filename : str 1202 path to the pose file 1203 default_clip_name : str 1204 the name to assign to a clip if it does not have a name in the raw data 1205 1206 Returns 1207 ------- 1208 data dictionary : dict 1209 a dictionary where the keys are clip ids and the values are keypoint dataframes (see above for details) 1210 metadata_dictionary : dict 1211 a dictionary where the keys are clip ids and the values are metadata objects (can be any additional information, 1212 like the annotator tag; for no metadata pass `None`) 1213 """
An implementation of dlc2action.data.InputStore for datasets where each input data file corresponds to one video
Inherited Members
- GeneralInputStore
- GeneralInputStore
- data_suffix
- get_folder
- remove
- key_objects
- load_from_key_objects
- to_ram
- get_original_coordinates
- create_subsample
- get_video_id
- get_clip_id
- get_clip_length
- get_clip_start_end
- get_clip_start
- get_visibility
- get_annotation_objects
- get_file_ids
- get_bodyparts
- get_coords
- get_n_frames
- get_likelihood
- get_indices
- get_tag
1216class SequenceInputStore(GeneralInputStore): 1217 """ 1218 An implementation of `dlc2action.data.InputStore` for datasets where input data files correspond to multiple videos 1219 """ 1220 1221 def _count_bodyparts( 1222 self, data: Dict, stripped_name: str, max_frames: Dict 1223 ) -> Dict: 1224 """ 1225 Create a visibility score dictionary (with a score from 0 to 1 assigned to each frame of each clip) 1226 """ 1227 1228 result = {stripped_name: {}} 1229 for ind in data.keys(): 1230 num_bp = len(data[ind].index.unique(level=1)) 1231 coords = ( 1232 data[ind].values.reshape(-1, num_bp, len(data[ind].columns))[ 1233 : max_frames[ind], :, 0 1234 ] 1235 != 0 1236 ) 1237 res = np.sum(coords, axis=1) 1238 result[stripped_name][ind] = res / coords.shape[1] 1239 return result 1240 1241 def _generate_features(self, data: Dict, name: str) -> Dict: 1242 """ 1243 Generate features for an individual 1244 """ 1245 1246 features = self.extractor.extract_features(data, name, prefix=None) 1247 if self.feature_suffix is not None: 1248 loaded_features = self._load_saved_features(name) 1249 for clip_id, feature_tensor in loaded_features.items(): 1250 if not isinstance(feature_tensor, torch.Tensor): 1251 feature_tensor = torch.tensor(feature_tensor) 1252 if self.convert_int_indices and ( 1253 isinstance(clip_id, int) or isinstance(clip_id, np.integer) 1254 ): 1255 clip_id = f"ind{clip_id}" 1256 key1 = f"{os.path.basename(name)}---{clip_id}" 1257 if key1 in features: 1258 try: 1259 key2 = list(features[key1].keys())[0] 1260 n_frames = features[key1][key2].shape[0] 1261 if feature_tensor.shape[0] != n_frames: 1262 n = feature_tensor.shape[0] - n_frames 1263 if ( 1264 abs(n) > 2 1265 and abs(feature_tensor.shape[1] - n_frames) <= 2 1266 ): 1267 feature_tensor = feature_tensor.T 1268 # If off by <=2 frames, just clip the end 1269 elif n > 0 and n <= 2: 1270 feature_tensor = feature_tensor[:n_frames, :] 1271 elif n < 0 and n >= -2: 1272 filler = feature_tensor[-2:-1, :] 1273 for i in range(n_frames - feature_tensor.shape[0]): 1274 feature_tensor = torch.cat( 1275 [feature_tensor, filler], 0 1276 ) 1277 else: 1278 raise RuntimeError( 1279 print( 1280 f"Number of frames in precomputed features with shape" 1281 f" {feature_tensor.shape} is inconsistent with generated features!" 1282 ) 1283 ) 1284 features[key1]["loaded"] = feature_tensor 1285 except ValueError: 1286 raise RuntimeError( 1287 print( 1288 "Individuals in precomputed features are inconsistent " 1289 "with generated features" 1290 ) 1291 ) 1292 return features 1293 1294 def _load_data(self) -> np.array: 1295 """ 1296 Load input data and generate data prompts 1297 """ 1298 1299 if self.video_order is None: 1300 return None 1301 1302 files = [] 1303 for f in self.file_paths: 1304 if os.path.basename(f) in self.video_order: 1305 files.append(f) 1306 1307 def make_data_dictionary(seq_tuple): 1308 seq_id, sequence = seq_tuple 1309 data, tag = self._get_data(seq_id, sequence, self.default_agent_name) 1310 if "loaded" in data.keys(): 1311 loaded_features = data.pop("loaded") 1312 data, min_frames, max_frames = self._filter(data) 1313 sample_df = list(data.values())[0] 1314 self.bodyparts = sorted(list(sample_df.index.unique(1))) 1315 data_dict = self._generate_features(data, seq_id) 1316 for key in data_dict.keys(): 1317 data_dict[key]["loaded"] = loaded_features 1318 bp_dict = self._count_bodyparts( 1319 data=data, stripped_name=seq_id, max_frames=max_frames 1320 ) 1321 min_frames = {seq_id: min_frames} # name is e.g. 20190707T1126-1226 1322 max_frames = {seq_id: max_frames} 1323 names, lengths, coords = self._make_trimmed_data(data_dict) 1324 return names, lengths, coords, bp_dict, min_frames, max_frames, tag 1325 1326 seq_tuples = [] 1327 for file in files: 1328 opened = self._open_file(file) 1329 seq_tuples += opened 1330 dict_list = p_map( 1331 make_data_dictionary, sorted(seq_tuples), num_cpus=self.num_cpus 1332 ) 1333 # dict_list = tqdm([make_data_dictionary(f) for f in files]) 1334 1335 self.visibility = {} 1336 self.min_frames = {} 1337 self.max_frames = {} 1338 self.original_coordinates = [] 1339 self.metadata = [] 1340 X = [] 1341 for ( 1342 names, 1343 lengths, 1344 coords, 1345 bp_dictionary, 1346 min_frames, 1347 max_frames, 1348 metadata, 1349 ) in dict_list: 1350 X += names 1351 self.original_coordinates += coords 1352 self.visibility.update(bp_dictionary) 1353 self.min_frames.update(min_frames) 1354 self.max_frames.update(max_frames) 1355 if metadata is not None: 1356 self.metadata += metadata 1357 del dict_list 1358 1359 if len(self.metadata) != len(self.original_coordinates): 1360 self.metadata = None 1361 else: 1362 self.metadata = np.array(self.metadata) 1363 self.min_frames = dict(self.min_frames) 1364 self.max_frames = dict(self.max_frames) 1365 self.original_coordinates = np.array(self.original_coordinates) 1366 return np.array(X) 1367 1368 @classmethod 1369 def get_file_ids( 1370 cls, 1371 filenames: Set = None, 1372 data_path: Union[str, Set] = None, 1373 file_paths: Set = None, 1374 *args, 1375 **kwargs, 1376 ) -> List: 1377 """ 1378 Process data parameters and return a list of ids of the videos that should 1379 be processed by the __init__ function 1380 1381 Parameters 1382 ---------- 1383 filenames : set, optional 1384 a set of string filenames to search for (only basenames, not the whole paths) 1385 data_path : str | set, optional 1386 the path to the folder where the pose and feature files are stored or a set of such paths 1387 (not passed if creating from key objects or from `file_paths`) 1388 file_paths : set, optional 1389 a set of string paths to the pose and feature files 1390 (not passed if creating from key objects or from `data_path`) 1391 1392 Returns 1393 ------- 1394 video_ids : list 1395 a list of video file ids 1396 """ 1397 1398 if file_paths is None: 1399 file_paths = [] 1400 if data_path is not None: 1401 if isinstance(data_path, str): 1402 data_path = [data_path] 1403 file_paths = [] 1404 for folder in data_path: 1405 file_paths += [os.path.join(folder, x) for x in os.listdir(folder)] 1406 ids = set() 1407 for f in file_paths: 1408 if os.path.basename(f) in filenames: 1409 ids.add(os.path.basename(f)) 1410 ids = sorted(ids) 1411 return ids 1412 1413 @abstractmethod 1414 def _open_file(self, filename: str) -> List: 1415 """ 1416 Open a file and make a list of sequences 1417 1418 The sequence objects should contain information about all clips in one video. The sequences and 1419 video ids will be processed in the `_get_data` function. 1420 1421 Parameters 1422 ---------- 1423 filename : str 1424 the name of the file 1425 1426 Returns 1427 ------- 1428 video_tuples : list 1429 a list of video tuples: `(video_id, sequence)` 1430 """ 1431 1432 @abstractmethod 1433 def _get_data( 1434 self, video_id: str, sequence, default_agent_name: str 1435 ) -> Tuple[Dict, Optional[Dict]]: 1436 """ 1437 Get the keypoint dataframes from a sequence 1438 1439 The sequences and video ids are generated in the `_open_file` function. 1440 In `data_dictionary`, the keys are clip ids and the values are `pandas` dataframes with two-level indices. 1441 The first level is the frame numbers and the second is the body part names. The dataframes should have from 1442 two to four columns labeled `"x"`, `"y"` and (optionally) `"z"` and `"likelihood"`. Each frame should have 1443 information on all the body parts. You don't have to filter the data in any way or fill the nans, it will 1444 be done automatically. 1445 1446 Parameters 1447 ---------- 1448 video_id : str 1449 the video id 1450 sequence 1451 an object containing information about all clips in one video 1452 default_agent_name 1453 1454 Returns 1455 ------- 1456 data dictionary : dict 1457 a dictionary where the keys are clip ids and the values are keypoint dataframes (see above for details) 1458 metadata_dictionary : dict 1459 a dictionary where the keys are clip ids and the values are metadata objects (can be any additional information, 1460 like the annotator tag; for no metadata pass `None`) 1461 """
An implementation of dlc2action.data.InputStore for datasets where input data files correspond to multiple videos
@classmethod
def
get_file_ids( cls, filenames: Set = None, data_path: Union[str, Set] = None, file_paths: Set = None, *args, **kwargs) -> List:
1368 @classmethod 1369 def get_file_ids( 1370 cls, 1371 filenames: Set = None, 1372 data_path: Union[str, Set] = None, 1373 file_paths: Set = None, 1374 *args, 1375 **kwargs, 1376 ) -> List: 1377 """ 1378 Process data parameters and return a list of ids of the videos that should 1379 be processed by the __init__ function 1380 1381 Parameters 1382 ---------- 1383 filenames : set, optional 1384 a set of string filenames to search for (only basenames, not the whole paths) 1385 data_path : str | set, optional 1386 the path to the folder where the pose and feature files are stored or a set of such paths 1387 (not passed if creating from key objects or from `file_paths`) 1388 file_paths : set, optional 1389 a set of string paths to the pose and feature files 1390 (not passed if creating from key objects or from `data_path`) 1391 1392 Returns 1393 ------- 1394 video_ids : list 1395 a list of video file ids 1396 """ 1397 1398 if file_paths is None: 1399 file_paths = [] 1400 if data_path is not None: 1401 if isinstance(data_path, str): 1402 data_path = [data_path] 1403 file_paths = [] 1404 for folder in data_path: 1405 file_paths += [os.path.join(folder, x) for x in os.listdir(folder)] 1406 ids = set() 1407 for f in file_paths: 1408 if os.path.basename(f) in filenames: 1409 ids.add(os.path.basename(f)) 1410 ids = sorted(ids) 1411 return ids
Process data parameters and return a list of ids of the videos that should be processed by the __init__ function
Parameters
filenames : set, optional
a set of string filenames to search for (only basenames, not the whole paths)
data_path : str | set, optional
the path to the folder where the pose and feature files are stored or a set of such paths
(not passed if creating from key objects or from file_paths)
file_paths : set, optional
a set of string paths to the pose and feature files
(not passed if creating from key objects or from data_path)
Returns
video_ids : list a list of video file ids
Inherited Members
- GeneralInputStore
- GeneralInputStore
- data_suffix
- get_folder
- remove
- key_objects
- load_from_key_objects
- to_ram
- get_original_coordinates
- create_subsample
- get_video_id
- get_clip_id
- get_clip_length
- get_clip_start_end
- get_clip_start
- get_visibility
- get_annotation_objects
- get_bodyparts
- get_coords
- get_n_frames
- get_likelihood
- get_indices
- get_tag
1464class DLCTrackStore(FileInputStore): 1465 """ 1466 DLC track data 1467 1468 Assumes the following file structure: 1469 ``` 1470 data_path 1471 ├── video1DLC1000.pickle 1472 ├── video2DLC400.pickle 1473 ├── video1_features.pt 1474 └── video2_features.pt 1475 ``` 1476 Here `data_suffix` is `{'DLC1000.pickle', 'DLC400.pickle'}` and `feature_suffix` (optional) is `'_features.pt'`. 1477 1478 The feature files should to be dictionaries where keys are clip IDs (e.g. animal names) and values are 1479 feature values (arrays of shape `(#frames, #features)`). If the arrays are shaped as `(#features, #frames)`, 1480 set `transpose_features` to `True`. 1481 1482 The files can be saved with `numpy.save()` (with `.npy` extension), `torch.save()` (with `.pt` extension) or 1483 with `pickle.dump()` (with `.pickle` or `.pkl` extension). 1484 """ 1485 1486 def _open_data( 1487 self, filename: str, default_agent_name: str 1488 ) -> Tuple[Dict, Optional[Dict]]: 1489 """ 1490 Load the keypoints from filename and organize them in a dictionary 1491 1492 In `data_dictionary`, the keys are clip ids and the values are `pandas` dataframes with two-level indices. 1493 The first level is the frame numbers and the second is the body part names. The dataframes should have from 1494 two to four columns labeled `"x"`, `"y"` and (optionally) `"z"` and `"likelihood"`. Each frame should have 1495 information on all the body parts. You don't have to filter the data in any way or fill the nans, it will 1496 be done automatically. 1497 1498 Parameters 1499 ---------- 1500 filename : str 1501 path to the pose file 1502 default_clip_name : str 1503 the name to assign to a clip if it does not have a name in the raw data 1504 1505 Returns 1506 ------- 1507 data dictionary : dict 1508 a dictionary where the keys are clip ids and the values are keypoint dataframes (see above for details) 1509 metadata_dictionary : dict 1510 a dictionary where the keys are clip ids and the values are metadata objects (can be any additional information, 1511 like the annotator tag; for no metadata pass `None`) 1512 """ 1513 1514 if filename.endswith("h5"): 1515 temp = pd.read_hdf(filename) 1516 temp = temp.droplevel("scorer", axis=1) 1517 else: 1518 temp = pd.read_csv(filename, header=[1, 2]) 1519 temp.columns.names = ["bodyparts", "coords"] 1520 if "individuals" not in temp.columns.names: 1521 old_idx = temp.columns.to_frame() 1522 old_idx.insert(0, "individuals", self.default_agent_name) 1523 temp.columns = pd.MultiIndex.from_frame(old_idx) 1524 df = temp.stack(["individuals", "bodyparts"]) 1525 idx = pd.MultiIndex.from_product( 1526 [df.index.levels[0], df.index.levels[1], df.index.levels[2]], 1527 names=df.index.names, 1528 ) 1529 df = df.reindex(idx).fillna(value=0) 1530 animals = sorted(list(df.index.levels[1])) 1531 dic = {} 1532 for ind in animals: 1533 coord = df.iloc[df.index.get_level_values(1) == ind].droplevel(1) 1534 coord = coord[["x", "y", "likelihood"]] 1535 dic[ind] = coord 1536 1537 return dic, None
DLC track data
Assumes the following file structure:
data_path
├── video1DLC1000.pickle
├── video2DLC400.pickle
├── video1_features.pt
└── video2_features.pt
Here data_suffix is {'DLC1000.pickle', 'DLC400.pickle'} and feature_suffix (optional) is '_features.pt'.
The feature files should to be dictionaries where keys are clip IDs (e.g. animal names) and values are
feature values (arrays of shape (#frames, #features)). If the arrays are shaped as (#features, #frames),
set transpose_features to True.
The files can be saved with numpy.save() (with .npy extension), torch.save() (with .pt extension) or
with pickle.dump() (with .pickle or .pkl extension).
Inherited Members
- GeneralInputStore
- GeneralInputStore
- data_suffix
- get_folder
- remove
- key_objects
- load_from_key_objects
- to_ram
- get_original_coordinates
- create_subsample
- get_video_id
- get_clip_id
- get_clip_length
- get_clip_start_end
- get_clip_start
- get_visibility
- get_annotation_objects
- get_file_ids
- get_bodyparts
- get_coords
- get_n_frames
- get_likelihood
- get_indices
- get_tag
1540class DLCTrackletStore(FileInputStore): 1541 """ 1542 DLC tracklet data 1543 1544 Assumes the following file structure: 1545 ``` 1546 data_path 1547 ├── video1DLC1000.pickle 1548 ├── video2DLC400.pickle 1549 ├── video1_features.pt 1550 └── video2_features.pt 1551 ``` 1552 Here `data_suffix` is `{'DLC1000.pickle', 'DLC400.pickle'}` and `feature_suffix` (optional) is `'_features.pt'`. 1553 1554 The feature files should to be dictionaries where keys are clip IDs (e.g. animal names) and values are 1555 feature values (arrays of shape `(#frames, #features)`). If the arrays are shaped as `(#features, #frames)`, 1556 set `transpose_features` to `True`. 1557 1558 The files can be saved with `numpy.save()` (with `.npy` extension), `torch.save()` (with `.pt` extension) or 1559 with `pickle.dump()` (with `.pickle` or `.pkl` extension). 1560 """ 1561 1562 def _open_data( 1563 self, filename: str, default_agent_name: str 1564 ) -> Tuple[Dict, Optional[Dict]]: 1565 """ 1566 Load the keypoints from filename and organize them in a dictionary 1567 1568 In `data_dictionary`, the keys are clip ids and the values are `pandas` dataframes with two-level indices. 1569 The first level is the frame numbers and the second is the body part names. The dataframes should have from 1570 two to four columns labeled `"x"`, `"y"` and (optionally) `"z"` and `"likelihood"`. Each frame should have 1571 information on all the body parts. You don't have to filter the data in any way or fill the nans, it will 1572 be done automatically. 1573 1574 Parameters 1575 ---------- 1576 filename : str 1577 path to the pose file 1578 default_clip_name : str 1579 the name to assign to a clip if it does not have a name in the raw data 1580 1581 Returns 1582 ------- 1583 data dictionary : dict 1584 a dictionary where the keys are clip ids and the values are keypoint dataframes (see above for details) 1585 metadata_dictionary : dict 1586 a dictionary where the keys are clip ids and the values are metadata objects (can be any additional information, 1587 like the annotator tag; for no metadata pass `None`) 1588 """ 1589 1590 output = {} 1591 with open(filename, "rb") as f: 1592 data_p = pickle.load(f) 1593 header = data_p["header"] 1594 bodyparts = header.unique("bodyparts") 1595 1596 keys = sorted([key for key in data_p.keys() if key != "header"]) 1597 min_frames = defaultdict(lambda: 10**5) 1598 max_frames = defaultdict(lambda: 0) 1599 for tr_id in keys: 1600 coords = {} 1601 fr_i = int(list(data_p[tr_id].keys())[0][5:]) - 1 1602 for frame in data_p[tr_id]: 1603 count = 0 1604 while int(frame[5:]) > fr_i + 1: 1605 count += 1 1606 fr_i = fr_i + 1 1607 if count <= 3: 1608 for bp, name in enumerate(bodyparts): 1609 coords[(fr_i, name)] = coords[(fr_i - 1, name)] 1610 else: 1611 for bp, name in enumerate(bodyparts): 1612 coords[(fr_i, name)] = np.zeros( 1613 coords[(fr_i - 1, name)].shape 1614 ) 1615 fr_i = int(frame[5:]) 1616 if fr_i > max_frames[f"ind{tr_id}"]: 1617 max_frames[f"ind{tr_id}"] = fr_i 1618 if fr_i < min_frames[f"ind{tr_id}"]: 1619 min_frames[f"ind{tr_id}"] = fr_i 1620 for bp, name in enumerate(bodyparts): 1621 coords[(fr_i, name)] = data_p[tr_id][frame][bp][:3] 1622 1623 output[f"ind{tr_id}"] = pd.DataFrame( 1624 data=coords, index=["x", "y", "likelihood"] 1625 ).T 1626 return output, None
DLC tracklet data
Assumes the following file structure:
data_path
├── video1DLC1000.pickle
├── video2DLC400.pickle
├── video1_features.pt
└── video2_features.pt
Here data_suffix is {'DLC1000.pickle', 'DLC400.pickle'} and feature_suffix (optional) is '_features.pt'.
The feature files should to be dictionaries where keys are clip IDs (e.g. animal names) and values are
feature values (arrays of shape (#frames, #features)). If the arrays are shaped as (#features, #frames),
set transpose_features to True.
The files can be saved with numpy.save() (with .npy extension), torch.save() (with .pt extension) or
with pickle.dump() (with .pickle or .pkl extension).
Inherited Members
- GeneralInputStore
- GeneralInputStore
- data_suffix
- get_folder
- remove
- key_objects
- load_from_key_objects
- to_ram
- get_original_coordinates
- create_subsample
- get_video_id
- get_clip_id
- get_clip_length
- get_clip_start_end
- get_clip_start
- get_visibility
- get_annotation_objects
- get_file_ids
- get_bodyparts
- get_coords
- get_n_frames
- get_likelihood
- get_indices
- get_tag
1629class PKUMMDInputStore(FileInputStore): 1630 """ 1631 PKU-MMD data 1632 1633 Assumes the following file structure: 1634 ``` 1635 data_path 1636 ├── 0073-R.txt 1637 ... 1638 └── 0274-L.txt 1639 ``` 1640 """ 1641 1642 data_suffix = ".txt" 1643 1644 def __init__( 1645 self, 1646 video_order: str = None, 1647 data_path: Union[str, Set] = None, 1648 file_paths: Set = None, 1649 feature_save_path: str = None, 1650 feature_extraction: str = "kinematic", 1651 len_segment: int = 128, 1652 overlap: int = 0, 1653 key_objects: Tuple = None, 1654 num_cpus: int = None, 1655 interactive: bool = False, 1656 feature_extraction_pars: Dict = None, 1657 *args, 1658 **kwargs, 1659 ) -> None: 1660 """ 1661 Parameters 1662 ---------- 1663 video_order : list, optional 1664 a list of video ids that should be processed in the same order (not passed if creating from key objects) 1665 data_path : str | set, optional 1666 the path to the folder where the pose and feature files are stored or a set of such paths 1667 (not passed if creating from key objects or from `file_paths`) 1668 file_paths : set, optional 1669 a set of string paths to the pose and feature files 1670 (not passed if creating from key objects or from `data_path`) 1671 feature_save_path : str, optional 1672 the path to the folder where pre-processed files are stored (not passed if creating from key objects) 1673 feature_extraction : str, default 'kinematic' 1674 the feature extraction method (run options.feature_extractors to see available options) 1675 len_segment : int, default 128 1676 the length of the segments in which the data should be cut (in frames) 1677 overlap : int, default 0 1678 the length of the overlap between neighboring segments (in frames) 1679 interactive : bool, default False 1680 if True, distances between two agents are included; if False, only the first agent features are computed 1681 key_objects : tuple, optional 1682 a tuple of key objects 1683 num_cpus : int, optional 1684 the number of cpus to use in data processing 1685 feature_extraction_pars : dict, optional 1686 parameters of the feature extractor 1687 """ 1688 1689 if feature_extraction_pars is None: 1690 feature_extraction_pars = {} 1691 feature_extraction_pars["interactive"] = interactive 1692 self.interactive = interactive 1693 super().__init__( 1694 video_order, 1695 data_path, 1696 file_paths, 1697 data_suffix=".txt", 1698 data_prefix=None, 1699 feature_suffix=None, 1700 convert_int_indices=False, 1701 feature_save_path=feature_save_path, 1702 canvas_shape=[1, 1, 1], 1703 len_segment=len_segment, 1704 overlap=overlap, 1705 feature_extraction=feature_extraction, 1706 ignored_clips=None, 1707 ignored_bodyparts=None, 1708 default_agent_name="ind0", 1709 key_objects=key_objects, 1710 likelihood_threshold=0, 1711 num_cpus=num_cpus, 1712 frame_limit=1, 1713 interactive=interactive, 1714 feature_extraction_pars=feature_extraction_pars, 1715 ) 1716 1717 def _open_data( 1718 self, filename: str, default_clip_name: str 1719 ) -> Tuple[Dict, Optional[Dict]]: 1720 """ 1721 Load the keypoints from filename and organize them in a dictionary 1722 1723 In `data_dictionary`, the keys are clip ids and the values are `pandas` dataframes with two-level indices. 1724 The first level is the frame numbers and the second is the body part names. The dataframes should have from 1725 two to four columns labeled `"x"`, `"y"` and (optionally) `"z"` and `"likelihood"`. Each frame should have 1726 information on all the body parts. You don't have to filter the data in any way or fill the nans, it will 1727 be done automatically. 1728 1729 Parameters 1730 ---------- 1731 filename : str 1732 path to the pose file 1733 default_clip_name : str 1734 the name to assign to a clip if it does not have a name in the raw data 1735 1736 Returns 1737 ------- 1738 data dictionary : dict 1739 a dictionary where the keys are clip ids and the values are keypoint dataframes (see above for details) 1740 metadata_dictionary : dict 1741 a dictionary where the keys are clip ids and the values are metadata objects (can be any additional information, 1742 like the annotator tag; for no metadata pass `None`) 1743 """ 1744 1745 keypoint_dict = {"0": [], "1": []} 1746 with open(filename) as f: 1747 for line in f.readlines(): 1748 line_data = list(map(float, line.split())) 1749 line_data = np.array(line_data) 1750 line_data = line_data.reshape((2, 25, 3))[:, :, [0, 2, 1]] 1751 for ind in keypoint_dict: 1752 keypoint_dict[ind].append(line_data[int(ind)]) 1753 for ind in keypoint_dict: 1754 data = np.stack(keypoint_dict[ind]) 1755 mi = pd.MultiIndex.from_product( 1756 [list(range(data.shape[0])), list(range(data.shape[1]))] 1757 ) 1758 data = data.reshape((-1, 3)) 1759 keypoint_dict[ind] = pd.DataFrame( 1760 data=data, index=mi, columns=["x", "y", "z"] 1761 ) 1762 if not self.interactive: 1763 keypoint_dict.pop("1") 1764 return keypoint_dict, None
PKU-MMD data
Assumes the following file structure:
data_path
├── 0073-R.txt
...
└── 0274-L.txt
PKUMMDInputStore( video_order: str = None, data_path: Union[str, Set] = None, file_paths: Set = None, feature_save_path: str = None, feature_extraction: str = 'kinematic', len_segment: int = 128, overlap: int = 0, key_objects: Tuple = None, num_cpus: int = None, interactive: bool = False, feature_extraction_pars: Dict = None, *args, **kwargs)
1644 def __init__( 1645 self, 1646 video_order: str = None, 1647 data_path: Union[str, Set] = None, 1648 file_paths: Set = None, 1649 feature_save_path: str = None, 1650 feature_extraction: str = "kinematic", 1651 len_segment: int = 128, 1652 overlap: int = 0, 1653 key_objects: Tuple = None, 1654 num_cpus: int = None, 1655 interactive: bool = False, 1656 feature_extraction_pars: Dict = None, 1657 *args, 1658 **kwargs, 1659 ) -> None: 1660 """ 1661 Parameters 1662 ---------- 1663 video_order : list, optional 1664 a list of video ids that should be processed in the same order (not passed if creating from key objects) 1665 data_path : str | set, optional 1666 the path to the folder where the pose and feature files are stored or a set of such paths 1667 (not passed if creating from key objects or from `file_paths`) 1668 file_paths : set, optional 1669 a set of string paths to the pose and feature files 1670 (not passed if creating from key objects or from `data_path`) 1671 feature_save_path : str, optional 1672 the path to the folder where pre-processed files are stored (not passed if creating from key objects) 1673 feature_extraction : str, default 'kinematic' 1674 the feature extraction method (run options.feature_extractors to see available options) 1675 len_segment : int, default 128 1676 the length of the segments in which the data should be cut (in frames) 1677 overlap : int, default 0 1678 the length of the overlap between neighboring segments (in frames) 1679 interactive : bool, default False 1680 if True, distances between two agents are included; if False, only the first agent features are computed 1681 key_objects : tuple, optional 1682 a tuple of key objects 1683 num_cpus : int, optional 1684 the number of cpus to use in data processing 1685 feature_extraction_pars : dict, optional 1686 parameters of the feature extractor 1687 """ 1688 1689 if feature_extraction_pars is None: 1690 feature_extraction_pars = {} 1691 feature_extraction_pars["interactive"] = interactive 1692 self.interactive = interactive 1693 super().__init__( 1694 video_order, 1695 data_path, 1696 file_paths, 1697 data_suffix=".txt", 1698 data_prefix=None, 1699 feature_suffix=None, 1700 convert_int_indices=False, 1701 feature_save_path=feature_save_path, 1702 canvas_shape=[1, 1, 1], 1703 len_segment=len_segment, 1704 overlap=overlap, 1705 feature_extraction=feature_extraction, 1706 ignored_clips=None, 1707 ignored_bodyparts=None, 1708 default_agent_name="ind0", 1709 key_objects=key_objects, 1710 likelihood_threshold=0, 1711 num_cpus=num_cpus, 1712 frame_limit=1, 1713 interactive=interactive, 1714 feature_extraction_pars=feature_extraction_pars, 1715 )
Parameters
video_order : list, optional
a list of video ids that should be processed in the same order (not passed if creating from key objects)
data_path : str | set, optional
the path to the folder where the pose and feature files are stored or a set of such paths
(not passed if creating from key objects or from file_paths)
file_paths : set, optional
a set of string paths to the pose and feature files
(not passed if creating from key objects or from data_path)
feature_save_path : str, optional
the path to the folder where pre-processed files are stored (not passed if creating from key objects)
feature_extraction : str, default 'kinematic'
the feature extraction method (run options.feature_extractors to see available options)
len_segment : int, default 128
the length of the segments in which the data should be cut (in frames)
overlap : int, default 0
the length of the overlap between neighboring segments (in frames)
interactive : bool, default False
if True, distances between two agents are included; if False, only the first agent features are computed
key_objects : tuple, optional
a tuple of key objects
num_cpus : int, optional
the number of cpus to use in data processing
feature_extraction_pars : dict, optional
parameters of the feature extractor
Inherited Members
- GeneralInputStore
- get_folder
- remove
- key_objects
- load_from_key_objects
- to_ram
- get_original_coordinates
- create_subsample
- get_video_id
- get_clip_id
- get_clip_length
- get_clip_start_end
- get_clip_start
- get_visibility
- get_annotation_objects
- get_file_ids
- get_bodyparts
- get_coords
- get_n_frames
- get_likelihood
- get_indices
- get_tag
1767class CalMS21InputStore(SequenceInputStore): 1768 """ 1769 CalMS21 data 1770 1771 Use the `'random:test_from_name:{name}'` and `'val-from-name:{val_name}:test-from-name:{test_name}'` 1772 partitioning methods with `'train'`, `'test'` and `'unlabeled'` names to separate into train, test and validation 1773 subsets according to the original files. For example, with `'val-from-name:test:test-from-name:unlabeled'` 1774 the data from the test file will go into validation and the unlabeled files will be the test. 1775 1776 Assumes the following file structure: 1777 ``` 1778 data_path 1779 ├── calms21_task1_train.npy 1780 ├── calms21_task1_test.npy 1781 ├── calms21_task1_test_features.npy 1782 ├── calms21_task1_test_features.npy 1783 ├── calms21_unlabeled_videos_part1.npy 1784 ├── calms21_unlabeled_videos_part1.npy 1785 ├── calms21_unlabeled_videos_part2.npy 1786 └── calms21_unlabeled_videos_part3.npy 1787 ``` 1788 """ 1789 1790 def __init__( 1791 self, 1792 video_order: List = None, 1793 data_path: Union[Set, str] = None, 1794 file_paths: Set = None, 1795 task_n: int = 1, 1796 include_task1: bool = True, 1797 feature_save_path: str = None, 1798 len_segment: int = 128, 1799 overlap: int = 0, 1800 feature_extraction: str = "kinematic", 1801 key_objects: Dict = None, 1802 treba_files: bool = False, 1803 num_cpus: int = None, 1804 feature_extraction_pars: Dict = None, 1805 *args, 1806 **kwargs, 1807 ) -> None: 1808 """ 1809 Parameters 1810 ---------- 1811 video_order : list, optional 1812 a list of video ids that should be processed in the same order (not passed if creating from key objects) 1813 data_path : str | set, optional 1814 the path to the folder where the pose and feature files are stored or a set of such paths 1815 (not passed if creating from key objects or from `file_paths`) 1816 file_paths : set, optional 1817 a set of string paths to the pose and feature files 1818 (not passed if creating from key objects or from `data_path`) 1819 task_n : [1, 2] 1820 the number of the task 1821 include_task1 : bool, default True 1822 include task 1 data to training set 1823 feature_save_path : str, optional 1824 the path to the folder where pre-processed files are stored (not passed if creating from key objects) 1825 len_segment : int, default 128 1826 the length of the segments in which the data should be cut (in frames) 1827 overlap : int, default 0 1828 the length of the overlap between neighboring segments (in frames) 1829 feature_extraction : str, default 'kinematic' 1830 the feature extraction method (see options.feature_extractors for available options) 1831 ignored_bodyparts : list, optional 1832 list of strings of bodypart names to ignore 1833 key_objects : tuple, optional 1834 a tuple of key objects 1835 treba_files : bool, default False 1836 if `True`, TREBA feature files will be loaded 1837 num_cpus : int, optional 1838 the number of cpus to use in data processing 1839 feature_extraction_pars : dict, optional 1840 parameters of the feature extractor 1841 """ 1842 1843 self.task_n = int(task_n) 1844 self.include_task1 = include_task1 1845 self.treba_files = treba_files 1846 if feature_extraction_pars is not None: 1847 feature_extraction_pars["interactive"] = True 1848 1849 super().__init__( 1850 video_order, 1851 data_path, 1852 file_paths, 1853 data_prefix=None, 1854 feature_suffix=None, 1855 convert_int_indices=False, 1856 feature_save_path=feature_save_path, 1857 canvas_shape=[1024, 570], 1858 len_segment=len_segment, 1859 overlap=overlap, 1860 feature_extraction=feature_extraction, 1861 ignored_clips=None, 1862 ignored_bodyparts=None, 1863 default_agent_name="ind0", 1864 key_objects=key_objects, 1865 likelihood_threshold=0, 1866 num_cpus=num_cpus, 1867 frame_limit=1, 1868 feature_extraction_pars=feature_extraction_pars, 1869 ) 1870 1871 @classmethod 1872 def get_file_ids( 1873 cls, 1874 task_n: int = 1, 1875 include_task1: bool = False, 1876 treba_files: bool = False, 1877 data_path: Union[str, Set] = None, 1878 file_paths=None, 1879 *args, 1880 **kwargs, 1881 ) -> Iterable: 1882 """ 1883 Process data parameters and return a list of ids of the videos that should 1884 be processed by the __init__ function 1885 1886 Parameters 1887 ---------- 1888 task_n : {1, 2, 3} 1889 the index of the CalMS21 challenge task 1890 include_task1 : bool, default False 1891 if `True`, the training file of the task 1 will be loaded 1892 treba_files : bool, default False 1893 if `True`, the TREBA feature files will be loaded 1894 filenames : set, optional 1895 a set of string filenames to search for (only basenames, not the whole paths) 1896 data_path : str | set, optional 1897 the path to the folder where the pose and feature files are stored or a set of such paths 1898 (not passed if creating from key objects or from `file_paths`) 1899 file_paths : set, optional 1900 a set of string paths to the pose and feature files 1901 (not passed if creating from key objects or from `data_path`) 1902 1903 Returns 1904 ------- 1905 video_ids : list 1906 a list of video file ids 1907 """ 1908 1909 task_n = int(task_n) 1910 if task_n == 1: 1911 include_task1 = False 1912 files = [] 1913 if treba_files: 1914 postfix = "_features" 1915 else: 1916 postfix = "" 1917 files.append(f"calms21_task{task_n}_train{postfix}.npy") 1918 files.append(f"calms21_task{task_n}_test{postfix}.npy") 1919 if include_task1: 1920 files.append(f"calms21_task1_train{postfix}.npy") 1921 for i in range(1, 5): 1922 files.append(f"calms21_unlabeled_videos_part{i}{postfix}.npy") 1923 filenames = set(files) 1924 return SequenceInputStore.get_file_ids(filenames, data_path, file_paths) 1925 1926 def _open_file(self, filename: str) -> List: 1927 """ 1928 Open a file and make a list of sequences 1929 1930 The sequence objects should contain information about all clips in one video. The sequences and 1931 video ids will be processed in the `_get_data` function. 1932 1933 Parameters 1934 ---------- 1935 filename : str 1936 the name of the file 1937 1938 Returns 1939 ------- 1940 video_tuples : list 1941 a list of video tuples: `(video_id, sequence)` 1942 """ 1943 1944 if os.path.basename(filename).startswith("calms21_unlabeled_videos"): 1945 mode = "unlabeled" 1946 elif os.path.basename(filename).startswith(f"calms21_task{self.task_n}_test"): 1947 mode = "test" 1948 else: 1949 mode = "train" 1950 data_dict = np.load(filename, allow_pickle=True).item() 1951 data = {} 1952 keys = list(data_dict.keys()) 1953 for key in keys: 1954 data.update(data_dict[key]) 1955 data_dict.pop(key) 1956 dict_list = [(f'{mode}--{k.split("/")[-1]}', v) for k, v in data.items()] 1957 return dict_list 1958 1959 def _get_data( 1960 self, video_id: str, sequence, default_agent_name: str 1961 ) -> Tuple[Dict, Optional[Dict]]: 1962 """ 1963 Get the keypoint dataframes from a sequence 1964 1965 The sequences and video ids are generated in the `_open_file` function. 1966 In `data_dictionary`, the keys are clip ids and the values are `pandas` dataframes with two-level indices. 1967 The first level is the frame numbers and the second is the body part names. The dataframes should have from 1968 two to four columns labeled `"x"`, `"y"` and (optionally) `"z"` and `"likelihood"`. Each frame should have 1969 information on all the body parts. You don't have to filter the data in any way or fill the nans, it will 1970 be done automatically. 1971 1972 Parameters 1973 ---------- 1974 video_id : str 1975 the video id 1976 sequence 1977 an object containing information about all clips in one video 1978 default_agent_name 1979 1980 Returns 1981 ------- 1982 data dictionary : dict 1983 a dictionary where the keys are clip ids and the values are keypoint dataframes (see above for details) 1984 metadata_dictionary : dict 1985 a dictionary where the keys are clip ids and the values are metadata objects (can be any additional information, 1986 like the annotator tag; for no metadata pass `None`) 1987 """ 1988 1989 if "metadata" in sequence: 1990 annotator = sequence["metadata"]["annotator-id"] 1991 else: 1992 annotator = 0 1993 bodyparts = [ 1994 "nose", 1995 "left ear", 1996 "right ear", 1997 "neck", 1998 "left hip", 1999 "right hip", 2000 "tail", 2001 ] 2002 columns = ["x", "y"] 2003 if "keypoints" in sequence: 2004 sequence = sequence["keypoints"] 2005 index = pd.MultiIndex.from_product([range(sequence.shape[0]), bodyparts]) 2006 data = { 2007 "mouse1": pd.DataFrame( 2008 data=(sequence[:, 0, :, :]).transpose((0, 2, 1)).reshape(-1, 2), 2009 columns=columns, 2010 index=index, 2011 ), 2012 "mouse2": pd.DataFrame( 2013 data=(sequence[:, 1, :, :]).transpose((0, 2, 1)).reshape(-1, 2), 2014 columns=columns, 2015 index=index, 2016 ), 2017 } 2018 else: 2019 sequence = sequence["features"] 2020 mice = sequence[:, :-32].reshape((-1, 2, 2, 7)) 2021 index = pd.MultiIndex.from_product([range(mice.shape[0]), bodyparts]) 2022 data = { 2023 "mouse1": pd.DataFrame( 2024 data=(mice[:, 0, :, :]).transpose((0, 2, 1)).reshape(-1, 2), 2025 columns=columns, 2026 index=index, 2027 ), 2028 "mouse2": pd.DataFrame( 2029 data=(mice[:, 1, :, :]).transpose((0, 2, 1)).reshape(-1, 2), 2030 columns=columns, 2031 index=index, 2032 ), 2033 "loaded": sequence[:, -32:], 2034 } 2035 metadata = {k: annotator for k in data.keys()} 2036 return data, metadata
CalMS21 data
Use the 'random:test_from_name:{name}' and 'val-from-name:{val_name}:test-from-name:{test_name}'
partitioning methods with 'train', 'test' and 'unlabeled' names to separate into train, test and validation
subsets according to the original files. For example, with 'val-from-name:test:test-from-name:unlabeled'
the data from the test file will go into validation and the unlabeled files will be the test.
Assumes the following file structure:
data_path
├── calms21_task1_train.npy
├── calms21_task1_test.npy
├── calms21_task1_test_features.npy
├── calms21_task1_test_features.npy
├── calms21_unlabeled_videos_part1.npy
├── calms21_unlabeled_videos_part1.npy
├── calms21_unlabeled_videos_part2.npy
└── calms21_unlabeled_videos_part3.npy
CalMS21InputStore( video_order: List = None, data_path: Union[Set, str] = None, file_paths: Set = None, task_n: int = 1, include_task1: bool = True, feature_save_path: str = None, len_segment: int = 128, overlap: int = 0, feature_extraction: str = 'kinematic', key_objects: Dict = None, treba_files: bool = False, num_cpus: int = None, feature_extraction_pars: Dict = None, *args, **kwargs)
1790 def __init__( 1791 self, 1792 video_order: List = None, 1793 data_path: Union[Set, str] = None, 1794 file_paths: Set = None, 1795 task_n: int = 1, 1796 include_task1: bool = True, 1797 feature_save_path: str = None, 1798 len_segment: int = 128, 1799 overlap: int = 0, 1800 feature_extraction: str = "kinematic", 1801 key_objects: Dict = None, 1802 treba_files: bool = False, 1803 num_cpus: int = None, 1804 feature_extraction_pars: Dict = None, 1805 *args, 1806 **kwargs, 1807 ) -> None: 1808 """ 1809 Parameters 1810 ---------- 1811 video_order : list, optional 1812 a list of video ids that should be processed in the same order (not passed if creating from key objects) 1813 data_path : str | set, optional 1814 the path to the folder where the pose and feature files are stored or a set of such paths 1815 (not passed if creating from key objects or from `file_paths`) 1816 file_paths : set, optional 1817 a set of string paths to the pose and feature files 1818 (not passed if creating from key objects or from `data_path`) 1819 task_n : [1, 2] 1820 the number of the task 1821 include_task1 : bool, default True 1822 include task 1 data to training set 1823 feature_save_path : str, optional 1824 the path to the folder where pre-processed files are stored (not passed if creating from key objects) 1825 len_segment : int, default 128 1826 the length of the segments in which the data should be cut (in frames) 1827 overlap : int, default 0 1828 the length of the overlap between neighboring segments (in frames) 1829 feature_extraction : str, default 'kinematic' 1830 the feature extraction method (see options.feature_extractors for available options) 1831 ignored_bodyparts : list, optional 1832 list of strings of bodypart names to ignore 1833 key_objects : tuple, optional 1834 a tuple of key objects 1835 treba_files : bool, default False 1836 if `True`, TREBA feature files will be loaded 1837 num_cpus : int, optional 1838 the number of cpus to use in data processing 1839 feature_extraction_pars : dict, optional 1840 parameters of the feature extractor 1841 """ 1842 1843 self.task_n = int(task_n) 1844 self.include_task1 = include_task1 1845 self.treba_files = treba_files 1846 if feature_extraction_pars is not None: 1847 feature_extraction_pars["interactive"] = True 1848 1849 super().__init__( 1850 video_order, 1851 data_path, 1852 file_paths, 1853 data_prefix=None, 1854 feature_suffix=None, 1855 convert_int_indices=False, 1856 feature_save_path=feature_save_path, 1857 canvas_shape=[1024, 570], 1858 len_segment=len_segment, 1859 overlap=overlap, 1860 feature_extraction=feature_extraction, 1861 ignored_clips=None, 1862 ignored_bodyparts=None, 1863 default_agent_name="ind0", 1864 key_objects=key_objects, 1865 likelihood_threshold=0, 1866 num_cpus=num_cpus, 1867 frame_limit=1, 1868 feature_extraction_pars=feature_extraction_pars, 1869 )
Parameters
video_order : list, optional
a list of video ids that should be processed in the same order (not passed if creating from key objects)
data_path : str | set, optional
the path to the folder where the pose and feature files are stored or a set of such paths
(not passed if creating from key objects or from file_paths)
file_paths : set, optional
a set of string paths to the pose and feature files
(not passed if creating from key objects or from data_path)
task_n : [1, 2]
the number of the task
include_task1 : bool, default True
include task 1 data to training set
feature_save_path : str, optional
the path to the folder where pre-processed files are stored (not passed if creating from key objects)
len_segment : int, default 128
the length of the segments in which the data should be cut (in frames)
overlap : int, default 0
the length of the overlap between neighboring segments (in frames)
feature_extraction : str, default 'kinematic'
the feature extraction method (see options.feature_extractors for available options)
ignored_bodyparts : list, optional
list of strings of bodypart names to ignore
key_objects : tuple, optional
a tuple of key objects
treba_files : bool, default False
if True, TREBA feature files will be loaded
num_cpus : int, optional
the number of cpus to use in data processing
feature_extraction_pars : dict, optional
parameters of the feature extractor
@classmethod
def
get_file_ids( cls, task_n: int = 1, include_task1: bool = False, treba_files: bool = False, data_path: Union[str, Set] = None, file_paths=None, *args, **kwargs) -> Iterable:
1871 @classmethod 1872 def get_file_ids( 1873 cls, 1874 task_n: int = 1, 1875 include_task1: bool = False, 1876 treba_files: bool = False, 1877 data_path: Union[str, Set] = None, 1878 file_paths=None, 1879 *args, 1880 **kwargs, 1881 ) -> Iterable: 1882 """ 1883 Process data parameters and return a list of ids of the videos that should 1884 be processed by the __init__ function 1885 1886 Parameters 1887 ---------- 1888 task_n : {1, 2, 3} 1889 the index of the CalMS21 challenge task 1890 include_task1 : bool, default False 1891 if `True`, the training file of the task 1 will be loaded 1892 treba_files : bool, default False 1893 if `True`, the TREBA feature files will be loaded 1894 filenames : set, optional 1895 a set of string filenames to search for (only basenames, not the whole paths) 1896 data_path : str | set, optional 1897 the path to the folder where the pose and feature files are stored or a set of such paths 1898 (not passed if creating from key objects or from `file_paths`) 1899 file_paths : set, optional 1900 a set of string paths to the pose and feature files 1901 (not passed if creating from key objects or from `data_path`) 1902 1903 Returns 1904 ------- 1905 video_ids : list 1906 a list of video file ids 1907 """ 1908 1909 task_n = int(task_n) 1910 if task_n == 1: 1911 include_task1 = False 1912 files = [] 1913 if treba_files: 1914 postfix = "_features" 1915 else: 1916 postfix = "" 1917 files.append(f"calms21_task{task_n}_train{postfix}.npy") 1918 files.append(f"calms21_task{task_n}_test{postfix}.npy") 1919 if include_task1: 1920 files.append(f"calms21_task1_train{postfix}.npy") 1921 for i in range(1, 5): 1922 files.append(f"calms21_unlabeled_videos_part{i}{postfix}.npy") 1923 filenames = set(files) 1924 return SequenceInputStore.get_file_ids(filenames, data_path, file_paths)
Process data parameters and return a list of ids of the videos that should be processed by the __init__ function
Parameters
task_n : {1, 2, 3}
the index of the CalMS21 challenge task
include_task1 : bool, default False
if True, the training file of the task 1 will be loaded
treba_files : bool, default False
if True, the TREBA feature files will be loaded
filenames : set, optional
a set of string filenames to search for (only basenames, not the whole paths)
data_path : str | set, optional
the path to the folder where the pose and feature files are stored or a set of such paths
(not passed if creating from key objects or from file_paths)
file_paths : set, optional
a set of string paths to the pose and feature files
(not passed if creating from key objects or from data_path)
Returns
video_ids : list a list of video file ids
Inherited Members
- GeneralInputStore
- data_suffix
- get_folder
- remove
- key_objects
- load_from_key_objects
- to_ram
- get_original_coordinates
- create_subsample
- get_video_id
- get_clip_id
- get_clip_length
- get_clip_start_end
- get_clip_start
- get_visibility
- get_annotation_objects
- get_bodyparts
- get_coords
- get_n_frames
- get_likelihood
- get_indices
- get_tag
2039class Numpy3DInputStore(FileInputStore): 2040 """ 2041 3D data 2042 2043 Assumes the data files to be `numpy` arrays saved in `.npy` format with shape `(#frames, #keypoints, 3)`. 2044 2045 Assumes the following file structure: 2046 ``` 2047 data_path 2048 ├── video1_suffix1.npy 2049 ├── video2_suffix2.npy 2050 ├── video1_features.pt 2051 └── video2_features.pt 2052 ``` 2053 Here `data_suffix` is `{'_suffix1.npy', '_suffix1.npy'}` and `feature_suffix` (optional) is `'_features.pt'`. 2054 2055 The feature files should to be dictionaries where keys are clip IDs (e.g. animal names) and values are 2056 feature values (arrays of shape `(#frames, #features)`). If the arrays are shaped as `(#features, #frames)`, 2057 set `transpose_features` to `True`. 2058 2059 The files can be saved with `numpy.save()` (with `.npy` extension), `torch.save()` (with `.pt` extension) or 2060 with `pickle.dump()` (with `.pickle` or `.pkl` extension). 2061 """ 2062 2063 def __init__( 2064 self, 2065 video_order: List = None, 2066 data_path: Union[Set, str] = None, 2067 file_paths: Set = None, 2068 data_suffix: Union[Set, str] = None, 2069 data_prefix: Union[Set, str] = None, 2070 feature_suffix: Union[Set, str] = None, 2071 convert_int_indices: bool = True, 2072 feature_save_path: str = None, 2073 canvas_shape: List = None, 2074 len_segment: int = 128, 2075 overlap: int = 0, 2076 feature_extraction: str = "kinematic", 2077 ignored_clips: List = None, 2078 ignored_bodyparts: List = None, 2079 default_agent_name: str = "ind0", 2080 key_objects: Dict = None, 2081 likelihood_threshold: float = 0, 2082 num_cpus: int = None, 2083 frame_limit: int = 1, 2084 feature_extraction_pars: Dict = None, 2085 centered: bool = False, 2086 **kwargs, 2087 ) -> None: 2088 """ 2089 Parameters 2090 ---------- 2091 video_order : list, optional 2092 a list of video ids that should be processed in the same order (not passed if creating from key objects 2093 data_path : str | set, optional 2094 the path to the folder where the pose and feature files are stored or a set of such paths 2095 (not passed if creating from key objects or from `file_paths`) 2096 file_paths : set, optional 2097 a set of string paths to the pose and feature files 2098 (not passed if creating from key objects or from `data_path`) 2099 data_suffix : str | set, optional 2100 the suffix or the set of suffices such that the pose files are named {video_id}{data_suffix} 2101 (not passed if creating from key objects or if irrelevant for the dataset) 2102 data_prefix : str | set, optional 2103 the prefix or the set of prefixes such that the pose files for different video views of the same 2104 clip are named {prefix}{sep}{video_id}{data_suffix} (not passed if creating from key objects 2105 or if irrelevant for the dataset) 2106 feature_suffix : str | set, optional 2107 the suffix or the set of suffices such that the additional feature files are named 2108 {video_id}{feature_suffix} (and placed at the data_path folder) 2109 convert_int_indices : bool, default True 2110 if `True`, convert any integer key `i` in feature files to `'ind{i}'` 2111 feature_save_path : str, optional 2112 the path to the folder where pre-processed files are stored (not passed if creating from key objects) 2113 canvas_shape : List, default [1, 1] 2114 the canvas size where the pose is defined 2115 len_segment : int, default 128 2116 the length of the segments in which the data should be cut (in frames) 2117 overlap : int, default 0 2118 the length of the overlap between neighboring segments (in frames) 2119 feature_extraction : str, default 'kinematic' 2120 the feature extraction method (see options.feature_extractors for available options) 2121 ignored_clips : list, optional 2122 list of strings of clip ids to ignore 2123 ignored_bodyparts : list, optional 2124 list of strings of bodypart names to ignore 2125 default_agent_name : str, default 'ind0' 2126 the agent name used as default in the pose files for a single agent 2127 key_objects : tuple, optional 2128 a tuple of key objects 2129 likelihood_threshold : float, default 0 2130 coordinate values with likelihoods less than this value will be set to 'unknown' 2131 num_cpus : int, optional 2132 the number of cpus to use in data processing 2133 frame_limit : int, default 1 2134 clips shorter than this number of frames will be ignored 2135 feature_extraction_pars : dict, optional 2136 parameters of the feature extractor 2137 """ 2138 2139 super().__init__( 2140 video_order, 2141 data_path, 2142 file_paths, 2143 data_suffix=data_suffix, 2144 data_prefix=data_prefix, 2145 feature_suffix=feature_suffix, 2146 convert_int_indices=convert_int_indices, 2147 feature_save_path=feature_save_path, 2148 canvas_shape=canvas_shape, 2149 len_segment=len_segment, 2150 overlap=overlap, 2151 feature_extraction=feature_extraction, 2152 ignored_clips=ignored_clips, 2153 ignored_bodyparts=ignored_bodyparts, 2154 default_agent_name=default_agent_name, 2155 key_objects=key_objects, 2156 likelihood_threshold=likelihood_threshold, 2157 num_cpus=num_cpus, 2158 frame_limit=frame_limit, 2159 feature_extraction_pars=feature_extraction_pars, 2160 centered=centered, 2161 ) 2162 2163 def _open_data( 2164 self, filename: str, default_clip_name: str 2165 ) -> Tuple[Dict, Optional[Dict]]: 2166 """ 2167 Load the keypoints from filename and organize them in a dictionary 2168 2169 In `data_dictionary`, the keys are clip ids and the values are `pandas` dataframes with two-level indices. 2170 The first level is the frame numbers and the second is the body part names. The dataframes should have from 2171 two to four columns labeled `"x"`, `"y"` and (optionally) `"z"` and `"likelihood"`. Each frame should have 2172 information on all the body parts. You don't have to filter the data in any way or fill the nans, it will 2173 be done automatically. 2174 2175 Parameters 2176 ---------- 2177 filename : str 2178 path to the pose file 2179 default_clip_name : str 2180 the name to assign to a clip if it does not have a name in the raw data 2181 2182 Returns 2183 ------- 2184 data dictionary : dict 2185 a dictionary where the keys are clip ids and the values are keypoint dataframes (see above for details) 2186 metadata_dictionary : dict 2187 a dictionary where the keys are clip ids and the values are metadata objects (can be any additional information, 2188 like the annotator tag; for no metadata pass `None`) 2189 """ 2190 2191 data = np.load(filename) 2192 bodyparts = [str(i) for i in range(data.shape[1])] 2193 clip_id = self.default_agent_name 2194 columns = ["x", "y", "z"] 2195 index = pd.MultiIndex.from_product([range(data.shape[0]), bodyparts]) 2196 data_dict = { 2197 clip_id: pd.DataFrame( 2198 data=data.reshape(-1, 3), columns=columns, index=index 2199 ) 2200 } 2201 return data_dict, None
3D data
Assumes the data files to be numpy arrays saved in .npy format with shape (#frames, #keypoints, 3).
Assumes the following file structure:
data_path
├── video1_suffix1.npy
├── video2_suffix2.npy
├── video1_features.pt
└── video2_features.pt
Here data_suffix is {'_suffix1.npy', '_suffix1.npy'} and feature_suffix (optional) is '_features.pt'.
The feature files should to be dictionaries where keys are clip IDs (e.g. animal names) and values are
feature values (arrays of shape (#frames, #features)). If the arrays are shaped as (#features, #frames),
set transpose_features to True.
The files can be saved with numpy.save() (with .npy extension), torch.save() (with .pt extension) or
with pickle.dump() (with .pickle or .pkl extension).
Numpy3DInputStore( video_order: List = None, data_path: Union[Set, str] = None, file_paths: Set = None, data_suffix: Union[Set, str] = None, data_prefix: Union[Set, str] = None, feature_suffix: Union[Set, str] = None, convert_int_indices: bool = True, feature_save_path: str = None, canvas_shape: List = None, len_segment: int = 128, overlap: int = 0, feature_extraction: str = 'kinematic', ignored_clips: List = None, ignored_bodyparts: List = None, default_agent_name: str = 'ind0', key_objects: Dict = None, likelihood_threshold: float = 0, num_cpus: int = None, frame_limit: int = 1, feature_extraction_pars: Dict = None, centered: bool = False, **kwargs)
2063 def __init__( 2064 self, 2065 video_order: List = None, 2066 data_path: Union[Set, str] = None, 2067 file_paths: Set = None, 2068 data_suffix: Union[Set, str] = None, 2069 data_prefix: Union[Set, str] = None, 2070 feature_suffix: Union[Set, str] = None, 2071 convert_int_indices: bool = True, 2072 feature_save_path: str = None, 2073 canvas_shape: List = None, 2074 len_segment: int = 128, 2075 overlap: int = 0, 2076 feature_extraction: str = "kinematic", 2077 ignored_clips: List = None, 2078 ignored_bodyparts: List = None, 2079 default_agent_name: str = "ind0", 2080 key_objects: Dict = None, 2081 likelihood_threshold: float = 0, 2082 num_cpus: int = None, 2083 frame_limit: int = 1, 2084 feature_extraction_pars: Dict = None, 2085 centered: bool = False, 2086 **kwargs, 2087 ) -> None: 2088 """ 2089 Parameters 2090 ---------- 2091 video_order : list, optional 2092 a list of video ids that should be processed in the same order (not passed if creating from key objects 2093 data_path : str | set, optional 2094 the path to the folder where the pose and feature files are stored or a set of such paths 2095 (not passed if creating from key objects or from `file_paths`) 2096 file_paths : set, optional 2097 a set of string paths to the pose and feature files 2098 (not passed if creating from key objects or from `data_path`) 2099 data_suffix : str | set, optional 2100 the suffix or the set of suffices such that the pose files are named {video_id}{data_suffix} 2101 (not passed if creating from key objects or if irrelevant for the dataset) 2102 data_prefix : str | set, optional 2103 the prefix or the set of prefixes such that the pose files for different video views of the same 2104 clip are named {prefix}{sep}{video_id}{data_suffix} (not passed if creating from key objects 2105 or if irrelevant for the dataset) 2106 feature_suffix : str | set, optional 2107 the suffix or the set of suffices such that the additional feature files are named 2108 {video_id}{feature_suffix} (and placed at the data_path folder) 2109 convert_int_indices : bool, default True 2110 if `True`, convert any integer key `i` in feature files to `'ind{i}'` 2111 feature_save_path : str, optional 2112 the path to the folder where pre-processed files are stored (not passed if creating from key objects) 2113 canvas_shape : List, default [1, 1] 2114 the canvas size where the pose is defined 2115 len_segment : int, default 128 2116 the length of the segments in which the data should be cut (in frames) 2117 overlap : int, default 0 2118 the length of the overlap between neighboring segments (in frames) 2119 feature_extraction : str, default 'kinematic' 2120 the feature extraction method (see options.feature_extractors for available options) 2121 ignored_clips : list, optional 2122 list of strings of clip ids to ignore 2123 ignored_bodyparts : list, optional 2124 list of strings of bodypart names to ignore 2125 default_agent_name : str, default 'ind0' 2126 the agent name used as default in the pose files for a single agent 2127 key_objects : tuple, optional 2128 a tuple of key objects 2129 likelihood_threshold : float, default 0 2130 coordinate values with likelihoods less than this value will be set to 'unknown' 2131 num_cpus : int, optional 2132 the number of cpus to use in data processing 2133 frame_limit : int, default 1 2134 clips shorter than this number of frames will be ignored 2135 feature_extraction_pars : dict, optional 2136 parameters of the feature extractor 2137 """ 2138 2139 super().__init__( 2140 video_order, 2141 data_path, 2142 file_paths, 2143 data_suffix=data_suffix, 2144 data_prefix=data_prefix, 2145 feature_suffix=feature_suffix, 2146 convert_int_indices=convert_int_indices, 2147 feature_save_path=feature_save_path, 2148 canvas_shape=canvas_shape, 2149 len_segment=len_segment, 2150 overlap=overlap, 2151 feature_extraction=feature_extraction, 2152 ignored_clips=ignored_clips, 2153 ignored_bodyparts=ignored_bodyparts, 2154 default_agent_name=default_agent_name, 2155 key_objects=key_objects, 2156 likelihood_threshold=likelihood_threshold, 2157 num_cpus=num_cpus, 2158 frame_limit=frame_limit, 2159 feature_extraction_pars=feature_extraction_pars, 2160 centered=centered, 2161 )
Parameters
video_order : list, optional
a list of video ids that should be processed in the same order (not passed if creating from key objects
data_path : str | set, optional
the path to the folder where the pose and feature files are stored or a set of such paths
(not passed if creating from key objects or from file_paths)
file_paths : set, optional
a set of string paths to the pose and feature files
(not passed if creating from key objects or from data_path)
data_suffix : str | set, optional
the suffix or the set of suffices such that the pose files are named {video_id}{data_suffix}
(not passed if creating from key objects or if irrelevant for the dataset)
data_prefix : str | set, optional
the prefix or the set of prefixes such that the pose files for different video views of the same
clip are named {prefix}{sep}{video_id}{data_suffix} (not passed if creating from key objects
or if irrelevant for the dataset)
feature_suffix : str | set, optional
the suffix or the set of suffices such that the additional feature files are named
{video_id}{feature_suffix} (and placed at the data_path folder)
convert_int_indices : bool, default True
if True, convert any integer key i in feature files to 'ind{i}'
feature_save_path : str, optional
the path to the folder where pre-processed files are stored (not passed if creating from key objects)
canvas_shape : List, default [1, 1]
the canvas size where the pose is defined
len_segment : int, default 128
the length of the segments in which the data should be cut (in frames)
overlap : int, default 0
the length of the overlap between neighboring segments (in frames)
feature_extraction : str, default 'kinematic'
the feature extraction method (see options.feature_extractors for available options)
ignored_clips : list, optional
list of strings of clip ids to ignore
ignored_bodyparts : list, optional
list of strings of bodypart names to ignore
default_agent_name : str, default 'ind0'
the agent name used as default in the pose files for a single agent
key_objects : tuple, optional
a tuple of key objects
likelihood_threshold : float, default 0
coordinate values with likelihoods less than this value will be set to 'unknown'
num_cpus : int, optional
the number of cpus to use in data processing
frame_limit : int, default 1
clips shorter than this number of frames will be ignored
feature_extraction_pars : dict, optional
parameters of the feature extractor
Inherited Members
- GeneralInputStore
- data_suffix
- get_folder
- remove
- key_objects
- load_from_key_objects
- to_ram
- get_original_coordinates
- create_subsample
- get_video_id
- get_clip_id
- get_clip_length
- get_clip_start_end
- get_clip_start
- get_visibility
- get_annotation_objects
- get_file_ids
- get_bodyparts
- get_coords
- get_n_frames
- get_likelihood
- get_indices
- get_tag
2204class LoadedFeaturesInputStore(GeneralInputStore): 2205 """ 2206 Non-pose feature files 2207 2208 The feature files should to be dictionaries where keys are clip IDs (e.g. animal names) and values are 2209 feature values (arrays of shape `(#frames, #features)`). If the arrays are shaped as `(#features, #frames)`, 2210 set `transpose_features` to `True`. 2211 2212 The files can be saved with `numpy.save()` (with `.npy` extension), `torch.save()` (with `.pt` extension) or 2213 with `pickle.dump()` (with `.pickle` or `.pkl` extension). 2214 2215 Assumes the following file structure: 2216 ``` 2217 data_path 2218 ├── video1_features.pt 2219 └── video2_features.pt 2220 ``` 2221 Here `feature_suffix` (optional) is `'_features.pt'`. 2222 """ 2223 2224 def __init__( 2225 self, 2226 video_order: List = None, 2227 data_path: Union[Set, str] = None, 2228 file_paths: Set = None, 2229 feature_suffix: Union[Set, str] = None, 2230 convert_int_indices: bool = True, 2231 feature_save_path: str = None, 2232 len_segment: int = 128, 2233 overlap: int = 0, 2234 ignored_clips: List = None, 2235 key_objects: Dict = None, 2236 num_cpus: int = None, 2237 frame_limit: int = 1, 2238 transpose_features: bool = False, 2239 **kwargs, 2240 ) -> None: 2241 """ 2242 Parameters 2243 ---------- 2244 video_order : list, optional 2245 a list of video ids that should be processed in the same order (not passed if creating from key objects 2246 data_path : str | set, optional 2247 the path to the folder where the pose and feature files are stored or a set of such paths 2248 (not passed if creating from key objects or from `file_paths`) 2249 file_paths : set, optional 2250 a set of string paths to the pose and feature files 2251 (not passed if creating from key objects or from `data_path`) 2252 feature_suffix : str | set, optional 2253 the suffix or the set of suffices such that the additional feature files are named 2254 {video_id}{feature_suffix} (and placed at the data_path folder) 2255 feature_save_path : str, optional 2256 the path to the folder where pre-processed files are stored (not passed if creating from key objects) 2257 len_segment : int, default 128 2258 the length of the segments in which the data should be cut (in frames) 2259 overlap : int, default 0 2260 the length of the overlap between neighboring segments (in frames) 2261 ignored_clips : list, optional 2262 list of strings of clip ids to ignore 2263 default_agent_name : str, default 'ind0' 2264 the agent name used as default in the pose files for a single agent 2265 key_objects : tuple, optional 2266 a tuple of key objects 2267 num_cpus : int, optional 2268 the number of cpus to use in data processing 2269 frame_limit : int, default 1 2270 clips shorter than this number of frames will be ignored 2271 feature_extraction_pars : dict, optional 2272 parameters of the feature extractor 2273 """ 2274 2275 super().__init__( 2276 video_order, 2277 data_path, 2278 file_paths, 2279 feature_suffix=feature_suffix, 2280 convert_int_indices=convert_int_indices, 2281 feature_save_path=feature_save_path, 2282 len_segment=len_segment, 2283 overlap=overlap, 2284 ignored_clips=ignored_clips, 2285 key_objects=key_objects, 2286 num_cpus=num_cpus, 2287 frame_limit=frame_limit, 2288 transpose_features=transpose_features, 2289 ) 2290 2291 def get_visibility( 2292 self, video_id: str, clip_id: str, start: int, end: int, score: int 2293 ) -> float: 2294 """ 2295 Get the fraction of the frames in that have a visibility score better than a hard_threshold 2296 2297 For example, in the case of keypoint data the visibility score can be the number of identified keypoints. 2298 2299 Parameters 2300 ---------- 2301 video_id : str 2302 the video id of the frames 2303 clip_id : str 2304 the clip id of the frames 2305 start : int 2306 the start frame 2307 end : int 2308 the end frame 2309 score : float 2310 the visibility score hard_threshold 2311 2312 Returns 2313 ------- 2314 frac_visible: float 2315 the fraction of frames with visibility above the hard_threshold 2316 """ 2317 2318 return 1 2319 2320 def _generate_features( 2321 self, video_id: str 2322 ) -> Tuple[Dict, Dict, Dict, Union[str, int]]: 2323 """ 2324 Generate features from the raw coordinates 2325 """ 2326 2327 features = defaultdict(lambda: {}) 2328 loaded_features = self._load_saved_features(video_id) 2329 min_frames = None 2330 max_frames = None 2331 video_tag = None 2332 for clip_id, feature_tensor in loaded_features.items(): 2333 if clip_id == "max_frames": 2334 max_frames = feature_tensor 2335 elif clip_id == "min_frames": 2336 min_frames = feature_tensor 2337 elif clip_id == "video_tag": 2338 video_tag = feature_tensor 2339 else: 2340 if not isinstance(feature_tensor, torch.Tensor): 2341 feature_tensor = torch.tensor(feature_tensor) 2342 if self.convert_int_indices and ( 2343 isinstance(clip_id, int) or isinstance(clip_id, np.integer) 2344 ): 2345 clip_id = f"ind{clip_id}" 2346 key = f"{os.path.basename(video_id)}---{clip_id}" 2347 features[key]["loaded"] = feature_tensor 2348 if min_frames is None: 2349 min_frames = {} 2350 for key, value in features.items(): 2351 video_id, clip_id = key.split("---") 2352 min_frames[clip_id] = 0 2353 if max_frames is None: 2354 max_frames = {} 2355 for key, value in features.items(): 2356 video_id, clip_id = key.split("---") 2357 max_frames[clip_id] = value["loaded"].shape[0] - 1 + min_frames[clip_id] 2358 return features, min_frames, max_frames, video_tag 2359 2360 def _load_data(self) -> np.array: 2361 """ 2362 Load input data and generate data prompts 2363 """ 2364 2365 if self.video_order is None: 2366 return None 2367 2368 files = [] 2369 for video_id in self.video_order: 2370 for f in self.file_paths: 2371 if f.endswith(tuple(self.feature_suffix)): 2372 bn = os.path.basename(f) 2373 if video_id == strip_suffix(bn, self.feature_suffix): 2374 files.append(f) 2375 2376 def make_data_dictionary(filename): 2377 name = strip_suffix(filename, self.feature_suffix) 2378 name = os.path.basename(name) 2379 data_dict, min_frames, max_frames, video_tag = self._generate_features(name) 2380 bp_dict = defaultdict(lambda: {}) 2381 for key, value in data_dict.items(): 2382 video_id, clip_id = key.split("---") 2383 bp_dict[video_id][clip_id] = 1 2384 min_frames = {name: min_frames} # name is e.g. 20190707T1126-1226 2385 max_frames = {name: max_frames} 2386 names, lengths, coords = self._make_trimmed_data(data_dict) 2387 return names, lengths, coords, bp_dict, min_frames, max_frames, video_tag 2388 2389 dict_list = p_map(make_data_dictionary, files, num_cpus=self.num_cpus) 2390 # dict_list = tqdm([make_data_dictionary(f) for f in files]) 2391 2392 self.visibility = {} 2393 self.min_frames = {} 2394 self.max_frames = {} 2395 self.original_coordinates = [] 2396 self.metadata = [] 2397 X = [] 2398 for ( 2399 names, 2400 lengths, 2401 coords, 2402 bp_dictionary, 2403 min_frames, 2404 max_frames, 2405 metadata, 2406 ) in dict_list: 2407 X += names 2408 self.original_coordinates += coords 2409 self.visibility.update(bp_dictionary) 2410 self.min_frames.update(min_frames) 2411 self.max_frames.update(max_frames) 2412 if metadata is not None: 2413 self.metadata += metadata 2414 del dict_list 2415 if len(self.metadata) != len(self.original_coordinates): 2416 self.metadata = None 2417 else: 2418 self.metadata = np.array(self.metadata) 2419 2420 self.min_frames = dict(self.min_frames) 2421 self.max_frames = dict(self.max_frames) 2422 self.original_coordinates = np.array(self.original_coordinates) 2423 return np.array(X) 2424 2425 @classmethod 2426 def get_file_ids( 2427 cls, 2428 data_path: Union[Set, str] = None, 2429 file_paths: Set = None, 2430 feature_suffix: Set = None, 2431 *args, 2432 **kwargs, 2433 ) -> List: 2434 """ 2435 Process data parameters and return a list of ids of the videos that should 2436 be processed by the __init__ function 2437 2438 Parameters 2439 ---------- 2440 data_suffix : set | str, optional 2441 the suffix (or a set of suffixes) of the input data files 2442 data_path : set | str, optional 2443 the path to the folder where the pose and feature files are stored or a set of such paths 2444 (not passed if creating from key objects or from `file_paths`) 2445 data_prefix : set | str, optional 2446 the prefix or the set of prefixes such that the pose files for different video views of the same 2447 clip are named {prefix}{sep}{video_id}{data_suffix} (not passed if creating from key objects 2448 or if irrelevant for the dataset) 2449 file_paths : set, optional 2450 a set of string paths to the pose and feature files 2451 feature_suffix : str | set, optional 2452 the suffix or the set of suffices such that the additional feature files are named 2453 {video_id}{feature_suffix} (and placed at the `data_path` folder or at `file_paths`) 2454 2455 Returns 2456 ------- 2457 video_ids : list 2458 a list of video file ids 2459 """ 2460 2461 if feature_suffix is None: 2462 feature_suffix = [] 2463 if isinstance(feature_suffix, str): 2464 feature_suffix = [feature_suffix] 2465 feature_suffix = tuple(feature_suffix) 2466 if file_paths is None: 2467 file_paths = [] 2468 if data_path is not None: 2469 if isinstance(data_path, str): 2470 data_path = [data_path] 2471 file_paths = [] 2472 for folder in data_path: 2473 file_paths += [os.path.join(folder, x) for x in os.listdir(folder)] 2474 ids = set() 2475 for f in file_paths: 2476 if f.endswith(feature_suffix): 2477 bn = os.path.basename(f) 2478 video_id = strip_suffix(bn, feature_suffix) 2479 ids.add(video_id) 2480 ids = sorted(ids) 2481 return ids
Non-pose feature files
The feature files should to be dictionaries where keys are clip IDs (e.g. animal names) and values are
feature values (arrays of shape (#frames, #features)). If the arrays are shaped as (#features, #frames),
set transpose_features to True.
The files can be saved with numpy.save() (with .npy extension), torch.save() (with .pt extension) or
with pickle.dump() (with .pickle or .pkl extension).
Assumes the following file structure:
data_path
├── video1_features.pt
└── video2_features.pt
Here feature_suffix (optional) is '_features.pt'.
LoadedFeaturesInputStore( video_order: List = None, data_path: Union[Set, str] = None, file_paths: Set = None, feature_suffix: Union[Set, str] = None, convert_int_indices: bool = True, feature_save_path: str = None, len_segment: int = 128, overlap: int = 0, ignored_clips: List = None, key_objects: Dict = None, num_cpus: int = None, frame_limit: int = 1, transpose_features: bool = False, **kwargs)
2224 def __init__( 2225 self, 2226 video_order: List = None, 2227 data_path: Union[Set, str] = None, 2228 file_paths: Set = None, 2229 feature_suffix: Union[Set, str] = None, 2230 convert_int_indices: bool = True, 2231 feature_save_path: str = None, 2232 len_segment: int = 128, 2233 overlap: int = 0, 2234 ignored_clips: List = None, 2235 key_objects: Dict = None, 2236 num_cpus: int = None, 2237 frame_limit: int = 1, 2238 transpose_features: bool = False, 2239 **kwargs, 2240 ) -> None: 2241 """ 2242 Parameters 2243 ---------- 2244 video_order : list, optional 2245 a list of video ids that should be processed in the same order (not passed if creating from key objects 2246 data_path : str | set, optional 2247 the path to the folder where the pose and feature files are stored or a set of such paths 2248 (not passed if creating from key objects or from `file_paths`) 2249 file_paths : set, optional 2250 a set of string paths to the pose and feature files 2251 (not passed if creating from key objects or from `data_path`) 2252 feature_suffix : str | set, optional 2253 the suffix or the set of suffices such that the additional feature files are named 2254 {video_id}{feature_suffix} (and placed at the data_path folder) 2255 feature_save_path : str, optional 2256 the path to the folder where pre-processed files are stored (not passed if creating from key objects) 2257 len_segment : int, default 128 2258 the length of the segments in which the data should be cut (in frames) 2259 overlap : int, default 0 2260 the length of the overlap between neighboring segments (in frames) 2261 ignored_clips : list, optional 2262 list of strings of clip ids to ignore 2263 default_agent_name : str, default 'ind0' 2264 the agent name used as default in the pose files for a single agent 2265 key_objects : tuple, optional 2266 a tuple of key objects 2267 num_cpus : int, optional 2268 the number of cpus to use in data processing 2269 frame_limit : int, default 1 2270 clips shorter than this number of frames will be ignored 2271 feature_extraction_pars : dict, optional 2272 parameters of the feature extractor 2273 """ 2274 2275 super().__init__( 2276 video_order, 2277 data_path, 2278 file_paths, 2279 feature_suffix=feature_suffix, 2280 convert_int_indices=convert_int_indices, 2281 feature_save_path=feature_save_path, 2282 len_segment=len_segment, 2283 overlap=overlap, 2284 ignored_clips=ignored_clips, 2285 key_objects=key_objects, 2286 num_cpus=num_cpus, 2287 frame_limit=frame_limit, 2288 transpose_features=transpose_features, 2289 )
Parameters
video_order : list, optional
a list of video ids that should be processed in the same order (not passed if creating from key objects
data_path : str | set, optional
the path to the folder where the pose and feature files are stored or a set of such paths
(not passed if creating from key objects or from file_paths)
file_paths : set, optional
a set of string paths to the pose and feature files
(not passed if creating from key objects or from data_path)
feature_suffix : str | set, optional
the suffix or the set of suffices such that the additional feature files are named
{video_id}{feature_suffix} (and placed at the data_path folder)
feature_save_path : str, optional
the path to the folder where pre-processed files are stored (not passed if creating from key objects)
len_segment : int, default 128
the length of the segments in which the data should be cut (in frames)
overlap : int, default 0
the length of the overlap between neighboring segments (in frames)
ignored_clips : list, optional
list of strings of clip ids to ignore
default_agent_name : str, default 'ind0'
the agent name used as default in the pose files for a single agent
key_objects : tuple, optional
a tuple of key objects
num_cpus : int, optional
the number of cpus to use in data processing
frame_limit : int, default 1
clips shorter than this number of frames will be ignored
feature_extraction_pars : dict, optional
parameters of the feature extractor
def
get_visibility( self, video_id: str, clip_id: str, start: int, end: int, score: int) -> float:
2291 def get_visibility( 2292 self, video_id: str, clip_id: str, start: int, end: int, score: int 2293 ) -> float: 2294 """ 2295 Get the fraction of the frames in that have a visibility score better than a hard_threshold 2296 2297 For example, in the case of keypoint data the visibility score can be the number of identified keypoints. 2298 2299 Parameters 2300 ---------- 2301 video_id : str 2302 the video id of the frames 2303 clip_id : str 2304 the clip id of the frames 2305 start : int 2306 the start frame 2307 end : int 2308 the end frame 2309 score : float 2310 the visibility score hard_threshold 2311 2312 Returns 2313 ------- 2314 frac_visible: float 2315 the fraction of frames with visibility above the hard_threshold 2316 """ 2317 2318 return 1
Get the fraction of the frames in that have a visibility score better than a hard_threshold
For example, in the case of keypoint data the visibility score can be the number of identified keypoints.
Parameters
video_id : str the video id of the frames clip_id : str the clip id of the frames start : int the start frame end : int the end frame score : float the visibility score hard_threshold
Returns
frac_visible: float the fraction of frames with visibility above the hard_threshold
@classmethod
def
get_file_ids( cls, data_path: Union[Set, str] = None, file_paths: Set = None, feature_suffix: Set = None, *args, **kwargs) -> List:
2425 @classmethod 2426 def get_file_ids( 2427 cls, 2428 data_path: Union[Set, str] = None, 2429 file_paths: Set = None, 2430 feature_suffix: Set = None, 2431 *args, 2432 **kwargs, 2433 ) -> List: 2434 """ 2435 Process data parameters and return a list of ids of the videos that should 2436 be processed by the __init__ function 2437 2438 Parameters 2439 ---------- 2440 data_suffix : set | str, optional 2441 the suffix (or a set of suffixes) of the input data files 2442 data_path : set | str, optional 2443 the path to the folder where the pose and feature files are stored or a set of such paths 2444 (not passed if creating from key objects or from `file_paths`) 2445 data_prefix : set | str, optional 2446 the prefix or the set of prefixes such that the pose files for different video views of the same 2447 clip are named {prefix}{sep}{video_id}{data_suffix} (not passed if creating from key objects 2448 or if irrelevant for the dataset) 2449 file_paths : set, optional 2450 a set of string paths to the pose and feature files 2451 feature_suffix : str | set, optional 2452 the suffix or the set of suffices such that the additional feature files are named 2453 {video_id}{feature_suffix} (and placed at the `data_path` folder or at `file_paths`) 2454 2455 Returns 2456 ------- 2457 video_ids : list 2458 a list of video file ids 2459 """ 2460 2461 if feature_suffix is None: 2462 feature_suffix = [] 2463 if isinstance(feature_suffix, str): 2464 feature_suffix = [feature_suffix] 2465 feature_suffix = tuple(feature_suffix) 2466 if file_paths is None: 2467 file_paths = [] 2468 if data_path is not None: 2469 if isinstance(data_path, str): 2470 data_path = [data_path] 2471 file_paths = [] 2472 for folder in data_path: 2473 file_paths += [os.path.join(folder, x) for x in os.listdir(folder)] 2474 ids = set() 2475 for f in file_paths: 2476 if f.endswith(feature_suffix): 2477 bn = os.path.basename(f) 2478 video_id = strip_suffix(bn, feature_suffix) 2479 ids.add(video_id) 2480 ids = sorted(ids) 2481 return ids
Process data parameters and return a list of ids of the videos that should be processed by the __init__ function
Parameters
data_suffix : set | str, optional
the suffix (or a set of suffixes) of the input data files
data_path : set | str, optional
the path to the folder where the pose and feature files are stored or a set of such paths
(not passed if creating from key objects or from file_paths)
data_prefix : set | str, optional
the prefix or the set of prefixes such that the pose files for different video views of the same
clip are named {prefix}{sep}{video_id}{data_suffix} (not passed if creating from key objects
or if irrelevant for the dataset)
file_paths : set, optional
a set of string paths to the pose and feature files
feature_suffix : str | set, optional
the suffix or the set of suffices such that the additional feature files are named
{video_id}{feature_suffix} (and placed at the data_path folder or at file_paths)
Returns
video_ids : list a list of video file ids
Inherited Members
- GeneralInputStore
- data_suffix
- get_folder
- remove
- key_objects
- load_from_key_objects
- to_ram
- get_original_coordinates
- create_subsample
- get_video_id
- get_clip_id
- get_clip_length
- get_clip_start_end
- get_clip_start
- get_annotation_objects
- get_bodyparts
- get_coords
- get_n_frames
- get_likelihood
- get_indices
- get_tag
2484class SIMBAInputStore(FileInputStore): 2485 """ 2486 SIMBA paper format data 2487 2488 Assumes the following file structure 2489 2490 ``` 2491 data_path 2492 ├── Video1.csv 2493 ... 2494 └── Video9.csv 2495 ``` 2496 Here `data_suffix` is `.csv`. 2497 """ 2498 2499 def __init__( 2500 self, 2501 video_order: List = None, 2502 data_path: Union[Set, str] = None, 2503 file_paths: Set = None, 2504 data_prefix: Union[Set, str] = None, 2505 feature_suffix: str = None, 2506 feature_save_path: str = None, 2507 canvas_shape: List = None, 2508 len_segment: int = 128, 2509 overlap: int = 0, 2510 feature_extraction: str = "kinematic", 2511 ignored_clips: List = None, 2512 ignored_bodyparts: List = None, 2513 key_objects: Tuple = None, 2514 likelihood_threshold: float = 0, 2515 num_cpus: int = None, 2516 normalize: bool = False, 2517 feature_extraction_pars: Dict = None, 2518 centered: bool = False, 2519 data_suffix: str = None, 2520 use_features: bool = False, 2521 *args, 2522 **kwargs, 2523 ) -> None: 2524 """ 2525 Parameters 2526 ---------- 2527 video_order : list, optional 2528 a list of video ids that should be processed in the same order (not passed if creating from key objects 2529 data_path : str | set, optional 2530 the path to the folder where the pose and feature files are stored or a set of such paths 2531 (not passed if creating from key objects or from `file_paths`) 2532 file_paths : set, optional 2533 a set of string paths to the pose and feature files 2534 (not passed if creating from key objects or from `data_path`) 2535 data_suffix : str | set, optional 2536 the suffix or the set of suffices such that the pose files are named {video_id}{data_suffix} 2537 (not passed if creating from key objects or if irrelevant for the dataset) 2538 data_prefix : str | set, optional 2539 the prefix or the set of prefixes such that the pose files for different video views of the same 2540 clip are named {prefix}{sep}{video_id}{data_suffix} (not passed if creating from key objects 2541 or if irrelevant for the dataset) 2542 feature_suffix : str | set, optional 2543 the suffix or the set of suffices such that the additional feature files are named 2544 {video_id}{feature_suffix} (and placed at the data_path folder) 2545 feature_save_path : str, optional 2546 the path to the folder where pre-processed files are stored (not passed if creating from key objects) 2547 canvas_shape : List, default [1, 1] 2548 the canvas size where the pose is defined 2549 len_segment : int, default 128 2550 the length of the segments in which the data should be cut (in frames) 2551 overlap : int, default 0 2552 the length of the overlap between neighboring segments (in frames) 2553 feature_extraction : str, default 'kinematic' 2554 the feature extraction method (see options.feature_extractors for available options) 2555 ignored_clips : list, optional 2556 list of strings of clip ids to ignore 2557 ignored_bodyparts : list, optional 2558 list of strings of bodypart names to ignore 2559 key_objects : tuple, optional 2560 a tuple of key objects 2561 likelihood_threshold : float, default 0 2562 coordinate values with likelihoods less than this value will be set to 'unknown' 2563 num_cpus : int, optional 2564 the number of cpus to use in data processing 2565 feature_extraction_pars : dict, optional 2566 parameters of the feature extractor 2567 """ 2568 2569 self.use_features = use_features 2570 if feature_extraction_pars is not None: 2571 feature_extraction_pars["interactive"] = True 2572 super().__init__( 2573 video_order=video_order, 2574 data_path=data_path, 2575 file_paths=file_paths, 2576 data_suffix=data_suffix, 2577 data_prefix=data_prefix, 2578 feature_suffix=feature_suffix, 2579 convert_int_indices=False, 2580 feature_save_path=feature_save_path, 2581 canvas_shape=canvas_shape, 2582 len_segment=len_segment, 2583 overlap=overlap, 2584 feature_extraction=feature_extraction, 2585 ignored_clips=ignored_clips, 2586 ignored_bodyparts=ignored_bodyparts, 2587 default_agent_name="", 2588 key_objects=key_objects, 2589 likelihood_threshold=likelihood_threshold, 2590 num_cpus=num_cpus, 2591 min_frames=0, 2592 normalize=normalize, 2593 feature_extraction_pars=feature_extraction_pars, 2594 centered=centered, 2595 ) 2596 2597 def _open_data( 2598 self, filename: str, default_clip_name: str 2599 ) -> Tuple[Dict, Optional[Dict]]: 2600 data = pd.read_csv(filename) 2601 output = {} 2602 column_dict = {"x": "x", "y": "y", "z": "z", "p": "likelihood"} 2603 columns = [x for x in data.columns if x.split("_")[-1] in column_dict] 2604 animals = sorted(set([x.split("_")[-2] for x in columns])) 2605 coords = sorted(set([x.split("_")[-1] for x in columns])) 2606 names = sorted(set(["_".join(x.split("_")[:-2]) for x in columns])) 2607 for animal in animals: 2608 data_dict = {} 2609 for i, row in data.iterrows(): 2610 for col_name in names: 2611 data_dict[(i, col_name)] = [ 2612 row[f"{col_name}_{animal}_{coord}"] for coord in coords 2613 ] 2614 output[animal] = pd.DataFrame(data_dict).T 2615 output[animal].columns = [column_dict[x] for x in coords] 2616 if self.use_features: 2617 columns_to_avoid = [ 2618 x 2619 for x in data.columns 2620 if x.split("_")[-1] in column_dict 2621 or x.split("_")[-1].startswith("prediction") 2622 ] 2623 columns_to_avoid += ["scorer", "frames", "video_no"] 2624 output["loaded"] = ( 2625 data[[x for x in data.columns if x not in columns_to_avoid]] 2626 .interpolate() 2627 .values 2628 ) 2629 return output, None
SIMBA paper format data
Assumes the following file structure
data_path
├── Video1.csv
...
└── Video9.csv
Here data_suffix is .csv.
SIMBAInputStore( video_order: List = None, data_path: Union[Set, str] = None, file_paths: Set = None, data_prefix: Union[Set, str] = None, feature_suffix: str = None, feature_save_path: str = None, canvas_shape: List = None, len_segment: int = 128, overlap: int = 0, feature_extraction: str = 'kinematic', ignored_clips: List = None, ignored_bodyparts: List = None, key_objects: Tuple = None, likelihood_threshold: float = 0, num_cpus: int = None, normalize: bool = False, feature_extraction_pars: Dict = None, centered: bool = False, data_suffix: str = None, use_features: bool = False, *args, **kwargs)
2499 def __init__( 2500 self, 2501 video_order: List = None, 2502 data_path: Union[Set, str] = None, 2503 file_paths: Set = None, 2504 data_prefix: Union[Set, str] = None, 2505 feature_suffix: str = None, 2506 feature_save_path: str = None, 2507 canvas_shape: List = None, 2508 len_segment: int = 128, 2509 overlap: int = 0, 2510 feature_extraction: str = "kinematic", 2511 ignored_clips: List = None, 2512 ignored_bodyparts: List = None, 2513 key_objects: Tuple = None, 2514 likelihood_threshold: float = 0, 2515 num_cpus: int = None, 2516 normalize: bool = False, 2517 feature_extraction_pars: Dict = None, 2518 centered: bool = False, 2519 data_suffix: str = None, 2520 use_features: bool = False, 2521 *args, 2522 **kwargs, 2523 ) -> None: 2524 """ 2525 Parameters 2526 ---------- 2527 video_order : list, optional 2528 a list of video ids that should be processed in the same order (not passed if creating from key objects 2529 data_path : str | set, optional 2530 the path to the folder where the pose and feature files are stored or a set of such paths 2531 (not passed if creating from key objects or from `file_paths`) 2532 file_paths : set, optional 2533 a set of string paths to the pose and feature files 2534 (not passed if creating from key objects or from `data_path`) 2535 data_suffix : str | set, optional 2536 the suffix or the set of suffices such that the pose files are named {video_id}{data_suffix} 2537 (not passed if creating from key objects or if irrelevant for the dataset) 2538 data_prefix : str | set, optional 2539 the prefix or the set of prefixes such that the pose files for different video views of the same 2540 clip are named {prefix}{sep}{video_id}{data_suffix} (not passed if creating from key objects 2541 or if irrelevant for the dataset) 2542 feature_suffix : str | set, optional 2543 the suffix or the set of suffices such that the additional feature files are named 2544 {video_id}{feature_suffix} (and placed at the data_path folder) 2545 feature_save_path : str, optional 2546 the path to the folder where pre-processed files are stored (not passed if creating from key objects) 2547 canvas_shape : List, default [1, 1] 2548 the canvas size where the pose is defined 2549 len_segment : int, default 128 2550 the length of the segments in which the data should be cut (in frames) 2551 overlap : int, default 0 2552 the length of the overlap between neighboring segments (in frames) 2553 feature_extraction : str, default 'kinematic' 2554 the feature extraction method (see options.feature_extractors for available options) 2555 ignored_clips : list, optional 2556 list of strings of clip ids to ignore 2557 ignored_bodyparts : list, optional 2558 list of strings of bodypart names to ignore 2559 key_objects : tuple, optional 2560 a tuple of key objects 2561 likelihood_threshold : float, default 0 2562 coordinate values with likelihoods less than this value will be set to 'unknown' 2563 num_cpus : int, optional 2564 the number of cpus to use in data processing 2565 feature_extraction_pars : dict, optional 2566 parameters of the feature extractor 2567 """ 2568 2569 self.use_features = use_features 2570 if feature_extraction_pars is not None: 2571 feature_extraction_pars["interactive"] = True 2572 super().__init__( 2573 video_order=video_order, 2574 data_path=data_path, 2575 file_paths=file_paths, 2576 data_suffix=data_suffix, 2577 data_prefix=data_prefix, 2578 feature_suffix=feature_suffix, 2579 convert_int_indices=False, 2580 feature_save_path=feature_save_path, 2581 canvas_shape=canvas_shape, 2582 len_segment=len_segment, 2583 overlap=overlap, 2584 feature_extraction=feature_extraction, 2585 ignored_clips=ignored_clips, 2586 ignored_bodyparts=ignored_bodyparts, 2587 default_agent_name="", 2588 key_objects=key_objects, 2589 likelihood_threshold=likelihood_threshold, 2590 num_cpus=num_cpus, 2591 min_frames=0, 2592 normalize=normalize, 2593 feature_extraction_pars=feature_extraction_pars, 2594 centered=centered, 2595 )
Parameters
video_order : list, optional
a list of video ids that should be processed in the same order (not passed if creating from key objects
data_path : str | set, optional
the path to the folder where the pose and feature files are stored or a set of such paths
(not passed if creating from key objects or from file_paths)
file_paths : set, optional
a set of string paths to the pose and feature files
(not passed if creating from key objects or from data_path)
data_suffix : str | set, optional
the suffix or the set of suffices such that the pose files are named {video_id}{data_suffix}
(not passed if creating from key objects or if irrelevant for the dataset)
data_prefix : str | set, optional
the prefix or the set of prefixes such that the pose files for different video views of the same
clip are named {prefix}{sep}{video_id}{data_suffix} (not passed if creating from key objects
or if irrelevant for the dataset)
feature_suffix : str | set, optional
the suffix or the set of suffices such that the additional feature files are named
{video_id}{feature_suffix} (and placed at the data_path folder)
feature_save_path : str, optional
the path to the folder where pre-processed files are stored (not passed if creating from key objects)
canvas_shape : List, default [1, 1]
the canvas size where the pose is defined
len_segment : int, default 128
the length of the segments in which the data should be cut (in frames)
overlap : int, default 0
the length of the overlap between neighboring segments (in frames)
feature_extraction : str, default 'kinematic'
the feature extraction method (see options.feature_extractors for available options)
ignored_clips : list, optional
list of strings of clip ids to ignore
ignored_bodyparts : list, optional
list of strings of bodypart names to ignore
key_objects : tuple, optional
a tuple of key objects
likelihood_threshold : float, default 0
coordinate values with likelihoods less than this value will be set to 'unknown'
num_cpus : int, optional
the number of cpus to use in data processing
feature_extraction_pars : dict, optional
parameters of the feature extractor
Inherited Members
- GeneralInputStore
- data_suffix
- get_folder
- remove
- key_objects
- load_from_key_objects
- to_ram
- get_original_coordinates
- create_subsample
- get_video_id
- get_clip_id
- get_clip_length
- get_clip_start_end
- get_clip_start
- get_visibility
- get_annotation_objects
- get_file_ids
- get_bodyparts
- get_coords
- get_n_frames
- get_likelihood
- get_indices
- get_tag