122 lines
4.2 KiB
Python
122 lines
4.2 KiB
Python
"""
|
|
Reads the files of the HUST Motor dataset and parses them into a usable format for machine learning tasks.
|
|
Uses the tag and frequency information to create a structured dataset.
|
|
"""
|
|
|
|
from enum import Enum
|
|
from os import listdir, makedirs
|
|
from os.path import exists, getmtime, join
|
|
import pickle
|
|
|
|
from dataset.datasetParser import DatasetParser
|
|
|
|
class MUSTMotorProblemType(Enum):
|
|
HEALTHY = "H"
|
|
BEARING_FAULT = "BF"
|
|
BOWED_ROTOR = "BOW"
|
|
BROKEN_ROTOR_BARS = "BROKEN"
|
|
ROTOR_MISALIGNMENT = "MISAL"
|
|
VOLTAGE_UNBALANCE = "UNBAL"
|
|
|
|
class MUSTMotorFrequency(Enum):
|
|
FREQ_5HZ = "5HZ"
|
|
FREQ_10HZ = "10HZ"
|
|
FREQ_20HZ = "20HZ"
|
|
FREQ_30HZ = "30HZ"
|
|
|
|
|
|
class MUSTMotorDatasetParser(DatasetParser):
|
|
DATASET_PATH = "dataset/HUSTMotor"
|
|
CACHE_PATH = "dataset/.cache/hust_motor_dataset.pkl"
|
|
DATA_COLUMNS = ("time", "x", "y", "z", "sound")
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def get_files(self):
|
|
"""
|
|
Retrieves the list of dataset files from the specified dataset path.
|
|
"""
|
|
files = listdir(self.DATASET_PATH)
|
|
# Get the text files
|
|
text_files = sorted(f for f in files if f.endswith('.txt'))
|
|
return text_files
|
|
|
|
def get_label_from_filename(self, filename):
|
|
"""
|
|
Extracts the label (problem type) from the given filename.
|
|
"""
|
|
filename = filename.split('.')[0] # Remove the file extension
|
|
filename = filename.split('_')
|
|
problem_type = MUSTMotorProblemType(filename[0]).value
|
|
frequency = MUSTMotorFrequency(filename[1]).value
|
|
return problem_type, frequency
|
|
|
|
def parse_dataset(self):
|
|
"""
|
|
Parses the dataset files into one record per measurement row.
|
|
|
|
Each record contains the time, three vibration axes, sound, and the
|
|
problem type and frequency encoded in the source filename.
|
|
"""
|
|
files = self.get_files()
|
|
if self._cache_is_current(files):
|
|
with open(self.CACHE_PATH, "rb") as cache_file:
|
|
print(f"Loading dataset from cache: {self.CACHE_PATH}")
|
|
return pickle.load(cache_file)
|
|
|
|
dataset = []
|
|
for file in files:
|
|
problem_type, frequency = self.get_label_from_filename(file)
|
|
data_started = False
|
|
with open(join(self.DATASET_PATH, file), encoding="utf-8") as dataset_file:
|
|
for line_number, line in enumerate(dataset_file, start=1):
|
|
if line.strip() == "Time (seconds) and Data Channels":
|
|
data_started = True
|
|
continue
|
|
|
|
if not data_started or not line.strip():
|
|
continue
|
|
|
|
values = line.split()
|
|
if len(values) != len(self.DATA_COLUMNS):
|
|
raise ValueError(
|
|
f"Expected {len(self.DATA_COLUMNS)} values in {file} "
|
|
f"at line {line_number}, got {len(values)}"
|
|
)
|
|
|
|
record = dict(zip(self.DATA_COLUMNS, map(float, values)))
|
|
record["problem_type"] = problem_type
|
|
record["frequency"] = frequency
|
|
dataset.append(record)
|
|
|
|
makedirs(join(self.DATASET_PATH, "..", ".cache"), exist_ok=True)
|
|
with open(self.CACHE_PATH, "wb") as cache_file:
|
|
pickle.dump(dataset, cache_file, protocol=pickle.HIGHEST_PROTOCOL)
|
|
|
|
return dataset
|
|
|
|
def get_dataset_motor_problems(self):
|
|
"""
|
|
Retrieves the list of unique problem types present in the dataset.
|
|
"""
|
|
files = self.get_files()
|
|
problems = set()
|
|
for file in files:
|
|
problem_type, _ = self.get_label_from_filename(file)
|
|
problems.add(problem_type)
|
|
return sorted(problems)
|
|
|
|
def _cache_is_current(self, files):
|
|
if not exists(self.CACHE_PATH):
|
|
return False
|
|
|
|
cache_timestamp = getmtime(self.CACHE_PATH)
|
|
return all(
|
|
cache_timestamp >= getmtime(join(self.DATASET_PATH, file))
|
|
for file in files
|
|
)
|
|
|
|
def get_raw_sampling_rate(self):
|
|
return 25600.0 # 25.6 kHz from dataset documentation
|