DeepLabCut/deeplabcut/pose_estimation_tensorflow/datasets/pose_base.py at main · DeepLabCut/DeepLabCut · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
#
# DeepLabCut Toolbox (deeplabcut.org)
# © A. & M.W. Mathis Labs
# https://github.com/DeepLabCut/DeepLabCut
#
# Please see AUTHORS for contributors.
# https://github.com/DeepLabCut/DeepLabCut/blob/master/AUTHORS
#
# Licensed under GNU Lesser General Public License v3.0
#


import abc

import numpy as np


class BasePoseDataset(metaclass=abc.ABCMeta):
    # TODO Finish implementing actual abstract class
    def __init__(self, cfg):
        self.cfg = cfg

    @abc.abstractmethod
    def load_dataset(self): ...

    @abc.abstractmethod
    def next_batch(self): ...

    def sample_scale(self):
        if self.cfg.get("deterministic", False):
            np.random.seed(42)
        scale = self.cfg["global_scale"]
        if "scale_jitter_lo" in self.cfg and "scale_jitter_up" in self.cfg:
            scale_jitter = np.random.uniform(self.cfg["scale_jitter_lo"], self.cfg["scale_jitter_up"])
            scale *= scale_jitter
        return scale