dlc2action.data.dataset

   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"""
   7Behavior dataset (class that manages high-level data interactions)
   8"""
   9
  10import warnings
  11from typing import Dict, Union, Tuple, List, Optional, Any
  12
  13from numpy import ndarray
  14from torch.utils.data import Dataset
  15import torch
  16from abc import ABC
  17import numpy as np
  18from copy import copy
  19from tqdm import tqdm
  20from collections import Counter
  21import inspect
  22from collections import defaultdict
  23from dlc2action.utils import (
  24    apply_threshold_hysteresis,
  25    apply_threshold,
  26    apply_threshold_max,
  27)
  28from dlc2action.data.base_store import InputStore, AnnotationStore
  29from copy import deepcopy, copy
  30import os
  31import pickle
  32from dlc2action import options
  33
  34
  35class BehaviorDataset(Dataset, ABC):
  36    """
  37    A generalized dataset class
  38
  39    Data and annotation are stored in separate InputStore and AnnotationStore objects; the dataset class
  40    manages their interactions.
  41    """
  42
  43    def __init__(
  44        self,
  45        data_type: str,
  46        annotation_type: str = "none",
  47        ssl_transformations: List = None,
  48        saved_data_path: str = None,
  49        input_store: InputStore = None,
  50        annotation_store: AnnotationStore = None,
  51        only_load_annotated: bool = False,
  52        recompute_annotation: bool = False,
  53        # mask: str = None,
  54        ids: List = None,
  55        **data_parameters,
  56    ) -> None:
  57        """
  58        Parameters
  59        ----------
  60        data_type : str
  61            the data type (see available types by running BehaviorDataset.data_types())
  62        annotation_type : str
  63            the annotation type (see available types by running BehaviorDataset.annotation_types())
  64        ssl_transformations : list
  65            a list of functions that take a sample dictionary as input and return an (ssl input, ssl target) tuple
  66        saved_data_path : str
  67            the path to a pre-computed pickled dataset
  68        input_store : InputStore
  69            a pre-computed input store
  70        annotation_store : AnnotationStore
  71            a precomputed annotation store
  72        only_load_annotated : bool
  73            if True, the input files that don't have a matching annotation file will be disregarded
  74        *data_parameters : dict
  75            parameters to initialize the input and annotation stores
  76        """
  77
  78        mask = None
  79        if len(data_parameters) == 0:
  80            recompute_annotation = False
  81        feature_extraction = data_parameters.get("feature_extraction")
  82        if feature_extraction is not None and not issubclass(
  83            options.input_stores[data_type],
  84            options.feature_extractors[feature_extraction].input_store_class,
  85        ):
  86            raise ValueError(
  87                f"The {feature_extraction} feature extractor does not work with "
  88                f"the {data_type} data type, please choose a suclass of "
  89                f"{options.feature_extractors[feature_extraction].input_store_class}"
  90            )
  91        if ssl_transformations is None:
  92            ssl_transformations = []
  93        self.ssl_transformations = ssl_transformations
  94        self.input_type = data_type
  95        self.annotation_type = annotation_type
  96        self.stats = None
  97        if mask is not None:
  98            with open(mask, "rb") as f:
  99                self.mask = pickle.load(f)
 100        else:
 101            self.mask = None
 102        self.ids = ids
 103        self.tag = None
 104        self.return_unlabeled = None
 105        # load saved key objects for annotation and input if they exist
 106        input_key_objects, annotation_key_objects = None, None
 107        if saved_data_path is not None:
 108            if os.path.exists(saved_data_path):
 109                with open(saved_data_path, "rb") as f:
 110                    input_key_objects, annotation_key_objects = pickle.load(f)
 111        # if the input or the annotation store need to be created, generate the common video order
 112        if len(data_parameters) > 0:
 113            input_files = options.input_stores[data_type].get_file_ids(
 114                **data_parameters
 115            )
 116            annotation_files = options.annotation_stores[annotation_type].get_file_ids(
 117                **data_parameters
 118            )
 119            if only_load_annotated:
 120                data_parameters["video_order"] = [
 121                    x for x in input_files if x in annotation_files
 122                ]
 123            else:
 124                data_parameters["video_order"] = input_files
 125            if len(data_parameters["video_order"]) == 0:
 126                raise RuntimeError(
 127                    "The length of file list is 0! Please check your data parameters!"
 128                )
 129        data_parameters["mask"] = self.mask
 130        # load or create the input store
 131        ok = False
 132        if input_store is not None:
 133            self.input_store = input_store
 134            ok = True
 135        elif input_key_objects is not None:
 136            try:
 137                self.input_store = self._load_input_store(data_type, input_key_objects)
 138                ok = True
 139            except:
 140                warnings.warn("Loading input store from key objects failed")
 141        if not ok:
 142            self.input_store = self._get_input_store(
 143                data_type, deepcopy(data_parameters)
 144            )
 145        # get the objects needed to create the annotation store (like a clip length dictionary)
 146        annotation_objects = self.input_store.get_annotation_objects()
 147        data_parameters.update(annotation_objects)
 148        # load or create the annotation store
 149        ok = False
 150        if annotation_store is not None:
 151            self.annotation_store = annotation_store
 152            ok = True
 153        elif (
 154            annotation_key_objects is not None
 155            and mask is None
 156            and not recompute_annotation
 157        ):
 158            try:
 159                self.annotation_store = self._load_annotation_store(
 160                    annotation_type, annotation_key_objects
 161                )
 162                ok = True
 163            except:
 164                warnings.warn("Loading annotation store from key objects failed")
 165        if not ok:
 166            self.annotation_store = self._get_annotation_store(
 167                annotation_type, deepcopy(data_parameters)
 168            )
 169        if (
 170            mask is None
 171            and annotation_type != "none"
 172            and not recompute_annotation
 173            and (
 174                self.annotation_store.get_original_coordinates()
 175                != self.input_store.get_original_coordinates()
 176            ).any()
 177        ):
 178            raise RuntimeError(
 179                "The clip orders in the annotation store and input store are different!"
 180            )
 181        # filter the data based on data parameters
 182        # print(f"1 {self.annotation_store.get_original_coordinates().shape=}")
 183        # print(f"1 {self.input_store.get_original_coordinates().shape=}")
 184        to_remove = self.annotation_store.filtered_indices()
 185        if len(to_remove) > 0:
 186            print(
 187                f"Filtering {100 * len(to_remove) / len(self.annotation_store):.2f}% of samples"
 188            )
 189        if len(self.input_store) == len(self.annotation_store):
 190            self.input_store.remove(to_remove)
 191        self.annotation_store.remove(to_remove)
 192        self.input_indices = list(range(len(self.input_store)))
 193        self.annotation_indices = list(range(len(self.input_store)))
 194        self.indices = list(range(len(self.input_store)))
 195        # print(f'{data_parameters["video_order"]=}')
 196        # print(f"{self.annotation_store.get_original_coordinates().shape=}")
 197        # print(f"{self.input_store.get_original_coordinates().shape=}")
 198        # count = 0
 199        # for i, (x, y) in enumerate(zip(
 200        #     self.annotation_store.get_original_coordinates(),
 201        #     self.input_store.get_original_coordinates(),
 202        # )):
 203        #     if (x != y).any():
 204        #         count += 1
 205        #         print({i})
 206        #         print(f"ann: {x}")
 207        #         print(f"inp: {y}")
 208        #         print("\n")
 209        #     if count > 50:
 210        #         break
 211        if annotation_type != "none" and (
 212            self.annotation_store.get_original_coordinates().shape
 213            != self.input_store.get_original_coordinates().shape
 214            or (
 215                self.annotation_store.get_original_coordinates()
 216                != self.input_store.get_original_coordinates()
 217            ).any()
 218        ):
 219            raise RuntimeError(
 220                "The clip orders in the annotation store and input store are different!"
 221            )
 222
 223    def __getitem__(self, item: int) -> Dict:
 224        idx = self._get_idx(item)
 225        input = deepcopy(self.input_store[idx])
 226        target = self.annotation_store[idx]
 227        tag = self.input_store.get_tag(idx)
 228        ssl_inputs, ssl_targets = self._get_SSL_targets(input)
 229        batch = {"input": input}
 230        for name, x in zip(
 231            ["target", "ssl_inputs", "ssl_targets", "tag"],
 232            [target, ssl_inputs, ssl_targets, tag],
 233        ):
 234            if x is not None:
 235                batch[name] = x
 236        batch["index"] = idx
 237        if self.stats is not None:
 238            for key in batch["input"].keys():
 239                key_name = key.split("---")[0]
 240                if key_name in self.stats:
 241                    batch["input"][key][:, batch["input"][key].sum(0) != 0] = (
 242                        (batch["input"][key] - self.stats[key_name]["mean"])
 243                        / (self.stats[key_name]["std"] + 1e-7)
 244                    )[:, batch["input"][key].sum(0) != 0]
 245        return batch
 246
 247    def __len__(self) -> int:
 248        return len(self.indices)
 249        # if self.annotation_type != "none":
 250        #     return self.annotation_store.get_len(return_unlabeled=self.return_unlabeled)
 251        # else:
 252        #     return len(self.input_store)
 253
 254    def get_tags(self) -> List:
 255        """
 256        Get a list of all meta tags
 257
 258        Returns
 259        -------
 260        tags: List
 261            a list of unique meta tag values
 262        """
 263
 264        return self.input_store.get_tags()
 265
 266    def save(self, save_path: str) -> None:
 267        """
 268        Save the dictionary
 269
 270        Parameters
 271        ----------
 272        save_path : str
 273            the path where the pickled file will be stored
 274        """
 275
 276        input_obj = self.input_store.key_objects()
 277        annotation_obj = self.annotation_store.key_objects()
 278        with open(save_path, "wb") as f:
 279            pickle.dump((input_obj, annotation_obj), f)
 280
 281    def to_ram(self) -> None:
 282        """
 283        Transfer the dataset to RAM
 284        """
 285
 286        self.input_store.to_ram()
 287        self.annotation_store.to_ram()
 288
 289    def generate_full_length_gt(self) -> Dict:
 290        if self.annotation_class() == "exclusive_classification":
 291            gt = torch.zeros((len(self), self.len_segment()))
 292        else:
 293            gt = torch.zeros(
 294                (len(self), len(self.behaviors_dict()), self.len_segment())
 295            )
 296        for i in range(len(self)):
 297            gt[i] = self.annotation_store[i]
 298        return self.generate_full_length_prediction(gt)
 299
 300    def generate_full_length_prediction(self, predicted: torch.Tensor) -> Dict:
 301        """
 302        Map predictions for the equal-length pieces to predictions for the original data
 303
 304        Probabilities are averaged over predictions on overlapping intervals.
 305
 306        Parameters
 307        ----------
 308        predicted: torch.Tensor
 309            a tensor of predicted probabilities of shape `(N, #classes, #frames)`
 310
 311        Returns
 312        -------
 313        full_length_prediction : dict
 314            a nested dictionary where first-level keys are video ids, second-level keys are clip ids and values are
 315            averaged probability tensors
 316        """
 317
 318        result = defaultdict(lambda: {})
 319        counter = defaultdict(lambda: {})
 320        coordinates = self.input_store.get_original_coordinates()
 321        for coords, prediction in zip(coordinates, predicted):
 322            l = self.input_store.get_clip_length_from_coords(coords)
 323            video_name = self.input_store.get_video_id(coords)
 324            clip_id = self.input_store.get_clip_id(coords)
 325            start, end = self.input_store.get_clip_start_end(coords)
 326            if clip_id not in result[video_name].keys():
 327                result[video_name][clip_id] = torch.zeros(*prediction.shape[:-1], l)
 328                counter[video_name][clip_id] = torch.zeros(*prediction.shape[:-1], l)
 329            result[video_name][clip_id][..., start:end] += (
 330                prediction.squeeze()[..., : end - start].detach().cpu()
 331            )
 332            counter[video_name][clip_id][..., start:end] += 1
 333        for video_name in result:
 334            for clip_id in result[video_name]:
 335                result[video_name][clip_id] /= counter[video_name][clip_id]
 336                result[video_name][clip_id][counter[video_name][clip_id] == 0] = -100
 337        result = dict(result)
 338        return result
 339
 340    def find_valleys(
 341        self,
 342        predicted: Union[torch.Tensor, Dict],
 343        threshold: float = 0.5,
 344        min_frames: int = 0,
 345        visibility_min_score: float = 0,
 346        visibility_min_frac: float = 0,
 347        main_class: int = 1,
 348        low: bool = True,
 349        predicted_error: torch.Tensor = None,
 350        error_threshold: float = 0.5,
 351        hysteresis: bool = False,
 352        threshold_diff: float = None,
 353        min_frames_error: int = None,
 354        smooth_interval: int = 1,
 355        cut_annotated: bool = False,
 356    ) -> Dict:
 357        """
 358        Find the intervals where the probability of a certain class is below or above a certain hard_threshold
 359
 360        Parameters
 361        ----------
 362        predicted : torch.Tensor | dict
 363            either a tensor of predictions for the data prompts or the output of
 364            `BehaviorDataset.generate_full_length_prediction`
 365        threshold : float, default 0.5
 366            the main hard_threshold
 367        min_frames : int, default 0
 368            the minimum length of the intervals
 369        visibility_min_score : float, default 0
 370            the minimum visibility score in the intervals
 371        visibility_min_frac : float, default 0
 372            fraction of the interval that has to have the visibility score larger than visibility_score_thr
 373        main_class : int, default 1
 374            the index of the class the function is inspecting
 375        low : bool, default True
 376            if True, the probability in the intervals has to be below the hard_threshold, and if False, it has to be above
 377        predicted_error : torch.Tensor, optional
 378            a tensor of error predictions for the data prompts
 379        error_threshold : float, default 0.5
 380            maximum possible probability of error at the intervals
 381        hysteresis: bool, default False
 382            if True, the function will apply a hysteresis hard_threshold with the soft hard_threshold defined by threshold_diff
 383        threshold_diff: float, optional
 384            the difference between the soft and hard hard_threshold if hysteresis is used; if hysteresis is True, low is False and threshold_diff is None, the soft hard_threshold condition is set to the main_class having a larger probability than other classes
 385        min_frames_error: int, optional
 386            if not None, the intervals will only be considered where the error probability is below error_threshold at at least min_frames_error consecutive frames
 387
 388        Returns
 389        -------
 390        valleys : dict
 391            a dictionary where keys are video ids and values are lists of (start, end, individual name) tuples that denote the chosen intervals
 392        """
 393
 394        result = defaultdict(lambda: [])
 395        if type(predicted) is not dict:
 396            predicted = self.generate_full_length_prediction(predicted)
 397        if predicted_error is not None:
 398            predicted_error = self.generate_full_length_prediction(predicted_error)
 399        elif min_frames_error is not None and min_frames_error != 0:
 400            # warnings.warn(
 401            #     f"The min_frames_error parameter is set to {min_frames_error} but no error prediction "
 402            #     f"is given! Setting min_frames_error to 0."
 403            # )
 404            min_frames_error = 0
 405        if low and hysteresis and threshold_diff is None:
 406            raise ValueError(
 407                "Cannot set low=True, hysteresis=True and threshold_diff=None! Please set threshold_diff."
 408            )
 409        if cut_annotated:
 410            masked_intervals_dict = self.get_annotated_intervals()
 411        else:
 412            masked_intervals_dict = None
 413        print("Valleys found:")
 414        for v_id in predicted:
 415            for clip_id in predicted[v_id].keys():
 416                if predicted_error is not None:
 417                    error_mask = predicted_error[v_id][clip_id][1, :] < error_threshold
 418                    if min_frames_error is not None:
 419                        output, indices, counts = torch.unique_consecutive(
 420                            error_mask, return_inverse=True, return_counts=True
 421                        )
 422                        wrong_indices = torch.where(
 423                            output * (counts < min_frames_error)
 424                        )[0]
 425                        if len(wrong_indices) > 0:
 426                            for i in wrong_indices:
 427                                error_mask[indices == i] = False
 428                else:
 429                    error_mask = None
 430                if masked_intervals_dict is not None:
 431                    masked_intervals = masked_intervals_dict[v_id][clip_id]
 432                else:
 433                    masked_intervals = None
 434                if not hysteresis:
 435                    res_indices_start, res_indices_end = apply_threshold(
 436                        predicted[v_id][clip_id][main_class, :],
 437                        threshold,
 438                        low,
 439                        error_mask,
 440                        min_frames,
 441                        smooth_interval,
 442                        masked_intervals,
 443                    )
 444                elif threshold_diff is not None:
 445                    if low:
 446                        soft_threshold = threshold + threshold_diff
 447                    else:
 448                        soft_threshold = threshold - threshold_diff
 449                    res_indices_start, res_indices_end = apply_threshold_hysteresis(
 450                        predicted[v_id][clip_id][main_class, :],
 451                        soft_threshold,
 452                        threshold,
 453                        low,
 454                        error_mask,
 455                        min_frames,
 456                        smooth_interval,
 457                        masked_intervals,
 458                    )
 459                else:
 460                    res_indices_start, res_indices_end = apply_threshold_max(
 461                        predicted[v_id][clip_id],
 462                        threshold,
 463                        main_class,
 464                        error_mask,
 465                        min_frames,
 466                        smooth_interval,
 467                        masked_intervals,
 468                    )
 469                start = self.input_store.get_clip_start(v_id, clip_id)
 470                result[v_id] += [
 471                    [i + start, j + start, clip_id]
 472                    for i, j in zip(res_indices_start, res_indices_end)
 473                    if self.input_store.get_visibility(
 474                        v_id, clip_id, i, j, visibility_min_score
 475                    )
 476                    > visibility_min_frac
 477                ]
 478            result[v_id] = sorted(result[v_id])
 479            print(f"    {v_id}: {len(result[v_id])}")
 480        return dict(result)
 481
 482    def valleys_union(self, valleys_list) -> Dict:
 483        """
 484        Find the intersection of two valleys dictionaries
 485
 486        Parameters
 487        ----------
 488        valleys_list : list
 489            a list of valleys dictionaries
 490
 491        Returns
 492        -------
 493        intersection : dict
 494            a new valleys dictionary with the intersection of the input intervals
 495        """
 496
 497        valleys_list = [x for x in valleys_list if x is not None]
 498        if len(valleys_list) == 1:
 499            return valleys_list[0]
 500        elif len(valleys_list) == 0:
 501            return {}
 502        union = {}
 503        keys_list = [set(valleys.keys()) for valleys in valleys_list]
 504        keys = set.union(*keys_list)
 505        for v_id in keys:
 506            res = []
 507            clips_list = [
 508                set([x[-1] for x in valleys[v_id]]) for valleys in valleys_list
 509            ]
 510            clips = set.union(*clips_list)
 511            for clip_id in clips:
 512                clip_intervals = [
 513                    x
 514                    for valleys in valleys_list
 515                    for x in valleys[v_id]
 516                    if x[-1] == clip_id
 517                ]
 518                v_len = self.input_store.get_clip_length(v_id, clip_id)
 519                arr = torch.zeros(v_len)
 520                for start, end, _ in clip_intervals:
 521                    arr[start:end] += 1
 522                output, indices, counts = torch.unique_consecutive(
 523                    arr > 0, return_inverse=True, return_counts=True
 524                )
 525                long_indices = torch.where(output)[0]
 526                res += [
 527                    (
 528                        (indices == i).nonzero(as_tuple=True)[0][0].item(),
 529                        (indices == i).nonzero(as_tuple=True)[0][-1].item(),
 530                        clip_id,
 531                    )
 532                    for i in long_indices
 533                ]
 534            union[v_id] = res
 535        return union
 536
 537    def valleys_intersection(self, valleys_list) -> Dict:
 538        """
 539        Find the intersection of two valleys dictionaries
 540
 541        Parameters
 542        ----------
 543        valleys_list : list
 544            a list of valleys dictionaries
 545
 546        Returns
 547        -------
 548        intersection : dict
 549            a new valleys dictionary with the intersection of the input intervals
 550        """
 551
 552        valleys_list = [x for x in valleys_list if x is not None]
 553        if len(valleys_list) == 1:
 554            return valleys_list[0]
 555        elif len(valleys_list) == 0:
 556            return {}
 557        intersection = {}
 558        keys_list = [set(valleys.keys()) for valleys in valleys_list]
 559        keys = set.intersection(*keys_list)
 560        for v_id in keys:
 561            res = []
 562            clips_list = [
 563                set([x[-1] for x in valleys[v_id]]) for valleys in valleys_list
 564            ]
 565            clips = set.intersection(*clips_list)
 566            for clip_id in clips:
 567                clip_intervals = [
 568                    x
 569                    for valleys in valleys_list
 570                    for x in valleys[v_id]
 571                    if x[-1] == clip_id
 572                ]
 573                v_len = self.input_store.get_clip_length(v_id, clip_id)
 574                arr = torch.zeros(v_len)
 575                for start, end, _ in clip_intervals:
 576                    arr[start:end] += 1
 577                output, indices, counts = torch.unique_consecutive(
 578                    arr, return_inverse=True, return_counts=True
 579                )
 580                long_indices = torch.where(output == 2)[0]
 581                res += [
 582                    (
 583                        (indices == i).nonzero(as_tuple=True)[0][0].item(),
 584                        (indices == i).nonzero(as_tuple=True)[0][-1].item(),
 585                        clip_id,
 586                    )
 587                    for i in long_indices
 588                ]
 589            intersection[v_id] = res
 590        return intersection
 591
 592    def partition_train_test_val(
 593        self,
 594        use_test: float = 0,
 595        split_path: str = None,
 596        method: str = "random",
 597        val_frac: float = 0,
 598        test_frac: float = 0,
 599        save_split: bool = False,
 600        normalize: bool = False,
 601        skip_normalization_keys: List = None,
 602        stats: Dict = None,
 603    ) -> Tuple:
 604        """
 605        Partition the dataset into three new datasets
 606
 607        Parameters
 608        ----------
 609        use_test : float, default 0
 610            The fraction of the test dataset to be used in training without labels
 611        split_path : str, optional
 612            The path to load the split information from (if `'file'` method is used) and to save it to
 613            (if `'save_split'` is `True`)
 614        method : {'random', 'random:test-from-name', 'random:test-from-name:{name}',
 615            'val-from-name:{val_name}:test-from-name:{test_name}',
 616            'random:equalize:segments', 'random:equalize:videos',
 617            'folders', 'time', 'time:strict', 'file'}
 618            The partitioning method:
 619            - `'random'`: sort videos into subsets randomly,
 620            - `'random:test-from-name'` (or `'random:test-from-name:{name}'`): sort videos into training and validation
 621                subsets randomly and create
 622                the test subset from the video ids that start with a speific substring (`'test'` by default, or `name`
 623                if provided),
 624            - `'random:equalize:segments'` and `'random:equalize:videos'`: sort videos into subsets randomly but
 625                making sure that for the rarest classes at least `0.8 * val_frac` of the videos/segments that contain
 626                occurences of the class get into the validation subset and `0.8 * test_frac` get into the test subset;
 627                this in ensured for all classes in order of increasing number of occurences until the validation and
 628                test subsets are full
 629            - `'val-from-name:{val_name}:test-from-name:{test_name}'`: create the validation and test
 630                subsets from the video ids that start with specific substrings (`val_name` for validation
 631                and `test_name` for test) and sort all other videos into the training subset
 632            - `'folders'`: read videos from folders named *test*, *train* and *val* into corresponding subsets,
 633            - `'time'`: split each video into training, validation and test subsequences,
 634            - `'time:strict'`: split each video into validation, test and training subsequences
 635                and throw out the last segments in validation and test (to get rid of overlaps),
 636            - `'file'`: split according to a split file.
 637        val_frac : float, default 0
 638            The fraction of the dataset to be used in validation
 639        test_frac : float, default 0
 640            The fraction of the dataset to be used in test
 641        save_split : bool, default False
 642            Save a split file if True
 643
 644        Returns
 645        -------
 646        train_dataset : BehaviorDataset
 647            train dataset
 648
 649        val_dataset : BehaviorDataset
 650            validation dataset
 651
 652        test_dataset : BehaviorDataset
 653            test dataset
 654        """
 655
 656        train_indices, test_indices, val_indices = self._partition_indices(
 657            split_path=split_path,
 658            method=method,
 659            val_frac=val_frac,
 660            test_frac=test_frac,
 661            save_split=save_split,
 662        )
 663        val_dataset = self._create_new_dataset(val_indices)
 664        test_dataset = self._create_new_dataset(test_indices)
 665        train_dataset = self._create_new_dataset(
 666            train_indices, ssl_indices=test_indices[: int(len(test_indices) * use_test)]
 667        )
 668        print("Number of samples:")
 669        print(f"    validation:")
 670        print(f"      {val_dataset.count_classes()}")
 671        print(f"    training:")
 672        print(f"      {train_dataset.count_classes()}")
 673        print(f"    test:")
 674        print(f"      {test_dataset.count_classes()}")
 675        if normalize:
 676            if stats is None:
 677                print("Computing normalization statistics...")
 678                stats = train_dataset.get_normalization_stats(skip_normalization_keys)
 679            else:
 680                print("Setting loaded normalization statistics...")
 681            train_dataset.set_normalization_stats(stats)
 682            val_dataset.set_normalization_stats(stats)
 683            test_dataset.set_normalization_stats(stats)
 684        return train_dataset, test_dataset, val_dataset
 685
 686    def class_weights(self, proportional=False) -> List:
 687        """
 688        Calculate class weights in inverse proportion to number of samples
 689        Returns
 690        -------
 691        weights: list
 692            a list of class weights
 693        """
 694
 695        items = sorted(
 696            [
 697                (k, v)
 698                for k, v in self.annotation_store.count_classes().items()
 699                if k != -100
 700            ]
 701        )
 702        if self.annotation_store.annotation_class() == "exclusive_classification":
 703            if not proportional:
 704                numerator = len(self.annotation_store)
 705            else:
 706                numerator = max([x[1] for x in items])
 707            weights = [numerator / (v + 1e-7) for _, v in items]
 708        else:
 709            items_zero = sorted(
 710                [
 711                    (k, v)
 712                    for k, v in self.annotation_store.count_classes(zeros=True).items()
 713                    if k != -100
 714                ]
 715            )
 716            if not proportional:
 717                numerators = defaultdict(lambda: len(self.annotation_store))
 718            else:
 719                numerators = {
 720                    item_one[0]: max(item_one[1], item_zero[1])
 721                    for item_one, item_zero in zip(items, items_zero)
 722                }
 723            weights = {}
 724            weights[0] = [numerators[k] / (v + 1e-7) for k, v in items_zero]
 725            weights[1] = [numerators[k] / (v + 1e-7) for k, v in items]
 726        return weights
 727
 728    def boundary_class_weight(self):
 729        if self.annotation_type != "none":
 730            f = self.annotation_store.data.flatten()
 731            _, inv = torch.unique_consecutive(f, return_inverse=True)
 732            boundary = torch.cat([torch.tensor([0]), torch.diff(inv)]).reshape(
 733                self.annotation_store.data.shape
 734            )
 735            boundary[..., 0] = 0
 736            cnt = Counter(boundary.flatten().numpy())
 737            return cnt[1] / cnt[0]
 738        else:
 739            return 0
 740
 741    def count_classes(self, bouts: bool = False) -> Dict:
 742        """
 743        Get a class counter dictionary
 744
 745        Parameters
 746        ----------
 747        bouts : bool, default False
 748            if `True`, instead of frame counts segment counts are returned
 749
 750        Returns
 751        -------
 752        count_dictionary : dict
 753            a dictionary with class indices as keys and frame or bout counts as values
 754        """
 755
 756        return self.annotation_store.count_classes(bouts=bouts)
 757
 758    def behaviors_dict(self) -> Dict:
 759        """
 760        Get a behavior dictionary
 761
 762        Returns
 763        -------
 764        dict
 765            behavior dictionary
 766        """
 767
 768        return self.annotation_store.behaviors_dict()
 769
 770    def bodyparts_order(self) -> List:
 771        try:
 772            return self.input_store.get_bodyparts()
 773        except:
 774            raise RuntimeError(
 775                f"The {self.input_type} input store does not have bodyparts implemented!"
 776            )
 777
 778    def features_shape(self) -> Dict:
 779        """
 780        Get the shapes of the input features
 781
 782        Returns
 783        -------
 784        shapes : Dict
 785            a dictionary with the shapes of the features
 786        """
 787
 788        sample = self.input_store[0]
 789        shapes = {k: v.shape for k, v in sample.items()}
 790        # for key, value in shapes.items():
 791        #     print(f'{key}: {value}')
 792        return shapes
 793
 794    def num_classes(self) -> int:
 795        """
 796        Get the number of classes in the data
 797
 798        Returns
 799        -------
 800        num_classes : int
 801            the number of classes
 802        """
 803
 804        return len(self.annotation_store.behaviors_dict())
 805
 806    def len_segment(self) -> int:
 807        """
 808        Get the segment length in the data
 809
 810        Returns
 811        -------
 812        len_segment : int
 813            the segment length
 814        """
 815
 816        sample = self.input_store[0]
 817        key = list(sample.keys())[0]
 818        return sample[key].shape[-1]
 819
 820    def set_ssl_transformations(self, ssl_transformations: List) -> None:
 821        """
 822        Set new SSL transformations
 823
 824        Parameters
 825        ----------
 826        ssl_transformations : list
 827            a list of functions that take a sample feature dictionary as input and output ssl_inputs and ssl_targets
 828            lists
 829        """
 830
 831        self.ssl_transformations = ssl_transformations
 832
 833    @classmethod
 834    def new(cls, *args, **kwargs):
 835        """
 836        Create a new object of the same class
 837
 838        Returns
 839        -------
 840        new_instance: BehaviorDataset
 841            a new instance of the same class
 842        """
 843
 844        return cls(*args, **kwargs)
 845
 846    @classmethod
 847    def get_parameters(cls, data_type: str, annotation_type: str) -> List:
 848        """
 849        Get parameters necessary for initialization
 850
 851        Parameters
 852        ----------
 853        data_type : str
 854            the data type
 855        annotation_type : str
 856            the annotation type
 857        """
 858
 859        input_features = options.input_stores[data_type].get_parameters()
 860        annotation_features = options.annotation_stores[
 861            annotation_type
 862        ].get_parameters()
 863        self_features = inspect.getfullargspec(cls.__init__).args
 864        return self_features + input_features + annotation_features
 865
 866    @staticmethod
 867    def data_types() -> List:
 868        """
 869        List available data types
 870
 871        Returns
 872        -------
 873        data_types : list
 874            available data types
 875        """
 876
 877        return list(options.input_stores.keys())
 878
 879    @staticmethod
 880    def annotation_types() -> List:
 881        """
 882        List available annotation types
 883
 884        Returns
 885        -------
 886        annotation_types : list
 887            available annotation types
 888        """
 889
 890        return list(options.annotation_stores.keys())
 891
 892    def _get_SSL_targets(self, input: Dict) -> Tuple[List, List]:
 893        """
 894        Get the SSL inputs and targets from a sample dictionary
 895        """
 896
 897        ssl_inputs = []
 898        ssl_targets = []
 899        for transform in self.ssl_transformations:
 900            ssl_input, ssl_target = transform(copy(input))
 901            ssl_inputs.append(ssl_input)
 902            ssl_targets.append(ssl_target)
 903        return ssl_inputs, ssl_targets
 904
 905    def _create_new_dataset(self, indices: List, ssl_indices: List = None):
 906        """
 907        Create a subsample of the dataset, with samples at ssl_indices losing the annotation
 908        """
 909
 910        if ssl_indices is None:
 911            ssl_indices = []
 912        input_store = self.input_store.create_subsample(indices, ssl_indices)
 913        annotation_store = self.annotation_store.create_subsample(indices, ssl_indices)
 914        new = self.new(
 915            data_type=self.input_type,
 916            annotation_type=self.annotation_type,
 917            ssl_transformations=self.ssl_transformations,
 918            annotation_store=annotation_store,
 919            input_store=input_store,
 920            ids=list(indices) + list(ssl_indices),
 921            recompute_annotation=False,
 922        )
 923        return new
 924
 925    def _load_input_store(self, data_type: str, key_objects: Tuple) -> InputStore:
 926        """
 927        Load input store from key objects
 928        """
 929
 930        input_store = options.input_stores[data_type](key_objects=key_objects)
 931        return input_store
 932
 933    def _load_annotation_store(
 934        self, annotation_type: str, key_objects: Tuple
 935    ) -> AnnotationStore:
 936        """
 937        Load annotation store from key objects
 938        """
 939
 940        annotation_store = options.annotation_stores[annotation_type](
 941            key_objects=key_objects
 942        )
 943        return annotation_store
 944
 945    def _get_input_store(self, data_type: str, data_parameters: Dict) -> InputStore:
 946        """
 947        Create input store from parameters
 948        """
 949
 950        data_parameters["key_objects"] = None
 951        input_store = options.input_stores[data_type](**data_parameters)
 952        return input_store
 953
 954    def _get_annotation_store(
 955        self, annotation_type: str, data_parameters: Dict
 956    ) -> AnnotationStore:
 957        """
 958        Create annotation store from parameters
 959        """
 960
 961        annotation_store = options.annotation_stores[annotation_type](**data_parameters)
 962        return annotation_store
 963
 964    def set_indexing_parameters(self, unlabeled: bool, tag: int) -> None:
 965        """
 966        Set the parameters that change the subset that is returned at `__getitem__`
 967
 968        Parameters
 969        ----------
 970        unlabeled : bool
 971            a pseudolabeling parameter; return only unlabeled samples if `True`, only labeled if `False` and
 972            all if `None`
 973        tag : int
 974            if not `None`, only samples with this meta tag will be returned
 975        """
 976
 977        if unlabeled != self.return_unlabeled:
 978            self.annotation_indices = self.annotation_store.get_indices(unlabeled)
 979            self.return_unlabeled = unlabeled
 980        if tag != self.tag:
 981            self.input_indices = self.input_store.get_indices(tag)
 982            self.tag = tag
 983        self.indices = [x for x in self.annotation_indices if x in self.input_indices]
 984
 985    def _get_idx(self, index: int) -> int:
 986        """
 987        Get index in full dataset
 988        """
 989
 990        return self.indices[index]
 991
 992        # return self.annotation_store.get_idx(
 993        #     index, return_unlabeled=self.return_unlabeled
 994        # )
 995
 996    def _partition_indices(
 997        self,
 998        split_path: str = None,
 999        method: str = "random",
1000        val_frac: float = 0,
1001        test_frac: float = 0,
1002        save_split: bool = False,
1003    ) -> Tuple[List, List, List]:
1004        """
1005        Partition indices into train, validation, test subsets
1006        """
1007
1008        if self.mask is not None:
1009            val_indices = self.mask["val_ids"]
1010            train_indices = [x for x in range(len(self)) if x not in val_indices]
1011            test_indices = []
1012        elif method == "random":
1013            videos = np.array(self.input_store.get_video_id_order())
1014            all_videos = list(set(videos))
1015            if len(all_videos) == 1:
1016                warnings.warn(
1017                    "There is only one video in the dataset, so train/val/test split is done on segments; "
1018                    'that might lead to overlaps, please consider using "time" or "time:strict" as the '
1019                    "partitioning method instead"
1020                )
1021                # Quick fix for single video: the problem with this is that the segments can overlap
1022                # length = int(self.input_store.get_original_coordinates()[-1][1])    # number of segments
1023                length = len(self.input_store.get_original_coordinates())
1024                val_len = int(val_frac * length)
1025                test_len = int(test_frac * length)
1026                all_indices = np.random.choice(np.arange(length), length, replace=False)
1027                val_indices = all_indices[:val_len]
1028                test_indices = all_indices[val_len : val_len + test_len]
1029                train_indices = all_indices[val_len + test_len :]
1030                coords = self.input_store.get_original_coordinates()
1031                if save_split:
1032                    self._save_partition(
1033                        coords[train_indices],
1034                        coords[val_indices],
1035                        coords[test_indices],
1036                        split_path,
1037                        coords=True,
1038                    )
1039            else:
1040                length = len(all_videos)
1041                val_len = int(val_frac * length)
1042                test_len = int(test_frac * length)
1043                validation = all_videos[:val_len]
1044                test = all_videos[val_len : val_len + test_len]
1045                training = all_videos[val_len + test_len :]
1046                train_indices = np.where(np.isin(videos, training))[0]
1047                val_indices = np.where(np.isin(videos, validation))[0]
1048                test_indices = np.where(np.isin(videos, test))[0]
1049                if save_split:
1050                    self._save_partition(training, validation, test, split_path)
1051        elif method.startswith("random:equalize"):
1052            counter = self.count_classes()
1053            counter = sorted(list([(v, k) for k, v in counter.items()]))
1054            classes = [x[1] for x in counter]
1055            indicator = {c: [] for c in classes}
1056            if method.endswith("videos"):
1057                videos = np.array(self.input_store.get_video_id_order())
1058                all_videos = list(set(videos))
1059                total_len = len(all_videos)
1060                for video_id in all_videos:
1061                    video_coords = np.where(videos == video_id)[0]
1062                    ann = torch.cat(
1063                        [self.annotation_store[i] for i in video_coords], dim=-1
1064                    )
1065                    for c in classes:
1066                        if self.annotation_class() == "nonexclusive_classification":
1067                            indicator[c].append(torch.sum(ann[c] == 1) > 0)
1068                        elif self.annotation_class() == "exclusive_classification":
1069                            indicator[c].append(torch.sum(ann == c) > 0)
1070                        else:
1071                            raise ValueError(
1072                                f"The random:equalize partition method is not implemented"
1073                                f"for the {self.annotation_class()} annotation class"
1074                            )
1075            elif method.endswith("segments"):
1076                total_len = len(self)
1077                for ann in self.annotation_store:
1078                    for c in classes:
1079                        if self.annotation_class() == "nonexclusive_classification":
1080                            indicator[c].append(torch.sum(ann[c] == 1) > 0)
1081                        elif self.annotation_class() == "exclusive_classification":
1082                            indicator[c].append(torch.sum(ann == c) > 0)
1083                        else:
1084                            raise ValueError(
1085                                f"The random:equalize partition method is not implemented"
1086                                f"for the {self.annotation_class()} annotation class"
1087                            )
1088            else:
1089                values = []
1090                for v in options.partition_methods.values():
1091                    values += v
1092                raise ValueError(
1093                    f"The {method} partition method is not recognized; please choose from {values}"
1094                )
1095            val_indices = []
1096            test_indices = []
1097            for c in classes:
1098                indicator[c] = np.array(indicator[c])
1099                ind = np.where(indicator[c])[0]
1100                np.random.shuffle(ind)
1101                c_sum = len(ind)
1102                in_val = np.sum(indicator[c][val_indices])
1103                in_test = np.sum(indicator[c][test_indices])
1104                while (
1105                    len(val_indices) < val_frac * total_len
1106                    and in_val < val_frac * c_sum * 0.8
1107                ):
1108                    first, ind = ind[0], ind[1:]
1109                    val_indices = list(set(val_indices).union({first}))
1110                    in_val = np.sum(indicator[c][val_indices])
1111                while (
1112                    len(test_indices) < test_frac * total_len
1113                    and in_test < test_frac * c_sum * 0.8
1114                ):
1115                    first, ind = ind[0], ind[1:]
1116                    test_indices = list(set(test_indices).union({first}))
1117                    in_test = np.sum(indicator[c][test_indices])
1118            if len(val_indices) < int(val_frac * total_len):
1119                left_val = int(val_frac * total_len) - len(val_indices)
1120            else:
1121                left_val = 0
1122            if len(test_indices) < int(test_frac * total_len):
1123                left_test = int(test_frac * total_len) - len(test_indices)
1124            else:
1125                left_test = 0
1126            indicator = np.ones(total_len)
1127            indicator[val_indices] = 0
1128            indicator[test_indices] = 0
1129            ind = np.where(indicator)[0]
1130            np.random.shuffle(ind)
1131            val_indices += list(ind[:left_val])
1132            test_indices += list(ind[left_val : left_val + left_test])
1133            train_indices = list(ind[left_val + left_test :])
1134            if save_split:
1135                if method.endswith("segments"):
1136                    coords = self.input_store.get_original_coordinates()
1137                    self._save_partition(
1138                        coords[train_indices],
1139                        coords[val_indices],
1140                        coords[test_indices],
1141                        coords[split_path],
1142                        coords=True,
1143                    )
1144                else:
1145                    all_videos = np.array(all_videos)
1146                    validation = all_videos[val_indices]
1147                    test = all_videos[test_indices]
1148                    training = all_videos[train_indices]
1149                    self._save_partition(training, validation, test, split_path)
1150        elif method.startswith("random:test-from-name"):
1151            split = method.split(":")
1152            if len(split) > 2:
1153                test_name = split[-1]
1154            else:
1155                test_name = "test"
1156            videos = np.array(self.input_store.get_video_id_order())
1157            all_videos = list(set(videos))
1158            test = []
1159            train_videos = []
1160            for x in all_videos:
1161                if x.startswith(test_name):
1162                    test.append(x)
1163                else:
1164                    train_videos.append(x)
1165            length = len(train_videos)
1166            val_len = int(val_frac * length)
1167            validation = train_videos[:val_len]
1168            training = train_videos[val_len:]
1169            train_indices = np.where(np.isin(videos, training))[0]
1170            val_indices = np.where(np.isin(videos, validation))[0]
1171            test_indices = np.where(np.isin(videos, test))[0]
1172            if save_split:
1173                self._save_partition(training, validation, test, split_path)
1174        elif method.startswith("val-from-name"):
1175            split = method.split(":")
1176            if split[2] != "test-from-name":
1177                raise ValueError(
1178                    f"The {method} partition method is not recognized, please choose from {options.partition_methods}"
1179                )
1180            val_name = split[1]
1181            test_name = split[-1]
1182            videos = np.array(self.input_store.get_video_id_order())
1183            all_videos = list(set(videos))
1184            test = []
1185            validation = []
1186            training = []
1187            for x in all_videos:
1188                if x.startswith(test_name):
1189                    test.append(x)
1190                elif x.startswith(val_name):
1191                    validation.append(x)
1192                else:
1193                    training.append(x)
1194            train_indices = np.where(np.isin(videos, training))[0]
1195            val_indices = np.where(np.isin(videos, validation))[0]
1196            test_indices = np.where(np.isin(videos, test))[0]
1197        elif method == "folders":
1198            folders = np.array(self.input_store.get_folder_order())
1199            videos = np.array(self.input_store.get_video_id_order())
1200            train_indices = np.where(np.isin(folders, ["training", "train"]))[0]
1201            if np.sum(np.isin(folders, ["validation", "val"])) > 0:
1202                val_indices = np.where(np.isin(folders, ["validation", "val"]))[0]
1203            else:
1204                train_videos = list(set(videos[train_indices]))
1205                val_len = int(val_frac * len(train_videos))
1206                validation = train_videos[:val_len]
1207                training = train_videos[val_len:]
1208                train_indices = np.where(np.isin(videos, training))[0]
1209                val_indices = np.where(np.isin(videos, validation))[0]
1210            test_indices = np.where(folders == "test")[0]
1211            if save_split:
1212                self._save_partition(
1213                    list(set(videos[train_indices])),
1214                    list(set(videos[val_indices])),
1215                    list(set(videos[test_indices])),
1216                    split_path,
1217                )
1218        elif method.startswith("leave-one-out"):
1219            n = int(method.split(":")[-1])
1220            videos = np.array(self.input_store.get_video_id_order())
1221            all_videos = sorted(list(set(videos)))
1222            validation = [all_videos.pop(n)]
1223            training = all_videos
1224            train_indices = np.where(np.isin(videos, training))[0]
1225            val_indices = np.where(np.isin(videos, validation))[0]
1226            test_indices = np.array([])
1227        elif method.startswith("time"):
1228            if method.endswith("strict"):
1229                len_segment = self.len_segment()
1230                # TODO: change this
1231                step = self.input_store.step
1232                num_removed = len_segment // step
1233            else:
1234                num_removed = 0
1235            videos = np.array(self.input_store.get_video_id_order())
1236            all_videos = set(videos)
1237            train_indices = []
1238            val_indices = []
1239            test_indices = []
1240            start = 0
1241            if len(method.split(":")) > 1 and method.split(":")[1] == "start-from":
1242                start = float(method.split(":")[2])
1243            for video_id in all_videos:
1244                video_indices = np.where(videos == video_id)[0]
1245                val_len = int(val_frac * len(video_indices))
1246                test_len = int(test_frac * len(video_indices))
1247                start_pos = int(start * len(video_indices))
1248                all_ind = np.ones(len(video_indices))
1249                val_indices += list(video_indices[start_pos : start_pos + val_len])
1250                all_ind[start_pos : start_pos + val_len] = 0
1251                if start_pos + val_len > len(video_indices):
1252                    p = start_pos + val_len - len(video_indices)
1253                    val_indices += list(video_indices[:p])
1254                    all_ind[:p] = 0
1255                else:
1256                    p = start_pos + val_len
1257                test_indices += list(video_indices[p : p + test_len])
1258                all_ind[p : p + test_len] = 0
1259                if p + test_len > len(video_indices):
1260                    p = test_len + p - len(video_indices)
1261                    test_indices += list(video_indices[:p])
1262                    all_ind[:p] = 0
1263                train_indices += list(video_indices[all_ind > 0])
1264                for _ in range(num_removed):
1265                    if len(val_indices) > 0:
1266                        val_indices.pop(-1)
1267                    if len(test_indices) > 0:
1268                        test_indices.pop(-1)
1269                    if start > 0 and len(train_indices) > 0:
1270                        train_indices.pop(-1)
1271        elif method == "file":
1272            if split_path is None:
1273                raise ValueError(
1274                    'You need to either set split_path or change partition method ("file" requires a file)'
1275                )
1276            with open(split_path) as f:
1277                train_line = f.readline()
1278                line = f.readline()
1279                while not line.startswith("Validation") and not line.startswith(
1280                    "Validataion"
1281                ):
1282                    line = f.readline()
1283                if line.startswith("Validation"):
1284                    validation = []
1285                    test = []
1286                    while True:
1287                        line = f.readline()
1288                        if line.startswith("Test") or len(line) == 0:
1289                            break
1290                        validation.append(line.rstrip())
1291                    while True:
1292                        line = f.readline()
1293                        if len(line) == 0:
1294                            break
1295                        test.append(line.rstrip())
1296                    type = train_line[9:-2]
1297                else:
1298                    line = f.readline()
1299                    validation = line.rstrip(",\n ").split(", ")
1300                    test = []
1301                    type = "videos"
1302            if type == "videos":
1303                videos = np.array(self.input_store.get_video_id_order())
1304                val_indices = np.where(np.isin(videos, validation))[0]
1305                test_indices = np.where(np.isin(videos, test))[0]
1306            elif type == "coords":
1307                coords = self.input_store.get_original_coordinates()
1308                video_ids = self.input_store.get_video_id_order()
1309                clip_ids = [self.input_store.get_clip_id(coord) for coord in coords]
1310                starts, ends = zip(
1311                    *[self.input_store.get_clip_start_end(coord) for coord in coords]
1312                )
1313                coords = np.array(
1314                    [
1315                        f"{video_id}---{clip_id}---{start}-{end}"
1316                        for video_id, clip_id, start, end in zip(
1317                            video_ids, clip_ids, starts, ends
1318                        )
1319                    ]
1320                )
1321                val_indices = np.where(np.isin(coords, validation))[0]
1322                test_indices = np.where(np.isin(coords, test))[0]
1323            else:
1324                raise ValueError("The split path has unrecognized format!")
1325            all = np.ones(len(self))
1326            all[val_indices] = 0
1327            all[test_indices] = 0
1328            train_indices = np.where(all)[0]
1329        else:
1330            raise ValueError(
1331                f"The {method} partition is not recognized, please choose from {options.partition_methods}"
1332            )
1333        return sorted(train_indices), sorted(test_indices), sorted(val_indices)
1334
1335    def _save_partition(
1336        self,
1337        training: List,
1338        validation: List,
1339        test: List,
1340        split_path: str,
1341        coords: bool = False,
1342    ) -> None:
1343        """
1344        Save a split file
1345        """
1346
1347        if coords:
1348            name = "coords"
1349            training_coords = []
1350            val_coords = []
1351            test_coords = []
1352            for coord in training:
1353                video_id = self.input_store.get_video_id(coord)
1354                clip_id = self.input_store.get_clip_id(coord)
1355                start, end = self.input_store.get_clip_start_end(coord)
1356                training_coords.append(f"{video_id}---{clip_id}---{start}-{end}")
1357            for coord in validation:
1358                video_id = self.input_store.get_video_id(coord)
1359                clip_id = self.input_store.get_clip_id(coord)
1360                start, end = self.input_store.get_clip_start_end(coord)
1361                val_coords.append(f"{video_id}---{clip_id}---{start}-{end}")
1362            for coord in test:
1363                video_id = self.input_store.get_video_id(coord)
1364                clip_id = self.input_store.get_clip_id(coord)
1365                start, end = self.input_store.get_clip_start_end(coord)
1366                test_coords.append(f"{video_id}---{clip_id}---{start}-{end}")
1367            training, validation, test = training_coords, val_coords, test_coords
1368        else:
1369            name = "videos"
1370        if split_path is not None:
1371            with open(split_path, "w") as f:
1372                f.write(f"Training {name}:\n")
1373                for x in training:
1374                    f.write(x + "\n")
1375                f.write(f"Validation {name}:\n")
1376                for x in validation:
1377                    f.write(x + "\n")
1378                f.write(f"Test {name}:\n")
1379                for x in test:
1380                    f.write(x + "\n")
1381
1382    def _get_intervals_from_ind(self, frame_indices: np.ndarray):
1383        """
1384        Get a list of intervals from a list of frame indices
1385
1386        Example: `[0, 1, 2, 5, 6, 8] -> [[0, 3], [5, 7], [8, 9]]`.
1387
1388        Parameters
1389        ----------
1390        frame_indices : np.ndarray
1391            a list of frame indices
1392
1393        Returns
1394        -------
1395        intervals : list
1396            a list of interval boundaries
1397        """
1398
1399        masked_intervals = []
1400        breaks = np.where(np.diff(frame_indices) != 1)[0]
1401        if len(frame_indices) > 0:
1402            start = frame_indices[0]
1403            for k in breaks:
1404                masked_intervals.append([start, frame_indices[k] + 1])
1405                start = frame_indices[k + 1]
1406            masked_intervals.append([start, frame_indices[-1] + 1])
1407        return masked_intervals
1408
1409    def get_intervals(self) -> Tuple[dict, Optional[list]]:
1410        """
1411        Get a list of intervals covered by the dataset in the original coordinates
1412
1413        Returns
1414        -------
1415        intervals : dict
1416            a nested dictionary where first-level keys are video ids, second-level keys are clip ids and
1417            values are lists of the intervals in `[start, end]` format
1418        """
1419
1420        counter = defaultdict(lambda: {})
1421        coordinates = self.input_store.get_original_coordinates()
1422        for coords in coordinates:
1423            l = self.input_store.get_clip_length_from_coords(coords)
1424            video_name = self.input_store.get_video_id(coords)
1425            clip_id = self.input_store.get_clip_id(coords)
1426            start, end = self.input_store.get_clip_start_end(coords)
1427            if clip_id not in counter[video_name]:
1428                counter[video_name][clip_id] = np.zeros(l)
1429            counter[video_name][clip_id][start:end] = 1
1430        result = {video_name: {} for video_name in counter}
1431        for video_name in counter:
1432            for clip_id in counter[video_name]:
1433                result[video_name][clip_id] = self._get_intervals_from_ind(
1434                    np.where(counter[video_name][clip_id])[0]
1435                )
1436        return result, self.ids
1437
1438    def get_unannotated_intervals(self, first_intervals=None) -> Dict:
1439        """
1440        Get a list of intervals in the original coordinates where there is no annotation
1441
1442        Returns
1443        -------
1444        intervals : dict
1445            a nested dictionary where first-level keys are video ids, second-level keys are clip ids and
1446            values are lists of the intervals in `[start, end]` format
1447        """
1448
1449        counter_value = 2
1450        if first_intervals is None:
1451            first_intervals = defaultdict(lambda: defaultdict(lambda: []))
1452            counter_value = 1
1453        counter = defaultdict(lambda: {})
1454        coordinates = self.input_store.get_original_coordinates()
1455        for i, coords in enumerate(coordinates):
1456            l = self.input_store.get_clip_length_from_coords(coords)
1457            ann = self.annotation_store[i]
1458            if (
1459                self.annotation_store.annotation_class()
1460                == "nonexclusive_classification"
1461            ):
1462                ann = ann[0, :]
1463            video_name = self.input_store.get_video_id(coords)
1464            clip_id = self.input_store.get_clip_id(coords)
1465            start, end = self.input_store.get_clip_start_end(coords)
1466            if clip_id not in counter[video_name]:
1467                counter[video_name][clip_id] = np.ones(l)
1468            counter[video_name][clip_id][start:end] = (ann[: end - start] == -100).int()
1469        result = {video_name: {} for video_name in counter}
1470        for video_name in counter:
1471            for clip_id in counter[video_name]:
1472                for start, end in first_intervals[video_name][clip_id]:
1473                    counter[video_name][clip_id][start:end] += 1
1474                result[video_name][clip_id] = self._get_intervals_from_ind(
1475                    np.where(counter[video_name][clip_id] == counter_value)[0]
1476                )
1477        return result
1478
1479    def get_annotated_intervals(self) -> Dict:
1480        """
1481        Get a list of intervals in the original coordinates where there is no annotation
1482
1483        Returns
1484        -------
1485        intervals : dict
1486            a nested dictionary where first-level keys are video ids, second-level keys are clip ids and
1487            values are lists of the intervals in `[start, end]` format
1488        """
1489
1490        if self.annotation_type == "none":
1491            return []
1492        counter_value = 1
1493        counter = defaultdict(lambda: {})
1494        coordinates = self.input_store.get_original_coordinates()
1495        for i, coords in enumerate(coordinates):
1496            l = self.input_store.get_clip_length_from_coords(coords)
1497            ann = self.annotation_store[i]
1498            video_name = self.input_store.get_video_id(coords)
1499            clip_id = self.input_store.get_clip_id(coords)
1500            start, end = self.input_store.get_clip_start_end(coords)
1501            if clip_id not in counter[video_name]:
1502                counter[video_name][clip_id] = np.zeros(l)
1503            if (
1504                self.annotation_store.annotation_class()
1505                == "nonexclusive_classification"
1506            ):
1507                counter[video_name][clip_id][start:end] = (
1508                    torch.sum(ann[:, : end - start] != -100, dim=0) > 0
1509                ).int()
1510            else:
1511                counter[video_name][clip_id][start:end] = (
1512                    ann[: end - start] != -100
1513                ).int()
1514        result = {video_name: {} for video_name in counter}
1515        for video_name in counter:
1516            for clip_id in counter[video_name]:
1517                result[video_name][clip_id] = self._get_intervals_from_ind(
1518                    np.where(counter[video_name][clip_id] == counter_value)[0]
1519                )
1520        return result
1521
1522    def get_ids(self) -> Dict:
1523        """
1524        Get a dictionary of all clip ids in the dataset
1525
1526        Returns
1527        -------
1528        ids : dict
1529            a dictionary where keys are video ids and values are lists of clip ids
1530        """
1531
1532        coordinates = self.input_store.get_original_coordinates()
1533        video_ids = np.array(self.input_store.get_video_id_order())
1534        id_set = set(video_ids)
1535        result = {}
1536        for video_id in id_set:
1537            coords = coordinates[video_ids == video_id]
1538            clip_ids = list({self.input_store.get_clip_id(c) for c in coords})
1539            result[video_id] = clip_ids
1540        return result
1541
1542    def get_len(self, video_id: str, clip_id: str) -> int:
1543        """
1544        Get the length of a specific clip
1545
1546        Parameters
1547        ----------
1548        video_id : str
1549            the video id
1550        clip_id : str
1551            the clip id
1552
1553        Returns
1554        -------
1555        length : int
1556            the length
1557        """
1558
1559        return self.input_store.get_clip_length(video_id, clip_id)
1560
1561    def get_confusion_matrix(
1562        self, prediction: torch.Tensor, confusion_type: str = "recall"
1563    ) -> Tuple[ndarray, list]:
1564        """
1565        Get a confusion matrix
1566
1567        Parameters
1568        ----------
1569        prediction : torch.Tensor
1570            a tensor of predicted class probabilities of shape `(#samples, #classes, #frames)`
1571        confusion_type : {"recall", "precision"}
1572            for datasets with non-exclusive annotation, if `type` is `"recall"`, only false positives are taken
1573            into account, and if `type` is `"precision"`, only false negatives
1574
1575        Returns
1576        -------
1577        confusion_matrix : np.ndarray
1578            a confusion matrix of shape `(#classes, #classes)` where `A[i, j] = F_ij/N_i`, `F_ij` is the number of
1579            frames that have the i-th label in the ground truth and a false positive j-th label in the prediction,
1580            `N_i` is the number of frames that have the i-th label in the ground truth
1581        classes : list
1582            a list of classes
1583        """
1584
1585        behaviors_dict = self.annotation_store.behaviors_dict()
1586        num_behaviors = len(behaviors_dict)
1587        confusion_matrix = np.zeros((num_behaviors, num_behaviors))
1588        if self.annotation_store.annotation_class() == "exclusive_classification":
1589            exclusive = True
1590            confusion_type = None
1591        elif self.annotation_store.annotation_class() == "nonexclusive_classification":
1592            exclusive = False
1593        else:
1594            raise RuntimeError(
1595                f"The {self.annotation_store.annotation_class()} annotation class is not recognized!"
1596            )
1597        for ann, p in zip(self.annotation_store, prediction):
1598            if exclusive:
1599                class_prediction = torch.max(p, dim=0)[1]
1600                for i in behaviors_dict.keys():
1601                    for j in behaviors_dict.keys():
1602                        confusion_matrix[i, j] += int(
1603                            torch.sum(class_prediction[ann == i] == j)
1604                        )
1605            else:
1606                class_prediction = (p > 0.5).int()
1607                for i in behaviors_dict.keys():
1608                    for j in behaviors_dict.keys():
1609                        if confusion_type == "recall":
1610                            pred = deepcopy(class_prediction[j])
1611                            if i != j:
1612                                pred[ann[j] == 1] = 0
1613                            confusion_matrix[i, j] += int(torch.sum(pred[ann[i] == 1]))
1614                        elif confusion_type == "precision":
1615                            annotation = deepcopy(ann[j])
1616                            if i != j:
1617                                annotation[class_prediction[j] == 1] = 0
1618                            confusion_matrix[i, j] += int(
1619                                torch.sum(annotation[class_prediction[i] == 1])
1620                            )
1621                        else:
1622                            raise ValueError(
1623                                f"The {confusion_type} type is not recognized; please choose from ['recall', 'precision']"
1624                            )
1625        counter = self.annotation_store.count_classes()
1626        for i in behaviors_dict.keys():
1627            if counter[i] != 0:
1628                if confusion_type == "recall" or confusion_type is None:
1629                    confusion_matrix[i, :] /= counter[i]
1630                else:
1631                    confusion_matrix[:, i] /= counter[i]
1632        return confusion_matrix, list(behaviors_dict.values()), confusion_type
1633
1634    def annotation_class(self) -> str:
1635        """
1636        Get the type of annotation ('exclusive_classification', 'nonexclusive_classification', more coming soon)
1637
1638        Returns
1639        -------
1640        annotation_class : str
1641            the type of annotation
1642        """
1643
1644        return self.annotation_store.annotation_class()
1645
1646    def set_normalization_stats(self, stats: Dict) -> None:
1647        """
1648        Set the stats to normalize data at runtime
1649
1650        Parameters
1651        ----------
1652        stats : dict
1653            a nested dictionary where first-level keys are feature key names, second-level keys are 'mean' and 'std'
1654            and values are the statistics in `torch` tensors of shape `(#features, 1)`
1655        """
1656
1657        self.stats = stats
1658
1659    def get_min_max_frames(self, video_id) -> Tuple[Dict, Dict]:
1660        coords = self.input_store.get_original_coordinates()
1661        clips = set(
1662            [
1663                self.input_store.get_clip_id(c)
1664                for c in coords
1665                if self.input_store.get_video_id(c) == video_id
1666            ]
1667        )
1668        min_frames = {}
1669        max_frames = {}
1670        for clip in clips:
1671            start = self.input_store.get_clip_start(video_id, clip)
1672            end = start + self.input_store.get_clip_length(video_id, clip)
1673            min_frames[clip] = start
1674            max_frames[clip] = end - 1
1675        return min_frames, max_frames
1676
1677    def get_normalization_stats(self, skip_keys=None) -> Dict:
1678        """
1679        Get mean and standard deviation for each key
1680
1681        Returns
1682        -------
1683        stats : dict
1684            a nested dictionary where first-level keys are feature key names, second-level keys are 'mean' and 'std'
1685            and values are the statistics in `torch` tensors of shape `(#features, 1)`
1686        """
1687
1688        stats = defaultdict(lambda: {})
1689        sums = defaultdict(lambda: 0)
1690        if skip_keys is None:
1691            skip_keys = []
1692        counter = defaultdict(lambda: 0)
1693        for sample in tqdm(self):
1694            for key, value in sample["input"].items():
1695                key_name = key.split("---")[0]
1696                if key_name not in skip_keys:
1697                    sums[key_name] += value[:, value.sum(0) != 0].sum(-1)
1698                counter[key_name] += torch.sum(value.sum(0) != 0)
1699        for key, value in sums.items():
1700            stats[key]["mean"] = (value / counter[key]).unsqueeze(-1)
1701        sums = defaultdict(lambda: 0)
1702        for sample in tqdm(self):
1703            for key, value in sample["input"].items():
1704                key_name = key.split("---")[0]
1705                if key_name not in skip_keys:
1706                    sums[key_name] += (
1707                        (value[:, value.sum(0) != 0] - stats[key_name]["mean"]) ** 2
1708                    ).sum(-1)
1709        for key, value in sums.items():
1710            stats[key]["std"] = np.sqrt(value.unsqueeze(-1) / counter[key])
1711        return stats