""" Multimodal PC Fault Detection - Dataset Module =============================================== Handles audio/visual data loading, synthetic visual generation, ESC-50 → fault mapping, and paired multimodal batching. """ import os import random import numpy as np import torch from torch.utils.data import Dataset from PIL import Image, ImageDraw, ImageFont from typing import Dict, Optional, Tuple import torchaudio import torchaudio.transforms as T from datasets import load_dataset from config import ( FAULT_CLASSES, ESC50_TO_FAULT, ESC50_CATEGORY_TO_TARGET, VISUAL_SYNTHESIS, DataConfig, ModelConfig ) class SyntheticVisualGenerator: """Generates synthetic PC diagnostic screen images for each fault class.""" def __init__(self, image_size: int = 224): self.image_size = image_size self.templates = VISUAL_SYNTHESIS def generate(self, fault_class_name: str, variation_seed: int = 0) -> Image.Image: rng = random.Random(variation_seed) template = self.templates[fault_class_name] base_r, base_g, base_b = template["color_dominant"] r = max(0, min(255, base_r + rng.randint(-30, 30))) g = max(0, min(255, base_g + rng.randint(-30, 30))) b = max(0, min(255, base_b + rng.randint(-30, 30))) img = Image.new("RGB", (self.image_size, self.image_size), (r, g, b)) draw = ImageDraw.Draw(img) texts = template["text_overlay"] y_offset = rng.randint(10, 40) for text in texts: text_color = (255, 255, 255) if (r + g + b) / 3 < 128 else (0, 0, 0) x = rng.randint(5, max(10, self.image_size // 4)) draw.text((x, y_offset), text, fill=text_color) y_offset += rng.randint(25, 45) if fault_class_name == "normal_operation": bar_y = rng.randint(100, 180) draw.rectangle([10, bar_y, int(self.image_size * rng.uniform(0.5, 0.95)), bar_y + 15], fill=(0, 200, 0)) elif fault_class_name == "system_crash": cx, cy = self.image_size // 2, self.image_size // 3 draw.text((cx - 10, cy), ":(", fill=(255, 255, 255)) elif fault_class_name == "overheating_fan": gauge_x = rng.randint(40, 120) draw.rectangle([gauge_x, 150, gauge_x + 30, 200], outline=(255, 255, 255)) fill_h = int(50 * rng.uniform(0.7, 1.0)) draw.rectangle([gauge_x + 2, 200 - fill_h, gauge_x + 28, 198], fill=(255, 50, 50)) elif fault_class_name == "storage_failure": bar_y = 160 draw.rectangle([20, bar_y, 200, bar_y + 12], outline=(255, 255, 255)) progress = rng.uniform(0.1, 0.4) draw.rectangle([22, bar_y + 2, int(22 + 176 * progress), bar_y + 10], fill=(255, 165, 0)) elif fault_class_name == "boot_failure": cx, cy = rng.randint(10, 60), rng.randint(140, 200) draw.rectangle([cx, cy, cx + 10, cy + 15], fill=(255, 255, 255)) img_array = np.array(img) noise = np.random.RandomState(variation_seed).randint(-15, 16, img_array.shape, dtype=np.int16) img_array = np.clip(img_array.astype(np.int16) + noise, 0, 255).astype(np.uint8) return Image.fromarray(img_array) class AudioPreprocessor: def __init__(self, config: DataConfig): self.config = config self.target_length = int(config.sample_rate * config.audio_duration) self.mel_transform = T.MelSpectrogram( sample_rate=config.sample_rate, n_fft=config.n_fft, hop_length=config.hop_length, n_mels=config.n_mels, f_min=config.fmin, f_max=config.fmax, ) self.amplitude_to_db = T.AmplitudeToDB(top_db=80) def process(self, waveform, sample_rate): if waveform.dim() > 1 and waveform.shape[0] > 1: waveform = waveform.mean(dim=0, keepdim=True) elif waveform.dim() == 1: waveform = waveform.unsqueeze(0) if sample_rate != self.config.sample_rate: waveform = T.Resample(orig_freq=sample_rate, new_freq=self.config.sample_rate)(waveform) if waveform.shape[-1] < self.target_length: waveform = torch.nn.functional.pad(waveform, (0, self.target_length - waveform.shape[-1])) else: waveform = waveform[:, :self.target_length] return self.amplitude_to_db(self.mel_transform(waveform)) def augment(self, log_mel, training=True): if not training: return log_mel log_mel = T.FrequencyMasking(freq_mask_param=self.config.freq_mask_max)(log_mel) return T.TimeMasking(time_mask_param=self.config.time_mask_max)(log_mel) class PCFaultDataset(Dataset): def __init__(self, config, model_config, split="train", vit_processor=None, ast_feature_extractor=None, augment=True): self.config, self.model_config, self.split = config, model_config, split self.augment = augment and (split == "train") self.vit_processor, self.ast_feature_extractor = vit_processor, ast_feature_extractor self.esc50 = load_dataset(config.esc50_dataset, split="train") self.samples = [] for idx, row in enumerate(self.esc50): category = row["category"] if category in ESC50_TO_FAULT: fault_label, fold = ESC50_TO_FAULT[category], row["fold"] if (split == "train" and fold != config.val_fold) or (split == "val" and fold == config.val_fold): self.samples.append({"audio_idx": idx, "fault_label": fault_label, "category": category}) self.visual_gen = SyntheticVisualGenerator(config.image_size) self.audio_preprocessor = AudioPreprocessor(config) label_counts = {} for s in self.samples: label_counts[s["fault_label"]] = label_counts.get(s["fault_label"], 0) + 1 print(f"\\n[PCFaultDataset] Split: {split}, Total: {len(self.samples)}") for label, count in sorted(label_counts.items()): print(f" Class {label} ({FAULT_CLASSES[label]}): {count}") def __len__(self): return len(self.samples) def __getitem__(self, idx): sample = self.samples[idx] fault_label, fault_name = sample["fault_label"], FAULT_CLASSES[sample["fault_label"]] esc50_row = self.esc50[sample["audio_idx"]] audio_data = esc50_row["audio"] if self.ast_feature_extractor is not None: audio_array = np.array(audio_data["array"], dtype=np.float32) sr = audio_data["sampling_rate"] if sr != 16000: wf = T.Resample(orig_freq=sr, new_freq=16000)(torch.tensor(audio_array).unsqueeze(0)) audio_array = wf.squeeze(0).numpy() audio_values = self.ast_feature_extractor(audio_array, sampling_rate=16000, return_tensors="pt")["input_values"].squeeze(0) else: waveform = torch.tensor(audio_data["array"], dtype=torch.float32) log_mel = self.audio_preprocessor.process(waveform, audio_data["sampling_rate"]) if self.augment: log_mel = self.audio_preprocessor.augment(log_mel, training=True) audio_values = log_mel.squeeze(0) visual_image = self.visual_gen.generate(fault_name, variation_seed=idx * 7 + fault_label) if self.vit_processor is not None: pixel_values = self.vit_processor(images=visual_image, return_tensors="pt")["pixel_values"].squeeze(0) else: img_array = np.array(visual_image.resize((224, 224))).astype(np.float32) / 255.0 img_array = (img_array - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225] pixel_values = torch.tensor(img_array).permute(2, 0, 1) return {"pixel_values": pixel_values, "audio_values": audio_values, "labels": torch.tensor(fault_label, dtype=torch.long)} def multimodal_collate_fn(batch): pixel_values = torch.stack([b["pixel_values"] for b in batch]) labels = torch.stack([b["labels"] for b in batch]) audio_list = [b["audio_values"] for b in batch] max_len = max(a.shape[-1] for a in audio_list) padded_audio = [torch.nn.functional.pad(a, (0, max_len - a.shape[-1])) if a.shape[-1] < max_len else a for a in audio_list] return {"pixel_values": pixel_values, "audio_values": torch.stack(padded_audio), "labels": labels}