File size: 1,975 Bytes
bbde80b |
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 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
import numpy as np
import os, sys
import pickle
import yaml
from easydict import EasyDict as edict
from typing import Any, IO
ROOT_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), '..')
class TextLogger:
def __init__(self, log_path):
self.log_path = log_path
with open(self.log_path, "w") as f:
f.write("")
def log(self, log):
with open(self.log_path, "a+") as f:
f.write(log + "\n")
class Loader(yaml.SafeLoader):
"""YAML Loader with `!include` constructor."""
def __init__(self, stream: IO) -> None:
"""Initialise Loader."""
try:
self._root = os.path.split(stream.name)[0]
except AttributeError:
self._root = os.path.curdir
super().__init__(stream)
def construct_include(loader: Loader, node: yaml.Node) -> Any:
"""Include file referenced at node."""
filename = os.path.abspath(os.path.join(loader._root, loader.construct_scalar(node)))
extension = os.path.splitext(filename)[1].lstrip('.')
with open(filename, 'r') as f:
if extension in ('yaml', 'yml'):
return yaml.load(f, Loader)
elif extension in ('json', ):
return json.load(f)
else:
return ''.join(f.readlines())
def get_config(config_path):
yaml.add_constructor('!include', construct_include, Loader)
with open(config_path, 'r') as stream:
config = yaml.load(stream, Loader=Loader)
config = edict(config)
_, config_filename = os.path.split(config_path)
config_name, _ = os.path.splitext(config_filename)
config.name = config_name
return config
def ensure_dir(path):
"""
create path by first checking its existence,
:param paths: path
:return:
"""
if not os.path.exists(path):
os.makedirs(path)
def read_pkl(data_url):
file = open(data_url,'rb')
content = pickle.load(file)
file.close()
return content |