ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-29 18:51:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fixed issues with vgg19 preprocessing. Added yaml support on config file

Browse Source
This commit is contained in:
Jaret Burkett
2023-07-19 06:12:57 -06:00
parent 8d8229dfc0
commit 6ada328d8d
6 changed files with 46 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
toolkit/config.py
View File

@@ -1,10 +1,11 @@
import os
import json
import oyaml as yaml
from collections import OrderedDict
from toolkit.paths import TOOLKIT_ROOT
possible_extensions = ['.json', '.jsonc']
possible_extensions = ['.json', '.jsonc', '.yaml', '.yml']
def get_cwd_abs_path(path):
@@ -49,8 +50,14 @@ def get_config(config_file_path):
if not real_config_path:
raise ValueError(f"Could not find config file {config_file_path}")
# load the config
with open(real_config_path, 'r') as f:
config = json.load(f, object_pairs_hook=OrderedDict)
# if we found it, check if it is a json or yaml file
if real_config_path.endswith('.json') or real_config_path.endswith('.jsonc'):
with open(real_config_path, 'r') as f:
config = json.load(f, object_pairs_hook=OrderedDict)
elif real_config_path.endswith('.yaml') or real_config_path.endswith('.yml'):
with open(real_config_path, 'r') as f:
config = yaml.load(f, Loader=yaml.FullLoader)
else:
raise ValueError(f"Config file {config_file_path} must be a json or yaml file")
return preprocess_config(config)