File size: 8,345 Bytes
3f69a1e | 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 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | """
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}
|