MCPcopy Index your code
hub / github.com/DeepLabCut/DeepLabCut / load_model

Function load_model

deeplabcut/pose_estimation_tensorflow/export.py:78–158  ·  view source on GitHub ↗

Loads a tensorflow session with a DLC model from the associated configuration Return a tensorflow session with DLC model given cfg and shuffle. Parameters: ----------- cfg : dict Configuration read from the project's main config.yaml file shuffle : int, optional

(cfg, shuffle=1, trainingsetindex=0, TFGPUinference=True, modelprefix="")

Source from the content-addressed store, hash-verified

76
77
78def load_model(cfg, shuffle=1, trainingsetindex=0, TFGPUinference=True, modelprefix=""):
79 """Loads a tensorflow session with a DLC model from the associated configuration
80 Return a tensorflow session with DLC model given cfg and shuffle.
81
82 Parameters:
83 -----------
84 cfg : dict
85 Configuration read from the project's main config.yaml file
86
87 shuffle : int, optional
88 which shuffle to use
89
90 trainingsetindex : int. optional
91 which training fraction to use, identified by its index
92
93 TFGPUinference : bool, optional
94 use tensorflow inference model? default = True
95
96 Returns:
97 --------
98 sess : tensorflow session
99 tensorflow session with DLC model from the provided configuration, shuffle, and trainingsetindex
100
101 checkpoint file path : string
102 the path to the checkpoint file associated with the loaded model
103 """
104
105 ########################
106 ### find snapshot to use
107 ########################
108
109 train_fraction = cfg["TrainingFraction"][trainingsetindex]
110 model_folder = os.path.join(
111 cfg["project_path"],
112 str(auxiliaryfunctions.get_model_folder(train_fraction, shuffle, cfg, modelprefix=modelprefix)),
113 )
114 os.path.normpath(model_folder + "/test/pose_cfg.yaml")
115 path_train_config = os.path.normpath(model_folder + "/train/pose_cfg.yaml")
116
117 try:
118 dlc_cfg = load_config(str(path_train_config))
119 # dlc_cfg_train = load_config(str(path_train_config))
120 except FileNotFoundError as e:
121 raise FileNotFoundError(
122 f"It seems the model for shuffle {shuffle} and trainFraction {train_fraction} does not exist."
123 ) from e
124
125 Snapshots = auxiliaryfunctions.get_snapshots_from_folder(
126 train_folder=Path(model_folder) / "train",
127 )
128
129 if cfg["snapshotindex"] == "all":
130 print("Snapshotindex is set to 'all' in the config.yaml file. Changing snapshot index to -1!")
131 snapshotindex = -1
132 else:
133 snapshotindex = cfg["snapshotindex"]
134
135 ####################################

Callers 1

export_modelFunction · 0.85

Calls 3

load_configFunction · 0.90
splitMethod · 0.45
getMethod · 0.45

Tested by

no test coverage detected