| """ |
| Data loading utilities for VETime anomaly detection. |
| |
| This module provides dataset classes and collate functions for loading and |
| preprocessing time series anomaly detection data. It supports: |
| - Loading preprocessed datasets from pickle files |
| - Converting time series to image representations on-the-fly |
| - Padding and batching sequences of variable lengths |
| - Random masking for self-supervised pretraining |
| """ |
| from typing import Tuple, List, Dict, Any, Optional, Union |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| import pickle |
| from torch.utils.data import Dataset |
| import random |
|
|
| from dataset.pre_image import ts2image_1d |
|
|
|
|
| class AnomalyDataset(Dataset): |
| """ |
| PyTorch Dataset for time series anomaly detection. |
| |
| This dataset class loads preprocessed time series data from pickle files |
| and optionally generates image representations on-the-fly. It supports |
| train/test split and filters out very short sequences. |
| |
| The dataset expects pickle files containing a list of dictionaries, where |
| each dictionary represents a sample with keys: |
| - 'time_series': numpy array of shape (L, C) |
| - 'normal_time_series': numpy array for normal reference |
| - 'labels': numpy array of anomaly labels (0=normal, 1=anomaly) |
| - 'attribute': metadata dictionary |
| |
| Args: |
| dataset_dir: Path to the pickle file containing the dataset. |
| patch_size: Size of patches for image generation. Used to determine |
| target image width. |
| gen_image: If True, generate image representations for all samples |
| during initialization. Default: True. |
| split: Data split to use. 'train' uses all data, 'test' uses the |
| last (1 - train_ratio) portion. Default: 'train'. |
| train_ratio: Ratio of data to use for training when split='test'. |
| Only used when split='test'. Default: 0.95. |
| seed: Random seed for shuffling indices. Default: 42. |
| name: Optional name identifier for the dataset. Default: None. |
| |
| Attributes: |
| data: List of sample dictionaries after filtering and splitting. |
| image_type: Type of image representation ('RGB'). |
| image_h: Height of each channel tile in generated images. |
| |
| """ |
|
|
| def __init__( |
| self, |
| dataset_dir: str, |
| patch_size: int, |
| gen_image: bool = True, |
| split: str = 'train', |
| train_ratio: float = 0.95, |
| seed: int = 42, |
| name: Optional[str] = None |
| ): |
| """ |
| Initialize the AnomalyDataset. |
| |
| Args: |
| dataset_dir: Path to the pickle file containing the dataset. |
| patch_size: Size of patches for image generation. |
| gen_image: If True, generate image representations during init. |
| split: Data split ('train' or 'test'). |
| train_ratio: Ratio of data for training split. |
| seed: Random seed for reproducibility. |
| name: Optional dataset name identifier. |
| """ |
| file_path = dataset_dir |
| self.image_h = patch_size |
| self.gen_image = gen_image |
| self.patch_size = patch_size |
| with open(file_path, 'rb') as f: |
| dataset = pickle.load(f) |
| random.seed(seed) |
| indices = list(range(len(dataset))) |
| random.shuffle(indices) |
| num_train = int(len(dataset) * train_ratio) |
| if split == 'train': |
| selected_indices = indices |
| elif split == 'test': |
| selected_indices = indices[num_train:] |
| else: |
| raise ValueError("split must be 'train' or 'test'") |
|
|
| self.data = [dataset[i] for i in selected_indices] |
| self.data = [x for x in self.data if len(x['time_series']) > 100] |
| self.data.sort(key=lambda x: len(x['time_series'])) |
|
|
| self.image_type = 'RGB' |
| self.name = name |
| if self.gen_image: |
| self.generate_image(self.data) |
|
|
| def __len__(self) -> int: |
| """ |
| Return the number of samples in the dataset. |
| |
| Returns: |
| int: Number of samples in the dataset. |
| """ |
| return len(self.data) |
|
|
| def generate_image(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]: |
| """ |
| Generate image representations for time series samples. |
| |
| This function converts each time series in the data list to an image |
| using ts2image_1d. The image width is determined by the sequence length |
| and patch_size, rounded up to the nearest multiple of patch_size. |
| |
| Args: |
| data: List of sample dictionaries. Each dictionary should contain |
| at least 'time_series' key. The image, period, and padding |
| value will be added in-place. |
| |
| Returns: |
| List[Dict[str, Any]]: The same data list with added keys: |
| - 'image': Generated image array of shape (3, C*h_size, width) |
| - 'period': Detected period (integer) |
| - 'padding_value': Padding values for the image |
| |
| Note: |
| This function modifies the input data list in-place and also |
| returns it for convenience. |
| """ |
| for idx, data0 in enumerate(data): |
| target_length = ((len(data0['time_series']) + self.patch_size - 1) // self.patch_size) * self.patch_size |
| img, period, padding_value = ts2image_1d(data0['time_series'], target_length, self.patch_size) |
| data[idx]['image'] = img |
| data[idx]['period'] = period |
| data[idx]['padding_value'] = padding_value |
| return data |
|
|
| def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Dict, int, torch.Tensor]: |
| """ |
| Get a single sample from the dataset. |
| |
| Args: |
| idx: Index of the sample to retrieve. |
| |
| Returns: |
| A tuple containing: |
| - time_series: Time series data as float32 tensor (L, C) |
| - normal_time_series: Normal reference time series (L, C) |
| - image: Image representation as float32 tensor (3, H, W) |
| - labels: Anomaly labels as long tensor (L,) |
| - attribute: Metadata dictionary |
| - period: Detected period (int) |
| - padding_value: Padding values as float32 tensor (3, C, 1) |
| """ |
| sample = self.data[idx] |
| img_tensor = torch.tensor(sample['image'], dtype=torch.float32) |
| time_series = torch.tensor(sample['time_series'], dtype=torch.float32) |
| normal_time_series = torch.tensor(sample['normal_time_series'], dtype=torch.float32) |
| labels = torch.tensor(sample['labels'], dtype=torch.long) |
| attribute = sample['attribute'] |
| period = sample['period'] |
| padding_value = torch.tensor(sample['padding_value'], dtype=torch.float32) |
| return time_series, normal_time_series, img_tensor, labels, attribute, period, padding_value |
|
|
|
|
| def collate_fn( |
| batch: List[Tuple], |
| patch_size: int |
| ) -> Dict[str, Union[torch.Tensor, List, Tuple]]: |
| """ |
| Collate function for batching anomaly detection samples. |
| |
| This function processes a batch of samples from AnomalyDataset and: |
| 1. Concatenates all time series and computes global mean/std for normalization |
| 2. Pads all sequences to the same length (multiple of patch_size) |
| 3. Generates attention masks for valid sequence positions |
| 4. Applies random masking for self-supervised learning |
| 5. Pads images to match the target width |
| |
| Args: |
| batch: List of samples from AnomalyDataset.__getitem__. Each sample is |
| a tuple of (time_series, normal_time_series, img_tensor, labels, |
| attribute, period, padding_value). |
| patch_size: Size of patches for masking and padding alignment. |
| |
| Returns: |
| A dictionary containing: |
| - 'time_series': Padded time series tensor (B, L_max, C) |
| - 'normal_time_series': Padded normal reference tensor (B, L_max, C) |
| - 'mask_time_series': Time series with random patches masked (B, L_max, C) |
| - 'image': Padded image tensor (B, 3, H, W_max) |
| - 'mask': Boolean mask indicating masked positions (B, L_max) |
| - 'labels': Padded label tensor (B, L_max) with -1 for padding |
| - 'attention_mask': Boolean mask for valid positions (B, L_max) |
| - 'period': Tuple of periods for each sample in batch |
| - 'padding_value': Tensor of padding values (B, 3, C, 1) |
| |
| Note: |
| - Time series are normalized using batch-wide statistics |
| - Labels are padded with -1 (ignored in loss computation) |
| - Random masking applies mask_ratio=0.3 to valid sequence regions only |
| """ |
| time_series_list, normal_time_series_list, img_tensor, labels_list, attribute_list, period, padding_value = zip(*batch) |
| |
| if time_series_list[0].ndim == 1: |
| time_series_tensors = [ts.unsqueeze(-1) for ts in time_series_list] |
| normal_time_series_tensors = [nts.unsqueeze(-1) for nts in normal_time_series_list] |
| else: |
| time_series_tensors = [ts for ts in time_series_list] |
| normal_time_series_tensors = [nts for nts in normal_time_series_list] |
|
|
| concatenated = torch.cat(time_series_tensors, dim=0) |
| mean = concatenated.mean(dim=0, keepdim=True) |
| std = concatenated.std(dim=0, keepdim=True) + 1e-4 |
| time_series_tensors = [(ts - mean) / std for ts in time_series_tensors] |
| normal_time_series_tensors = [(nts - mean) / std for nts in normal_time_series_tensors] |
|
|
| labels = [label for label in labels_list] |
| lengths = [t.size(0) for t in labels] |
| max_len = max(lengths) |
| max_idx = lengths.index(max_len) |
| target_length = ((max_len + patch_size - 1) // patch_size) * patch_size |
|
|
| def padding_to_target_length(list0, value): |
| original_tensor = list0[max_idx] |
| pad_shape = [0, 0] * original_tensor.dim() |
| pad_shape[-1] = target_length - max_len |
| padded_tensor = torch.nn.functional.pad(original_tensor, pad=pad_shape, mode='constant', value=value) |
| list0[max_idx] = padded_tensor |
| return torch.nn.utils.rnn.pad_sequence(list0, batch_first=True, padding_value=value) |
|
|
| padded_time_series = padding_to_target_length(time_series_tensors, 0.0) |
| normal_time_series_tensors = padding_to_target_length(normal_time_series_tensors, 0.0) |
| padded_labels = padding_to_target_length(labels, -1) |
|
|
| image_inputs = image_right_padding(img_tensor, target_length, padding_value) |
| sequence_lengths = [ts.size(0) for ts in time_series_tensors] |
| B, max_seq_len, num_features = padded_time_series.shape |
| attention_mask = torch.ones(B, max_seq_len, dtype=torch.bool) |
|
|
| for i, length in enumerate(sequence_lengths): |
| attention_mask[i, length:] = False |
|
|
| mask_time_series, mask = create_random_mask(padded_time_series, attention_mask, patch_size) |
| normal_time_series_tensors, mask = create_random_mask(normal_time_series_tensors, attention_mask, patch_size) |
|
|
| return { |
| 'time_series': padded_time_series, |
| 'normal_time_series': normal_time_series_tensors, |
| 'mask_time_series': mask_time_series, |
| 'image': image_inputs, |
| 'mask': mask, |
| 'labels': padded_labels, |
| 'attention_mask': attention_mask, |
| 'period': period, |
| 'padding_value': padding_value, |
| } |
|
|
|
|
| def image_right_padding( |
| imgs: List[torch.Tensor], |
| max_width: int, |
| p_values: torch.Tensor |
| ) -> torch.Tensor: |
| """ |
| Pad images on the right side to match target width. |
| |
| This function extends images that are shorter than max_width by padding |
| on the right side. The padding uses the provided padding values to maintain |
| consistency with the time series padding strategy. |
| |
| Args: |
| imgs: List of image tensors, each of shape (3, H, W_i) where W_i may |
| vary across samples. |
| max_width: Target width for all images. Images with W < max_width will |
| be padded; images with W >= max_width remain unchanged. |
| p_values: Tensor of padding values with shape (B, 3, C, 1) or compatible. |
| Each sample's padding value is used to fill its padded region. |
| |
| Returns: |
| torch.Tensor: Stacked tensor of padded images with shape (B, 3, H, max_width). |
| All images have the same width after processing. |
| |
| Note: |
| - Padding is applied only on the right side (width dimension) |
| - Padding values are transposed to match the image channel format |
| """ |
| padded_images = [] |
| for i in range(len(imgs)): |
| img = imgs[i] |
| C, H_size, W = img.shape |
| p_value = p_values[i] |
| if max_width > W: |
| right_padding = max_width - img.shape[2] |
| padding = (0, right_padding, 0, 0) |
| padded_img = F.pad(img.unsqueeze(0), padding, mode='constant', value=0).squeeze(0) |
| padded_img[:, :, W:] = p_value.T[:, :, None] |
| else: |
| padded_img = img |
| padded_images.append(padded_img) |
| return torch.stack(padded_images) |
|
|
|
|
| def create_random_mask( |
| time_series: torch.Tensor, |
| attention_mask: torch.Tensor, |
| patch_size: int = 14, |
| mask_ratio: float = 0.3 |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| """ |
| Create random mask for time series patches in self-supervised learning. |
| |
| This function generates a binary mask that randomly masks patches of the |
| time series while respecting the attention mask (only valid sequence |
| positions are masked). The masked positions are replaced with small |
| Gaussian noise for reconstruction-based pretraining. |
| |
| The masking strategy: |
| 1. Divide the valid sequence into patches of size `patch_size` |
| 2. Randomly select `mask_ratio` fraction of patches to mask |
| 3. Apply mask only to valid positions (respecting attention_mask) |
| 4. Replace masked positions with Gaussian noise (std=0.1) |
| |
| Args: |
| time_series: Input time series tensor of shape (B, L, C) where B is |
| batch size, L is sequence length, C is number of features. |
| attention_mask: Boolean tensor of shape (B, L) indicating valid positions |
| (True = valid, False = padding). |
| patch_size: Size of each patch for masking. The sequence is divided |
| into non-overlapping patches of this size. Default: 14. |
| mask_ratio: Ratio of patches to mask within valid regions. Must be in |
| [0, 1]. Default: 0.3. |
| |
| Returns: |
| A tuple containing: |
| masked_time_series: Time series with masked positions replaced by |
| small Gaussian noise. Same shape as input (B, L, C). |
| mask: Boolean tensor of shape (B, L) indicating masked positions |
| (True = masked, False = visible). |
| """ |
| batch_size, seq_len, num_features = time_series.shape |
| mask = torch.zeros(batch_size, seq_len, dtype=torch.bool) |
| valid_mask = attention_mask |
| |
| for i in range(batch_size): |
| valid_length = valid_mask[i].sum().item() |
| num_valid_patches = (valid_length + patch_size - 1) // patch_size |
| num_masked = max(1, int(num_valid_patches * mask_ratio)) if mask_ratio > 0 else 0 |
| num_masked = min(num_masked, num_valid_patches) |
|
|
| if num_masked > 0: |
| masked_patches = torch.randperm(num_valid_patches)[:num_masked] |
| for j in masked_patches: |
| start_idx = j * patch_size |
| end_idx = min((j + 1) * patch_size, valid_length) |
| mask[i, start_idx:end_idx] = 1 |
| |
| mask = mask & valid_mask |
| masked_time_series = time_series.clone() |
| mask_expanded = mask.unsqueeze(-1).expand(-1, -1, num_features) |
| masked_time_series[mask_expanded] = torch.randn_like(masked_time_series[mask_expanded]) * 0.1 |
|
|
| return masked_time_series, mask |
|
|
