-
Notifications
You must be signed in to change notification settings - Fork 6
Expand file tree
/
Copy pathio_utils.py
More file actions
132 lines (106 loc) · 3.61 KB
/
io_utils.py
File metadata and controls
132 lines (106 loc) · 3.61 KB
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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
# This file is part of the CoverageControl library
#
# Author: Saurav Agarwal
# Contact: [email protected], [email protected]
# Repository: https://github.com/KumarRobotics/CoverageControl
#
# Copyright (c) 2024, Saurav Agarwal
#
# The CoverageControl library is free software: you can redistribute it and/or
# modify it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or (at your
# option) any later version.
#
# The CoverageControl library is distributed in the hope that it will be
# useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along with
# CoverageControl library. If not, see <https://www.gnu.org/licenses/>.
## @file io_utils.py
# @brief The module provides utility functions for loading data from files
"""
The module provides utility functions for loading data from files
"""
import os
import sys
import pathlib
import torch
import yaml
if sys.version_info[1] < 11:
import tomli as tomllib
else:
import tomllib
## @ingroup python_api
class IOUtils:
"""
Class provides the following utility functions:
- load_tensor
- load_yaml
- load_toml
"""
@staticmethod
def sanitize_path(path_str: str) -> str:
"""
Function to sanitize a path string
"""
return os.path.normpath(os.path.expanduser(os.path.expandvars(path_str)))
@staticmethod
def load_tensor(path: (str, pathlib.Path)) -> torch.tensor:
"""
Function to load a tensor from a file
Can load tensors from jit script format files
Args:
path (str): Path to the file
Returns:
tensor: The loaded tensor
None: If the file does not exist
Raises:
FileNotFoundError: If the file does not exist
"""
if isinstance(path, pathlib.Path):
path = str(path)
# Throw error if path does not exist
path = IOUtils.sanitize_path(path)
if not os.path.exists(path):
raise FileNotFoundError(f"IOUtils::load_tensor: File not found: {path}")
# Load data
data = torch.load(path, weights_only=True)
# Extract tensor if data is in jit script format
if isinstance(data, torch.jit.ScriptModule):
tensor = list(data.parameters())[0]
else:
tensor = data
return tensor
@staticmethod
def load_yaml(path: str) -> dict:
"""
Function to load a yaml file
Args:
path (str): Path to the file
Returns:
data: The loaded data
Raises:
FileNotFoundError: If the file does not exist
"""
path = IOUtils.sanitize_path(path)
# Throw error if path does not exist
if not os.path.exists(path):
raise FileNotFoundError(f"IOUtils::load_yaml File not found: {path}")
# Load data
with open(path, "rb") as f:
data = yaml.load(f, Loader=yaml.FullLoader)
return data
@staticmethod
def load_toml(path: str) -> dict: # Throw error if path does not exist
"""
Function to load a toml file
"""
path = IOUtils.sanitize_path(path)
if not os.path.exists(path):
raise FileNotFoundError(f"IOUtils::load_toml: File not found: {path}")
# Load data
with open(path, "rb") as f:
data = tomllib.load(f)
return data