Initial model

dataset/HUSTMotorDatasetParser.py

@@ -0,0 +1,121 @@
+"""
+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