diff --git a/element_deeplabcut/model.py b/element_deeplabcut/model.py index 2e032ea..45ecdb3 100644 --- a/element_deeplabcut/model.py +++ b/element_deeplabcut/model.py @@ -7,7 +7,7 @@ import os import cv2 import csv -import yaml +import ruamel.yaml as yaml import inspect import importlib import numpy as np @@ -240,7 +240,7 @@ def extract_new_body_parts(cls, dlc_config: dict, verbose: bool = True): ), f"dlc_config is neither dict nor filepath\n Check: {dlc_config_fp}" if dlc_config_fp.suffix in (".yml", ".yaml"): with open(dlc_config_fp, "rb") as f: - dlc_config = yaml.safe_load(f) + dlc_config = yaml.load(f) # -- Check and insert new BodyPart -- assert "bodyparts" in dlc_config, f"Found no bodyparts section in {dlc_config}" tracked_body_parts = cls.fetch("body_part") @@ -382,7 +382,7 @@ def insert_new_model( ) if dlc_config_fp.suffix in (".yml", ".yaml"): with open(dlc_config_fp, "rb") as f: - dlc_config = yaml.safe_load(f) + dlc_config = yaml.load(f) if isinstance(params, dict): dlc_config.update(params)